论文地址为：

https://arxiv.org/pdf/2106.04560.pdf

一、简介

视觉 Transformer(ViT) 等基于注意力的神经网络最近在许多计算机视觉基准测试中取得了最先进的结果。比例是获得出色结果的主要因素，因此，了解模型的 scaling 属性是有效设计的关键。

虽然已经研究了扩展 Transformer 语言模型的规律，但尚不清楚 Vision Transformers 如何扩展。

为了解决这个问题，作者向上和向下扩展 ViT 模型和数据，并描述错误率、数据和计算之间的关系。

在此过程中，作者改进了 ViT 的架构和训练，减少了内存消耗并提高了结果模型的准确性。

结果，作者成功地训练了一个具有 20 亿个参数的 ViT 模型，该模型在 ImageNet 上达到了 90.45% 的 Top-1 准确率。该模型在小样本学习上也表现良好，例如，在 ImageNet 上每类只有 10 个 examples 的情况下可以达到 84.86% 的 Top-1 准确率。

在这个过程中，研究者创建了一个改进的大规模训练方案，探索了训练超参数以及发现微妙的选择，大幅改善小样本迁移性能。具体来说，研究者发现非常强的 $L_2$ 正则化，仅应用于最终的线性预测层，导致学习到的视觉表征具有很强的小样本学习能力。

例如，在 ImageNet 数据集（有 1000 个类）上，每个类只有一个示例，该研究的最佳模型达到 69.52% 的准确率；如果每个类有 10 个示例，准确率达到了 84.86%。此外，该研究大大减少了提出的原始 ViT 模型的内存占用，通过特定于硬件的体系架构更改和不同的优化器来实现这一点。

结果表明，该研究训练了一个具有 20 亿个参数的模型，在 ImageNet 数据集上达到了新的 SOTA 性能 90.45% 的准确率。可以看到，在 SOTA 基准排行榜上，ViT-G/14 模型的 Top-1 准确率已经超越了谷歌之前提出的 Meta Pseduo Labels 模型。

二、如何做到的？

研究者展示了对 ViT 模型及训练的改进，这些改进大多数易于实现，并显著提升了内存使用率和模型质量。如此一来，研究者可以单独使用数据并行化训练 ViT-G/14 模型，并在单个 TPUv3 上实现整个模型拟合。

三、扩展数据

该研究使用专有的 JFT-3B 数据集，它是 JFT-300M 数据集的更大规模版本，在之前的许多大型计算机视觉模型工作中使用过 [31, 18, 11]。该数据集由近 30 亿张图像组成，标注为包含 30K 类别的层级类别，并且由于采用半自动标注，所以标注是有噪音的。所有的测试数据也均从 JFT-3B 中移除。

下图展示了数据集从 JFT-300M 到 JFT-3B 过程中对模型性能的影响。可以观察到，更大规模的 JFT-3B 数据集可以得到更好的模型，所以 JFT-300M 数据集的过拟合并不是实现性能提升的唯一原因。

四、「head」 的解耦权重衰减

权重衰减对低数据情况下的模型自适应具有重大影响。

研究者在中等规模程度上研究了这一现象，并发现可以从模型中最终线性层（「head」）和剩余权重（「body」）的权重衰减强度解耦中获益。

下图展示了这一效果。研究者在 JFT-300M 上训练了一个 ViT-B/32 模型，每个单元格对应不同 head/body 权重衰减值的性能。

他们观察到的有趣的一点是：尽管提升了迁移性能，但 head 中高权重衰减却降低了预训练（上游）任务的性能。

五、通过移除 [class] token 节省内存

对于 ViT 模型，当前的 TPU 硬件将 token 维数填充为 128 的倍数，这可能导致高度 50% 的内存开销。

为了解决这一问题，研究者尝试探索「使用额外 [class] token」的替代方法。

具体地，他们对全局平均池化（ GAP）和多头注意力池化（MAP）进行评估以聚合来自所有 patch token 的表示，并将 MAP 中 head 的数量与模型其他部分中注意力 head 的数量设为相同。

为了进一步简化 head 设计，研究者原始 ViT 论文中出现的、最终预测层之前的最终非线性映射。

为了选择最佳 head，研究者对 [class] token 和 GAP / MAP head 进行了并排比较，结果如上图（右）所示。他们发现，所有 head 的表现类似，同时 GAP 和 MAP 由于进行了填充（padding）考虑，因而具备更高的内存效率。

此外，非线性映射还可以安全地进行移除。因此，研究者选择了 MAP head，这是因为其表现力最强，并且能够生成最统一的架构。

六、实验结果

在接下来的实验中，研究者在多达 30 亿个弱标记图像上训练了几个 ViT 模型。

研究者改变架构大小、训练图像的数量和训练持续时间。所有模型都在 TPUv3 上训练，因此总计算量是以 TPUv3 核每天（core-days）数来衡量的。

为了评估由模型学习的表征的质量，研究者进行了一下测量：

（i）通过在固定权值上训练线性分类器来进行小样本迁移；

（ii）通过对所有数据的整个模型进行微调，将其传递给多个基准测试任务。

6.1 将计算、模型和数据一起扩展

下图显示了在 ImageNet 上的 10-shot 线性评估和微调评估，并进行了高层级的观测。

首先，将计算、模型和数据一起扩展可以提高表征质量。

第二，模型尺寸会影响表征质量。

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

下图显示了在预训练期间「可见」的图像总数（批大小乘以 step 数）的表征质量。

除了在公共验证集上进行 ImageNet 微调和线性 10-shot 结果外，研究者还报告了在 ImageNet-v2 测试集上的 ImageNet 微调模型的结果，作为鲁棒泛化的指标。

下图中展示了对 30 亿张图像进行预训练的三个 ViT 模型。

实验可得，有足够的数据，以较少的步骤训练一个更大的模型是可取的。

6.2 ViT-G/14 结果

该研究训练了一个大型的视觉 Transformer，ViT-G/14，它包含近 20 亿个参数。

实验评估了 ViT-G/14 模型在一系列下游任务中的应用，并将其与 SOTA 结果进行了比较。研究者在 ImaegNet 上进行微调，并报告 ImageNet、ImageNet-v2、ReaL 和 ObjectNet 的准确率。此外，该研究还报告了在 VTAB-1k 基准上的迁移学习效果，该基准包括 19 个不同的任务。

显示了在 ImageNet 上进行小样本学习（few-shot learning）结果。由结果可得，ViT-G/14 比以前最好的 ViT-H/14 模型表现优异（超过 5%），达到 84.86% 的准确率，每类 10 个例子。

下表展示了其他基准上的结果。

ViT-G/14 在 ImageNet 数据集上实现了 90.45% 的 Top-1 准确率，成为新的 SOTA。

此外，在 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%。