来自德克萨斯大学奥斯汀分校、悉尼科技大学和谷歌的研究者提出了一个无需训练就能自动扩展框架 As-ViT,其能以高效和有原则的方式自动发现和扩展 ViT。
当前 Vision Transformers (ViT)领域有两个主要的痛点:1、缺少对 ViT 进行设计和扩展的有效方法;2、训练 ViT 的计算成本比卷积网络要大得多。
为了解决这两个问题,来自得克萨斯大学奥斯汀分校、悉尼科技大学和谷歌的研究者提出了 As-ViT(Auto-scaling Vision Transformers),这是一个无需训练的 ViT 自动扩展框架,它能以高效且有原则的方式自动设计和扩展 ViT。
论文链接:https://arxiv.org/abs/2202.11921
具体来说,研究人员首先利用无训练搜索过程设计了 ViT 拓扑的「种子」,这种极快的搜索是通过对 ViT 网络复杂性的全面研究来实现的,从而产生了与真实准确度的强 Kendall-tau 相关性。其次,从「种子」拓扑开始,通过将宽度 / 深度增加到不同的 ViT 层来自动化 ViT 的扩展规则,实现了在一次运行中具有不同数量参数的一系列架构。最后,基于 ViT 在早期训练阶段可以容忍粗粒度 tokenization 的经验,该研究提出了一种渐进式 tokenization 策略来更快、更节约地训练 ViT。
作为统一的框架,As-ViT 在分类(ImageNet-1k 上 83.5% 的 top1)和检测(COCO 上 52.7% 的 mAP)任务上实现了强大的性能,无需任何手动调整或扩展 ViT 架构,端到端模型设计和扩展过程在一块 V100 GPU 上只需 12 小时。
具有网络复杂度的 ViT 自动设计和扩展
为加快 ViT 设计并避免繁琐的手动工作,该研究希望以高效、自动化和有原则的 ViT 搜索和扩展为目标。具体来说有两个问题需要解决:1)在训练成本最小甚至为零的情况下,如何高效地找到最优的 ViT 架构拓扑?2)如何扩大 ViT 拓扑的深度和宽度以满足模型尺寸的不同需求?
扩展 ViT 的拓扑空间
在设计和扩展之前,首先是为 As-ViT 扩展的拓扑搜索空间:首先将输入图像嵌入到 1/4 尺度分辨率的块中,并采用逐级空间缩减和通道加倍策略。这是为了方便密集预测任务,例如需要多尺度特征的检测。
通过流形传播评估初始化时的 ViT 复杂性
ViT 训练速度很慢,因此,通过评估训练模型的准确率来进行架构搜索的成本将高得让人难以承受。最近学界出现很多用基于 ReLU 的 CNN 的免训练神经架构搜索方法,利用局部线性图 (Mellor et al., 2020)、梯度敏感性 (Abdelfattah et al., 2021)、线性区域数量 (Chen et al., 2021e;f) 或网络拓扑(Bhardwaj et al., 2021)等方式。
然而 ViT 配备了更复杂的非线性函数如 self-attention、softmax 和 GeLU。因此需要以更一般的方式衡量其学习能力。在新研究中,研究者考虑通过 ViT 测量流形传播的复杂性,以估计复杂函数可以如何被 ViT 逼近。直观地说,一个复杂的网络可以在其输出层将一个简单的输入传播到一个复杂的流形中,因此可能具有很强的学习能力。在 UT Austin 的工作中,他们通过 ViT 映射简单圆输入的多种复杂性:h(θ) = √ N [u^0 cos(θ) + u^1 sin(θ)]。这里,N 是 ViT 输入的维度(例如,对于 ImageNet 图像,N = 3 × 224 × 224),u^0 和 u^1 形成了圆所在的 R^N 的二维子空间的标准正交基。
搜索 ViT 拓扑奖励
研究者提出了基于 L^E 的免训练搜索(算法 1),大多数 NAS(神经架构搜索)方法将单路径或超级网络的准确率或损失值评估为代理推理。当应用于 ViT 时,这种基于训练的搜索将需要更多的计算成本。对于采样的每个架构,这里不是训练 ViT,而是计算 L^E 并将其视为指导搜索过程的奖励。
除了 L^E,还包括 NTK 条件数 κΘ = λ_max/λ_min ,以指示 ViT 的可训练性(Chen et al., 2021e; Xiao et al., 2019; Yang, 2020; Hron et al., 2020)。λ_max 和 λ_min 是 NTK 矩阵 Θ 的最大和最小特征值。
搜索使用强化学习方法,策略被定为联合分类分布,并通过策略梯度进行更新,该研究将策略更新为 500 step,观察到足以使策略收敛(熵从 15.3 下降到 5.7)。搜索过程非常快:在 ImageNet-1k 数据集上只有七个 GPU 小时 (V100),这要归功于绕过 ViT 训练的 L^E 的简单计算。为了解决 L^E 和 κΘ 的不同大小,该研究通过它们的相对值范围对它们进行归一化(算法 1 中的第 5 行)。
表 3 总结了新搜索方法的 ViT 拓扑统计数据。我们可以看到 L^E 和 κΘ 高度偏好:(1)具有重叠的 token (K_1∼K_4 都大于 stride ),以及(2)在更深层中更大的 FFN 扩展率(E_1 < E_2 < E_3 < E_4)。在注意力分裂和正面数量上没有发现 L^E 和 κΘ 的明显偏好。