ViT 自主的原则型扩展
得到最优拓扑后,接下来要解决的一个问题是:如何平衡网络的深度和宽度?
目前,对于 ViT 扩展没有这样的经验法则。最近的工作试图扩大或增长不同大小的卷积网络以满足各种资源限制(Liu et al., 2019a; Tan & Le, 2019)。然而,为了自动找到一个有原则的扩展规则,训练 ViT 将花费巨大的计算成本。也可以搜索不同的 ViT 变体(如第 3.3 节中所述),但这需要多次运行。相反,「向上扩展,scaling-up」是在一个实验中生成多个模型变体的更自然的方式。因此,该研究试图以一种免训练且有原则的有效方法将搜索到的基本「种子」ViT 扩展到更大的模型。算法 2 中描述了这种自动扩展方法:
初始架构的每个阶段都有一个注意力块,初始隐藏维度 C = 32。每次迭代找出最佳深度和宽度,以进行进一步向上扩展。对于深度,该研究尝试找出要加深哪个阶段(即,在哪个阶段添加一个注意力块);对于宽度,该研究尝试发现最佳扩展比(即,将通道数扩大到什么程度)。
扩展轨迹如下图 3 所示。比较自主扩展和随机扩展,研究者发现扩展原则更喜欢舍弃深度来换取更多宽度,使用更浅但更宽的网络。这种扩展更类似于 Zhai et al. (2021) 开发的规则。相比之下,ResNet 和 Swin Transformer (Liu et al., 2021) 选择更窄更深。
通过渐进灵活的 re-tokenization 进行高效的 ViT 训练
该研究通过提出渐进灵活的 re-tokenization 训练策略来提供肯定的答案。为了在训练期间更新 token 的数量而不影响线性投影中权重的形状,该研究在第一个线性投影层中采用不同的采样粒度。以第一个投影核 K_1 = 4 且 stride = 4 为例:训练时研究者逐渐将第一个投影核的 (stride, dilation) 对逐渐变为 (16, 5), (8, 2) 和 (4 , 1),保持权重的形状和架构不变。
这种 re-tokenization 的策略激发了 ViT 的课程学习(curriculum learning):训练开始时引入粗采样以显着减少 token 的数量。换句话说,As-ViT 在早期训练阶段以极低的计算成本(仅全分辨率训练的 13.2% FLOPs)快速从图像中学习粗略信息。在训练的后期阶段,该研究逐渐切换到细粒度采样,恢复完整的 token 分辨率,并保持有竞争力的准确率。如图 4 所示,当在早期训练阶段使用粗采样训练 ViT 时,它仍然可以获得很高的准确率,同时需要极低的计算成本。不同采样粒度之间的转换引入了性能的跳跃,最终网络恢复了具有竞争力的最终性能。
如图 4 所示,当 ViT 在早期训练阶段使用粗采样训练 ViT 时,它仍然可以获得很高的准确率,同时需要极低的计算成本。不同采样粒度之间的转换引入了性能的跳跃,最终网络恢复了具有竞争力的最终性能。
实验
AS-VIT:自动扩展 VIT
该研究在表 4 中展示了搜索到的 As-ViT 拓扑。这种架构在第一个投影(tokenization)step 和三个重新嵌入 step 中,促进了 token 之间的强烈重叠。FFN 扩展比首先变窄,然后在更深的层变宽。利用少量注意力拆分来更好地聚合全局信息。
图像分类
下表 5 展示了 As-ViT 与其他模型的比较。与之前基于 Transformer 和基于 CNN 的架构相比,As-ViT 以相当数量的参数和 FLOP 实现了 SOTA 性能。
高效训练
研究者调整了表 6 中为每个 token 减少阶段的时期,并将结果显示在表 6 中。标准训练需要 42.8 TPU 天,而高效训练可节省高达 56.2% 的训练 FLOP 和 41.1% 的训练 TPU 天,仍然达到很高的准确率。
拓扑和扩展的贡献
为了更好地验证搜索型拓扑和扩展规则的贡献,该研究进行了更多的消融研究(表 7)。首先,在扩展之前直接训练搜索到的拓扑。该研究搜索的种子拓扑优于图 2 中 87 个随机拓扑中的最佳拓扑。
第二,该研究将基于复杂度的规则与「随机扩展 + As-ViT 拓扑」进行比较。在不同的扩展下,该研究的自动扩展也优于随机扩展。
COCO 数据集上的目标检测
该研究将 As-ViT 与标准 CNN 和之前的 Transformer 网络进行了比较。比较是通过仅更改主干而其他设置未更改来进行的。从下表 8 的结果可以看出,As-ViT 也可以捕获多尺度特征并实现最先进的检测性能,尽管它是在 ImageNet 上设计的,并且它的复杂性是为分类而测量的。