[论文解读] Scaling Vision Transformers

Xiaohua Zhai, А. И. Колесников+2 more|arXiv (Cornell University)|Jun 8, 2021

Advanced Neural Network Applications被引用 15

一句话总结

本论文通过在100万至30亿张图像的数据集上训练参数量从500万到20亿不等的Vision Transformers（ViT）模型，并使用最多10,000个TPUv3核心日的计算资源，研究了ViT的缩放规律。论文引入了架构和训练方法的改进，特别是对最终层施加强L2正则化，使一个20亿参数的ViT在ImageNet上实现了90.45%的top-1准确率，创下新的最先进水平，且在每类仅10个样本的情况下也达到了84.86%的准确率。

ABSTRACT

Attention-based neural networks such as the Vision Transformer (ViT) have recently attained state-of-the-art results on many computer vision benchmarks. Scale is a primary ingredient in attaining excellent results, therefore, understanding a model's scaling properties is a key to designing future generations effectively. While the laws for scaling Transformer language models have been studied, it is unknown how Vision Transformers scale. To address this, we scale ViT models and data, both up and down, and characterize the relationships between error rate, data, and compute. Along the way, we refine the architecture and training of ViT, reducing memory consumption and increasing accuracy of the resulting models. As a result, we successfully train a ViT model with two billion parameters, which attains a new state-of-the-art on ImageNet of 90.45% top-1 accuracy. The model also performs well for few-shot transfer, for example, reaching 84.86% top-1 accuracy on ImageNet with only 10 examples per class.

研究动机与目标

• 理解Vision Transformers在模型大小、数据量和计算资源方面的缩放规律，特别是与NLP中的Transformer进行比较。
• 识别能够提升表示质量与少样本迁移性能的架构和训练配方改进。
• 训练一个具有20亿参数的大规模ViT模型，并在ImageNet和少样本迁移基准上评估其性能。
• 刻画性能-计算前沿，并识别在缩放过程中（模型大小、数据、计算）的瓶颈。
• 通过硬件感知设计和优化器选择，降低ViT的内存消耗，同时保持或提升准确率。

提出的方法

• 在TPUv3硬件上，将ViT模型从500万参数扩展到20亿参数，在ImageNet-21k和最多30亿张弱标签图像上进行训练。
• 通过在ImageNet、CIFAR-100、Oxford IIIT Pets和Caltech-UCSD Birds上的线性10-shot评估和完整微调来衡量表示质量。
• 仅对最终的线性分类层施加强L2正则化，以提升少样本迁移性能。
• 通过调整学习率、冷却步数和权重衰减等超参数，优化训练过程，以提升收敛性和泛化能力。
• 通过硬件特定的架构修改和改进的优化器，降低内存消耗，从而实现20亿参数ViT模型的训练。
• 使用饱和幂律模型来近似不同计算、模型和数据规模下的性能-计算前沿。

实验结果

研究问题

• RQ1 当增加模型大小、训练数据和计算预算时，ViT模型的准确率如何变化？
• RQ2 在不同配置下，ViT缩放的主要瓶颈是模型容量、数据集大小还是计算资源？
• RQ3 对最终层施加强L2正则化如何影响ViT模型的少样本迁移性能？
• RQ4 架构和训练配方的改进是否能降低内存消耗，同时支持更大规模的ViT模型？
• RQ5 在NLP Transformer中观察到的缩放规律在视觉领域在多大程度上具有普适性？

主要发现

• 一个20亿参数的Vision Transformer在ImageNet上实现了90.45%的top-1准确率，创下新的最先进水平。
• 在每类仅10个样本的情况下，ViT-G模型在ImageNet上的10-shot线性评估中达到84.86%的top-1准确率。
• 仅使用每类一个样本时，模型准确率达到69.52%，展现出强大的少样本迁移能力。
• 更大的模型在训练数据超过10亿张后仍能从更多数据中显著获益；3000万至3亿张图像不足以使最大模型达到饱和。
• 当训练时间过长时，较小的模型会迅速达到饱和并偏离性能-计算前沿，表明存在过拟合或计算效率低下问题。
• 通过硬件特定的架构修改和优化器改进，内存消耗显著降低，从而实现了20亿参数ViT模型的训练。