https://openreview.net/pdf?id=YicbFdNTTy

AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

Abstract:

Transformer 架构早已在自然语言处理任务中得到广泛应用，但在计算机视觉领域中仍然受到限制。在计算机视觉领域，注意力要么与卷积网络结合使用，要么用来代替卷积网络的某些组件，同时保持其整体架构不变。

该研究表明，对 CNN 的依赖不是必需的，当直接应用于图像块序列时，transformer 也能很好地执行图像分类任务。该研究基于大量数据进行模型预训练，并迁移至多个图像识别基准数据集（ImageNet、CIFAR-100、VTAB 等），结果表明 Vision Transformer（ViT）模型可以获得与当前最优卷积网络相媲美的结果，而其训练所需的计算资源大大减少。

1 Introduction

1.1 NLP 领域中的 Transformer VS 计算机视觉领域中的 CNN

基于自注意力的架构，尤其 Transformer，已经成为 NLP 领域的首选模型。该主流方法基于大型文本语料库进行预训练，然后针对较小的任务特定数据集进行微调。由于 Transformer 的计算效率和可扩展性，基于它甚至可以训练出参数超过 100B 的模型。随着模型和数据集的增长，性能仍然没有饱和的迹象。

然而，在计算机视觉中，卷积架构仍然占主导地位。受 NLP 成功的启发，多项计算机视觉研究尝试将类 CNN 架构与自注意力相结合，有的甚至完全代替了卷积。后者虽然在理论上有效，但由于其使用了专门的注意力模式，因此尚未在现代硬件加速器上有效地扩展。因此，在大规模图像识别任务中，经典的类 ResNet 架构仍然是最先进的。

1.2 Transformer 向视觉领域的跨界融合

受到 NLP 领域中 Transformer 缩放成功的启发，这项研究尝试将标准 Transformer 直接应用于图像，并尽可能减少修改。为此，该研究将图像分割成多个图像块（patch），并将这些图像块的线性嵌入序列作为 Transformer 的输入。然后用 NLP 领域中处理 token 的方式处理图像块，并以监督的方式训练图像分类模型。

在中等规模的数据集（如 ImageNet）上训练时，这样的模型产生的结果并不理想，准确率比同等大小的 ResNet 低几个百分点。这个看似令人沮丧的结果是可以预料的：Transformer 缺少一些 CNN 固有的归纳偏置，例如平移同变性和局部性，因此在数据量不足的情况下进行训练后，Transformer 不能很好地泛化。

但是，如果在大型数据集（14M-300M 张图像）上训练模型，则情况大为不同。该研究发现大规模训练胜过归纳偏置。在足够大的数据规模上进行预训练并迁移到数据点较少的任务时，Transformer 可以获得出色的结果。

该研究提出的 Vision Transformer 在 JFT-300M 数据集上进行预训练，在多个图像识别基准上接近或超过了 SOTA 水平，在 ImageNet 上达到了 88.36% 的准确率，在 ImageNet ReaL 上达到了 90.77% 的准确率，在 CIFAR-100 上达到了 94.55% 的准确率，在 VTAB 基准 19 个任务中达到了 77.16% 的准确率。

2 模型和方法

研究者尽可能地遵循原始 Transformer 的设计。这种故意为之的简单设置具有以下优势，即可扩展 NLP Transformer 架构和相应的高效实现几乎可以实现开箱即用。研究者想要证明，当进行适当地扩展时，该方法足以超越当前最优的卷积神经网络。

2.1 Vision Transformer（ViT）

该研究提出的 Vision Transformer 架构遵循原版 Transformer 架构。下图 1 为模型架构图。

标准 Transformer 接收 1D 序列的 token 嵌入为输入。为了处理 2D 图像，研究者将图像 x ∈ R^H×W×C 变形为一系列的扁平化 2D patch x_p ∈ R^N×(P^2 ·C)，其中 (H, W) 表示原始图像的分辨率，(P, P) 表示每个图像 patch 的分辨率。然后，N = HW/P^2 成为 Vision Transformer 的有效序列长度。

Vision Transformer 在所有层使用相同的宽度，所以一个可训练的线性投影将每个向量化 patch 映射到模型维度 D 上（公式 1），相应的输出被称为 patch 嵌入。

与 BERT 的 [class] token 类似，研究者在一系列嵌入 patch （z_0^0 = x_class）之前预先添加了一个可学习嵌入，它在 Transformer 编码器（z_0^L ）输出中的状态可以作为图像表示 y（公式 4）。在预训练和微调阶段，分类头（head）依附于 z_L^0。

位置嵌入被添加到 patch 嵌入中以保留位置信息。研究者尝试了位置嵌入的不同 2D 感知变体，但与标准 1D 位置嵌入相比并没有显著的增益。所以，编码器以联合嵌入为输入。

Transformer 编码器由多个交互层的多头自注意力（MSA）和 MLP 块组成（公式 2、3）。每个块之前应用 Layernorm（LN），而残差连接在每个块之后应用。MLP 包含两个呈现 GELU 非线性的层。

作为将图像分割成 patch 的一种替代方案，输出序列可以通过 ResNet 的中间特征图来形成。在这个混合模型中，patch 嵌入投影（公式 1）被早期阶段的 ResNet 取代。ResNet 的其中一个中间 2D 特征图被扁平化处理成一个序列，映射到 Transformer 维度，然后馈入并作为 Transformer 的输入序列。最后，如上文所述，将分类输入嵌入和位置嵌入添加到 Transformer 输入中。

2.2 微调和更高分辨率

研究者在大型数据集上预训练 ViT 模型，并针对更小规模的下游任务对模型进行微调。为此，研究者移除了预训练预测头，并添加了一个零初始化的 D × K 前馈层，其中 K 表示下游类的数量。与预训练相比，在更高分辨率时进行微调通常更有益处。当馈入更高分辨率的图像时，研究者保持 patch 大小不变，从而得到更大的有效序列长度。

