Scaling Vision Transformers to 22 Billion Parameters

主要贡献

• Vision Transformer（ViT）的大规模扩展：尽管Transformer架构在自然语言处理（NLP）领域取得了巨大成功，但在计算机视觉（CV）领域，尤其是图像和视频建模方面，尚未实现与语言模型相当的规模扩展。论文提出了ViT-22B，这是一个具有220亿参数的Vision Transformer模型，是目前最大的密集ViT模型。

• 训练稳定性和效率：在扩展ViT模型到22亿参数的过程中，作者遇到了训练不稳定性的问题。为了解决这个问题，他们提出了一些架构上的改进，如并行层、查询/键（QK）归一化以及省略某些偏置项，这些改进使得模型能够在保持训练稳定性的同时，实现高效的并行训练。

• 下游任务的性能提升：通过在各种下游任务上的广泛实验，论文展示了ViT-22B在图像分类、语义分割、单目深度估计和视频分类等任务上的性能提升。特别是，即使在作为冻结特征提取器的情况下，ViT-22B也能在ImageNet上达到89.5%的准确率。

• 模型的公平性、鲁棒性和人类视觉对齐：论文还探讨了模型规模增加对模型公平性、鲁棒性和与人类视觉感知对齐的影响。研究发现，随着模型规模的增加，ViT-22B在这些方面的表现有所改善，例如在形状/纹理偏差方面达到了前所未有的87%。

• 模型的可解释性和透明度：论文通过使用特征归因分析方法，如集成梯度（Integrated Gradients），来理解ViT-22B如何做出预测，这有助于提高模型的透明度和可解释性。

Method

论文通过以下几个关键步骤解决了在计算机视觉领域扩展Vision Transformer（ViT）模型的问题：

1. 架构改进：

• 并行层：将注意力和多层感知器（MLP）块并行处理，而不是顺序处理，以实现额外的并行化。

  $\begin{array}{rl}& y^{\prime }=\text{LayerNorm}(x),\\ & y=x+\text{MLP}(y^{\prime })+\text{Attention}(y^{\prime }).\end{array}$

• QK-Norm：在计算注意力权重之前，对查询（Q）和键（K）进行LayerNorm归一化，以防止注意力权重的发散，从而提高训练稳定性。

  $\begin{array}{r}1\\ \mathbf{softmax}\left[\frac{1}{\sqrt{d}}\mathrm{LN}(X{W}^Q)(\mathrm{LN}(X{W}^K){)}^T\right],\end{array}$

• 省略偏置项：在QKV投影和LayerNorms中移除偏置项，以提高硬件利用率，同时保持或提高模型质量。

1. 训练基础设施和效率：

• 使用JAX和FLAX库实现ViT-22B，利用模型和数据并行来处理大规模模型。

• 开发异步并行线性操作，以最大化矩阵乘法单元的利用率，同时最小化通信开销。

• 参数分片，允许在多个设备上分布模型参数，以适应更大的模型和批量大小。

3. 实验和评估：

• 在大规模数据集（如扩展的JFT数据集）上训练ViT-22B，并在多种下游任务上评估其性能。

• 使用线性探测、锁定图像调整（Locked-image Tuning）和冻结图像特征提取器等技术，将ViT-22B应用于图像分类、语义分割、单目深度估计和视频分类任务。

• 分析模型在公平性、人类视觉对齐、鲁棒性、可靠性和校准方面的表现。

4. 模型蒸馏：

• 通过知识蒸馏技术，将ViT-22B的知识压缩到较小的ViT模型中，以提高模型的可用性和部署效率。
• 通过这些方法，论文成功地训练了一个规模达到220亿参数的ViT模型，并在多个视觉任务上展示了其有效性。此外，论文还探讨了模型规模增加带来的其他好处，如改善公平性、提高与人类视觉感知的对齐程度以及增强模型的鲁棒性。