Scaling Vision Transformers

论文：https://arxiv.org/abs/2106.04560

1简介

视觉Transformer(ViT)等基于注意力的神经网络最近在许多计算机视觉基准测试中取得了最先进的结果。比例是获得出色结果的主要因素，因此，了解模型的scaling属性是有效设计的关键。虽然已经研究了扩展Transformer语言模型的规律，但尚不清楚Vision Transformers如何扩展。

为了解决这个问题，作者向上和向下扩展ViT模型和数据，并描述错误率、数据和计算之间的关系。在此过程中，作者改进了ViT的架构和训练，减少了内存消耗并提高了结果模型的准确性。结果，作者成功地训练了一个具有20亿个参数的ViT模型，该模型在ImageNet上达到了90.45%的Top-1准确率。该模型在小样本学习上也表现良好，例如，在ImageNet上每类只有10个examples的情况下可以达到84.86%的Top-1准确率。

 

2主要结论

作者首先介绍了关于扩展趋势的主要结果。在接下来的实验中在多达30亿个weakly-labelled图像上训练了几个ViT模型。作者期间改变了体系结构的大小，训练图像的数量和训练时间。为了评估模型学习的质量，作者测量了一下方面：

• 通过在冻结的权值上训练线性分类器来实现的few-shot transfer；

• 通过对所有数据进行fine-tuning来实现对整个模型的transfer。

2.1 将计算、模型和数据一起放大

下图显示了ImageNet上的10-shot线性评估和微调评估。

 

对于模型大小和数据大小的每个组合，作者预训练了不同数量的Steps。在图中，连接点表示经过不同步骤训练的同一模型。

结论：

• 首先，将计算、模型和数据一起放大可以提高模型表征能力。

• 其次，模型的大小会影响表征能力。

• 第三，大模型受益于额外的数据，甚至需要超过1 Billion图像。

2.2 Double-saturating power law

对于超过2个数量级的计算，计算和性能之间的关系遵循power-law()，在log-log图上形成一条直线。然而，在compute spectrum的2端都观察到了饱和。在计算的最后最大模型误差率并不趋向于零错误率。如果按照作者的观察进行推断，无限容量模型将得到一个非零误差。

其实之前生成模型也观察到了这种类似的效应,作者将这种残差称为任务的不可约熵。通过作者绘制错误率图像，信息论的解释并不适用，但作者的观察支持ImageNet的基本性能上限的概念。根据该定律，这个饱和对应于错误率的一个附加常数c:。

在compute spectrum的前面看到小模型的饱和；最小模型的性能优于幂律模型的预测。出现这种饱和是因为即使是普通的解决方案也可以实现非零误差。例如，预测大多数类(几乎为零计算)将获得与其在测试集中出现频率相关的精度。在生成模型中没有观察到这个下界，要么是因为它们的最小模型大到足以避免这个区域，要么是因为对数损失在性能比精度更差的情况下达到饱和(最终会达到饱和)。这个饱和对应于x轴上的位移中的d。这个常数表明零计算模型仍将获得非零精度。

2.3 大模型的样本效率更高

下图显示了预处理过程中“seen”的图像总数(批量大小乘以step数)的表征质量。除了ImageNet微调和公共验证集上的线性10-shot结果，作者还给出了ImageNet微调模型的结果，其中ImageNet-v2测试集作为泛化性的指标。在这幅图中展示了对30亿张图像进行预训练的3个ViT模型的结果。

 

作者观察到更大的模型样本效率更高，可以用更少的图像达到相同的错误率水平。对于10-shot Ti/16模型需要看到近100倍以上的图像来匹配L/16模型的表征质量。当进行微调时，这个因数从100降低到大约20。实验结果表明，在有足够的数据的情况下，训练一个更大的模型以较少的step训练是更好的。

2.4 ViT-G/14结果

ViT-G/14比以前最好的ViT-H/14模型有很大的优势(超过5%)，每类10个样本的准确率达到84.86%。每类10张图片不到1%ImageNet数据(每个类13个例子)，通常用于自监督和半监督学习。作为参考，下图显示了2个最先进的自监督学习模型，SimCLR v2和BYOL，使用了1%的ImageNet数据。但是请注意，这些方法有很大的不同:ViT-G/14使用大量缺乏监督的数据，并且只进行一次预训练，然后转移到不同的任务中。同时，自监督学习模型使用无标记但领域内的数据进行训练，并以单个任务为目标。

 

 