ViT 模型可以处理任意序列长度（取决于内存约束），但预训练位置嵌入或许不再具有意义。所以，研究者根据预训练位置嵌入在原始图像中的位置，对它们进行 2D 插值操作。需要注意的是，只有在分辨率调整和 patch 提取中，才能将 2D 图像的归纳偏置手动注入到 ViT 模型中。

3 实验

3.1 模型

该研究进行了大量实验，并使用了多个 ViT 模型变体，参见下表 1：

3.2 SOTA 模型的性能对比

研究者首先将最大的 ViT 模型（在 JFT-300M 数据集上预训练的 ViT-H/14 和 ViT-L/16）与 SOTA CNN 模型进行对比，结果参见下表 2。

表 2：ViT 模型与 SOTA 模型在流行图像分类基准数据集上的性能对比。

从上表中可以看出，规模较小的 ViT-L/16 模型在所有数据集上的性能堪比或者超过 BiT-L，同时它需要的算力也少得多。较大的 ViTH-14 模型进一步提升了性能，尤其在更具挑战性的数据集上，如 ImageNet、CIFAR-100 和 VTAB。ViTH-14 模型在所有数据集上的性能匹配或超过 SOTA，甚至在某些情况下大幅超过 SOTA 模型（如在 CIFAR-100 数据集上的性能高出 1%）。在 ImageNet 数据集上，ViT 模型的性能比 Noisy Student 低了大约 0.1%，不过在具备更干净 ReaL 标签的 ImageNet 数据集上，ViT 的性能超过 SOTA 模型。

下图 2 将 VTAB 任务分解为多个组，并对比了 ViT 与 SOTA 方法的性能，这些方法包括 BiT、VIVI 和 S4L。

在 Natural 任务中，ViT-H/14 的性能略低于 BiT-R152x4；在 Specialized 任务中，ViT 的性能超过 BiT 等方法；而在 Structured 任务中，ViT 显著优于其他方法。

3.3 预训练数据要求

Vision Transformer 在大型 JFT-300M 数据集上进行预训练后表现出了优秀的性能。在 ViT 的归纳偏置少于 ResNet 的情况下，数据集规模的重要性几何呢？该研究进行了一些实验。

首先，在规模逐渐增加的数据集（ImageNet、ImageNet-21k 和 JFT300M）上预训练 ViT 模型。下图 3 展示了模型在 ImageNet 数据集上的性能：

下表 3 展示了模型在 ImageNet、ImageNet-21k 和 JFT300M 数据集上的性能对比情况。在前两个规模较小的数据集上，ViT-Large 模型的性能不如 ViT-Base，但在规模较大的 JFT300M 数据集上，大模型展现出了优势。这说明，随着数据集规模的增大，较大的 ViT 模型变体优于较小的模型。

其次，研究者在 JFT300M 数据集的 9M、30M 和 90M 随机子集以及完整数据集上进行了模型训练。结果参见下图 4：

从图中可以看到，在较小的数据集和相同的计算成本下，Vision Transformer 比 ResNet 更加过拟合。该结果强化了这一认知：卷积归纳偏置对于规模较小的数据集较为有用，但对于较大的数据集而言，学习相关模式就足够了，甚至更加有效。

3.4 可扩展性研究

研究人员对不同模型执行了受控可扩展性研究（scaling study）。下图 5 展示了模型在不同预训练计算成本情况下的迁移性能：

实验结果表明：

Vision Transformer 在性能 / 算力权衡中显著优于 ResNet。

混合模型在较小计算成本的情况下略优于 ViT，但在计算成本较高时，这一现象消失。该结果令人吃惊。

Vision Transformer 在实验尝试的算力范围内似乎并未饱和，未来可以进行更多可扩展性研究。

3.5 ViT 如何处理图像数据？

为了了解 ViT 处理图像数据的过程，研究者分析了其内部表示。

ViT 的第一层将扁平化后的图像块线性投影至低维空间（公式 1），下图（左）展示了学得嵌入滤波器的主要组件。投影后，将学得的位置嵌入添加至图像块表示。下图（中）展示了模型学习编码图像内的距离，表明距离越近的图像块更有可能具备更相似的位置嵌入。自注意力允许 ViT 集成整个图像的信息，即使最低层也不例外。研究者调查了 ViT 网络利用这一能力的程度。具体而言，该研究计算图像空间中的平均距离（基于注意力权重）参见下图右。「注意力距离」类似于 CNN 中的感受野大小。

ViT 模型关注与分类具备语义相关性的图像区域，参见图 6：

4 CONCLUSION

我们已经探索了Transformers 在图像识别中的直接应用。不像以前的工作在计算机视觉使用self-attention ，我们没有引入 image-specific inductive biases 的结构。相反，我们将一个图像解释为一系列patches，并使用标准的Transformer 对其进行处理，就像在 nlp 中使用的那样。这种简单但可扩展的策略在与大型数据集的预训练相结合时，效果惊人地好。因此，Vision Transformer 匹配或超过许多图像分类数据集的最新水平，同时具有相对便宜的预训练。虽然初步结果令人鼓舞，但仍存在许多挑战。一种是将 vit 应用于其他计算机视觉任务，如检测和分割。我们的研究结果，加上 carion 等人(2020)的研究结果，表明了这种方法的前景。另一个，继续探索自我监督的预训练方法。我们的初步实验表明，自我监督的预训练有所改善，但自我监督的预训练与大规模监督的预训练还有很大的差距。最后可以进一步扩大ViT的规模，随着模型尺寸的增加，性能似乎还没有饱和。