表1显示了其余基准测试的结果。ViT-G/14在ImageNet上达到90.45%的top-1精度，设置新的艺术状态。在ImageNet-v2上，ViT-G/14比基于EfficientNet-L2的Noisy Student模型改进了3%。对于ReaL, ViT-G/14仅比ViT-H和BiT-L略胜一筹，这再次表明ImageNet分类任务很可能达到饱和点。对于ObjectNet来说，ViT-G/14比BiT-L表现要好很多，比Noisy Student好2%，但比CLIP落后2%。请注意，与其他方法不同，CLIP不会在ImageNet上进行微调，而是直接在ObjectNet上进行评估，这可能提高了它的健壮性。最后，当将ViT-G/14模型转移到VTAB时，它在所有任务中只使用一个超参数就能得到一致更好的结果。

3Method details

作者对ViT模型和训练过程提出了一些改进。这些改进大都很容易实现，并且可以显著提高内存利用率和模型质量。它们允许单独使用数据并行性来训练ViT-G/14，整个模型拟合在一个单独的模型上TPUv3核心。

3.1 “Head”解耦权重衰减

在低数据条件下，权值衰减对model adaptation有显著影响。作者对这一现象进行了mid-size规模的研究。

作者发现，对于模型中的最后一个线性层(“Head”)和其余的权重(“Body”)可以从解耦权值衰减强度中获益。

 

上图展示了这种效应:在JFT-300M上训练一个collection ViT-B/32模型，每个cell对应不同的head/body权重衰减值。对角线对应于对2个衰变使用相同的值。可以观察到最佳的性能表现是非对角线的(例如，head和body的权值衰减是解耦的)。有趣的是，作者观察到head的权值衰减尽管提高了迁移性能，但是也降低了pre-training(upstream)任务的表现。

对于这种现象，作者还没有一个完整的解释。然而，如果假设Head的权值衰减越大，代表就会在类别之间有更大的差距，从而更好地few-shot adaptation。这与支持向量机背后的主要思想相似。这种大的衰减使得上游预训练中更难获得高精确度，但作者的主要目标是高质量的迁移。

3.2 通过删除[class] token来节省内存

我们都知道最大的VIT模型使用14×14个patch, 224×224的图像产生了256个visual “tokens”，每个token对应一个图像patch。在此之上，ViT模型有一个额外的[class] token，它用于产生最终的全局表征，最终使得toekn的总数达到256+1=257个。

对于ViT模型，当前的TPU硬件将token维度设置为128的倍数，这可能导致高达50%的内存开销。为了克服这个问题，作者研究了使用额外的[class] token的替代方法。评估了全局平均池(GAP)和多头注意力池(MAP)来聚合来自所有patch token的表示。将MAP中的heads数设置为与模型其他部分中的注意heads数相等。为了进一步简化head设计，在最终预测层之前去掉了最终的非线性投影。

 

为了选择最好的head，作者对[class] token和GAP/MAP head进行并行比较。结果上图所示。作者发现所有heads的表现都是相似的，而由于前面提到的padding考虑，GAP和MAP的内存效率更高。作者还观察到非线性映射可以移除。因此，作者选择了MAP  head，因为它是最具表现力和泛化性。

3.3 Scaling up data

在本研究中作者使用了专有的JFT-3B数据集，这是JFT-300M数据集的一个更大的版本，JFT-300M数据集在以前的大规模计算机视觉模型研究中使用过。该数据集由近30亿张图像组成，通过半自动管道标注了约30k个标签的类层次结构。因此，数据和相关的标签是有噪声的。这里忽略标签的层次方面，只使用分配的标签作为目标，通过sigmoid交叉熵损失进行多标签分类。

 

如图，左边显示了整个线性10-shot的ImageNet性能评估。作者观察到，JFT-3B甚至在JFT-300M未完全训练一个epoch之前就产生了一个更好的模型。因此，JFT-300M在small B/32模型以及larger L/16上过拟合并不是改善的唯一原因。

作者将模型微调到完整的ImageNet数据集(右)，并确认这些改进转移到完整的微调设置。总体而言，无论是小型还是大型模型，数据集的改变提高了ImageNet的1%左右。JFT-300M和JFT-3B的训练行为除了表现上的改善外，还有相似之处。最重要的是，JFT-3B可以让我们在不担心过拟合和正则化的情况下进一步扩大规模。

3.4 Memory-efficient optimizers

当训练大型模型时，模型参数所需的存储成为瓶颈。本文设计最大的模型ViT-G，大约有20亿个参数，占用8GiB的显存。更糟糕的是，通常用于训练transformer的Adam优化器为每个参数存储2个额外的浮点标量，这导致了额外的2倍开销(额外的16GiB)。为了解决Adam优化器带来的开销，作者做了2个修改：

Adam with half-precision momentum

通过经验观察到，半精度(bfloat16类型)的动量存储不会影响训练，对结果也没有影响。这可以将优化器的开销从原来的2倍减少到1.5倍。值得注意的是，使用半精度存储第2个动量会导致性能显著下降。

Adafactor优化器

上述Adam优化器仍然会导致较大的内存开销。因此，作者将注意力转向adfactor优化器，它使用秩1因式分解存储第2动量。从实用的角度来看，这将导致微不足道的内存开销。然而，作者并没有直接使用Adafactor优化器，而是做了以下修改:

• 1 重新将半精确引入第1个动量，而推荐的设置根本不使用第1个动量。

• 2 禁用了相对于权重Norm的学习率的缩放，在Adafactor默认开启的。

由此产生的优化器只引入了50%的内存开销。作者观察到，这2个建议的优化器的性能与最初的Adam优化器相当，甚至稍还要好一些。

3.5 Learning-rate schedule

在研究中，作者希望使用几个不同的训练时间训练一个模型，以衡量模型大小和训练时间之间的权衡。当使用线性衰减时，每个训练时间都需要从头开始运行自己的训练，这是一种低效的方法。

通过探索学习率计划来解决这个问题，类似于开始的warmup阶段，包括训练结束的cooldown阶段，在这个阶段学习率线性趋于零。在warmup和cooldown阶段之间，学习率不应该很快下降到零。这可以通过使用一个常量，或一个倒数平方根schedule的主要训练部分来实现。

 

上图述了这些选项中的几个，在大约200k、400k和500k步之后有一个cooldown时间。图的上半部分显示了这些选项及其间隔期的验证分数(越高越好)，以及2个线性时间表。

当事先知道训练时间而不打算继续训练时，线性计划仍然是可取的，所有3种选择都相当接近，其优点是允许不确定的训练，并从一次跑步中评估多个训练时间。对于每一个schedule作者优化了学习率和exact shape。作者也简单地尝试了循环学习率计划，但是它们似乎表现得更差，因此，选择了平方根倒数schedule。

3.6 选择模型的维度

ViT模型有许多参数可以用来控制模型的shape，作者参考原始版本的完整细节。简单地说，这些包括patch-size，编码器block的数量(深度)，patch embeddings和self-attention(宽度)，attention heads的数量，MLP块的维度(MLP-宽度)。

在此基础上，作者依赖XLA编译器来优化模型以提高运行时速度和内存占用。在底层，XLA使用复杂的启发式方法将模型编译为特定硬件的代码，以最佳的方式权衡内存和速度。因此，很难预测哪个模型配置适合一个设备的内存。

因此，作者运行了一个广泛的模拟，其中实例化了大量不同shape的ViT，并试图训练它们几个step。作者通过改变深度、宽度、头部和mlp宽度，但保持patch大小在14像素。通过这种方式，可以测量它们的速度，以及给定的模型是否适合设备的内存。

上图总结了这个模拟的结果。每个块对应一个模型配置，块的阴影对应其训练速度(越亮越快)。橙色块显示哪个没有任何修改原始ViT模型适合。绿色块还包括第3.2节中描述的内存节约，以及第3.4节中描述的半精度Adam。最后，蓝色块是修改的AdaFactor优化器。通过修改和实验能够适应深度最多为100个编码器块的thin ViT模型。

 

最初的Vision Transformer 结论是伸缩所有方面是最有效的(深度、宽度、mlp宽度和patch大小)同时以相似的数量。遵循这一建议，并在相应的内存容量限制下为ViT-g和ViT-G选择形状，并在表中进行总结。