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

论文链接：https://arxiv.org/pdf/2010.11929 

代码链接：https://github.com/google-research/vision_transformer 

摘要

尽管 Transformer 架构已成为自然语言处理任务的事实标准，但其在计算机视觉领域的应用仍然有限。在视觉领域，注意力机制要么与卷积神经网络结合使用，要么用于替换卷积神经网络的某些组件，同时保持其整体结构不变。我们证明，这种对卷积神经网络的依赖并非必要，直接将纯 Transformer 应用于图像块序列也能在图像分类任务中表现出色。当使用大量数据进行预训练并迁移到多个中小型图像识别基准数据集（ImageNet、CIFAR-100、VTAB等）时，Vision Transformer（ViT）与最先进的卷积神经网络相比取得了优异的成绩，同时训练所需的计算资源也显著减少。

1.介绍

基于自注意力机制的架构，特别是 Transformer 模型，已成为自然语言处理（NLP）领域的首选模型。主流方法是在大型文本语料库上进行预训练，然后在规模较小的特定任务数据集上进行微调。得益于 Transformer 模型的计算效率和可扩展性，训练规模空前的模型已成为可能，其参数量已超过100B。尽管模型和数据集规模不断增长，但性能仍未出现饱和的迹象。

然而，在计算机视觉领域，卷积架构仍然占据主导地位。受自然语言处理（NLP）领域成功经验的启发，许多研究尝试将类卷积神经网络（CNN）架构与自注意力机制相结合，其中一些研究甚至完全用自注意力机制替代了卷积层。尽管后一种模型在理论上效率很高，但由于使用了专门的注意力模式，它们尚未能在现代硬件加速器上有效扩展。因此，在大规模图像识别领域，经典的类残差网络（ResNet）架构仍然是最先进的。

受 Transformer 在自然语言处理（NLP）领域扩展性成功的启发，我们尝试将标准的 Transformer 模型直接应用于图像，并尽可能减少修改。为此，我们将图像分割成多个图像块，并将这些图像块的线性嵌入序列作为 Transformer 的输入。图像块的处理方式与NLP应用中的 token 相同。我们采用有监督学习的方式训练图像分类模型。

当使用诸如 ImageNet 之类的中等规模数据集进行训练，且未采用强正则化时，这些模型的准确率仅为几个百分点，比规模相近的 ResNet 模型略低。这种看似令人失望的结果其实在意料之中：Transformer 模型缺乏 CNN 模型固有的一些归纳偏置，例如平移等变性和局部性。因此在数据量不足的情况下，其泛化能力较差。

然而，如果模型在更大的数据集（1400万至3亿张图像）上进行训练，情况则有所不同。我们发现，大规模训练优于归纳偏置。我们的 Vision Transformer（ViT）在足够大的数据集上进行预训练并迁移到数据点较少的任务时，取得了优异的成绩。当在公开的 ImageNet-21k 数据集或内部的 JFT-300M 数据集上进行预训练时，ViT 在多个图像识别基准测试中达到或超过了当前最佳水平。具体而言，最佳模型在 ImageNet 上的准确率达到88.55%，在 ImageNet-Real 数据集上达到 90.72%，在 CIFAR-100 数据集上达到 94.55%，在包含 19 个任务的 VTAB 测试套件上达到 77.63%。

2.相关工作

Vaswani et al. (2017) 提出了 Transformer 模型用于机器翻译，此后 Transformer 已成为许多自然语言处理任务中最先进的方法。大型 Transformer 模型通常在大型语料库上进行预训练，然后针对具体任务进行微调：BERT 使用去噪自监督预训练任务，而 GPT 系列则使用语言建模作为其预训练任务。

将自注意力机制直接应用于图像处理，需要每个像素关注其他所有像素。由于像素数量的平方级成本，这种方法无法扩展到实际的输入图像规模。因此，为了将 Transformer 应用于图像处理，过去人们尝试了多种近似方法。Parmar et al. (2018) 仅在每个 qeury 像素的局部邻域内应用自注意力机制，而不是全局应用。这种局部多头点积自注意力模块可以完全替代卷积运算。在另一项研究中，Sparse Transformer 采用了可扩展的全局自注意力近似方法，使其能够应用于图像处理。另一种扩展注意力机制的方法是将其应用于不同大小的块中，极端情况下甚至可以仅沿单个轴应用。许多此类专门的注意力架构在计算机视觉任务中展现出良好的效果，但需要在硬件加速器上进行复杂的工程设计才能高效实现。

与我们模型最相关的是 Cordonnier et al. (2020) 的模型，该模型从输入图像中提取 2×2 大小的图像块，并在此基础上应用完全自注意力机制。该模型与 ViT 非常相似，但我们的工作更进一步，证明了大规模预训练可以使普通 Transformer 模型与最先进的 CNN 模型相媲美（甚至更胜一筹）。此外， Cordonnier et al. (2020) 使用了 2×2 像素的小图像块，这使得该模型仅适用于小分辨率图像，而我们的模型也能处理中等分辨率图像。

人们对将卷积神经网络 (CNN) 与自注意力机制相结合也产生了浓厚的兴趣，例如通过增强特征图进行图像分类，或者通过使用自注意力机制进一步处理 CNN 的输出，例如用于目标检测、视频处理、图像分类、无监督目标发现或统一的文本视觉任务。

另一个近期提出的相关模型是 image GPT (iGPT)，它在降低图像分辨率和色彩空间后，将 Transformer 应用于图像像素。该模型以无监督的方式作为生成模型进行训练，生成的表示可以进行微调或线性探测以提升分类性能，在 ImageNet 数据集上达到了 72% 的最高准确率。

我们的工作进一步丰富了探索比标准 ImageNet 数据集更大规模图像识别的论文集。利用额外的数据源，我们在标准基准测试中取得了最先进的结果。此外，Sun et al. (2017) 研究了 CNN 性能如何随数据集规模扩展，Kolesnikov et al. (2020) 和 Djolonga et al. (2020) 则对 CNN 在 ImageNet-21k 和 JFT-300M 等大规模数据集上的迁移学习进行了实证探索。我们也关注后两个数据集，但训练的是 Transformer 模型，而不是先前工作中使用的基于 ResNet 的模型。

3.METHOD

在模型设计方面，我们尽可能地遵循原始的Transformer模型。这种有意简化的设置的优势在于，可扩展的 NLP Transformer 架构及其高效实现几乎可以开箱即用。

3.1 VISION TRANSFORMER (VIT)

图 1 展示了模型的概览。标准的 Transformer 接收一维 token 嵌入序列作为输入。为了处理二维图像，我们将图像 $\textbf x ∈ \mathbb R^{H×W×C}$ 重塑为一系列展平的二维图像块 $\textbf x_p ∈ \mathbb R^{N×(P^2·C)}$，其中 $(H, W)$ 是原始图像的分辨率，$C$ 是通道数，$(P, P)$ 是每个图像块的分辨率，$N = HW/P^2$ 是最终得到的图像块数量，同时也是 Transformer 的有效输入序列长度。Transformer 在所有层中都使用恒定的潜在向量大小 $D$，因此我们将图像块展平，并使用可训练的线性投影（公式 1）映射到 $D$ 维空间。我们将此投影的输出称为图像块嵌入。

与 BERT 的 $[class]$ token 类似，我们在嵌入图像块序列前添加一个可学习的嵌入向量 $(\textbf z^0_0 = x_{class})$，其在 Transformer 编码器输出端的状态 ($\textbf z^0_L$) 即为图像表示 $\textbf y$（公式 4）。在预训练和微调过程中，$\textbf z^0_L$ 都连接了一个分类头。预训练时，分类头由一个具有一个隐藏层的多层感知器 (MLP) 实现；微调时，分类头则由一个线性层实现。

为了保留位置信息，我们在图像块嵌入中添加了位置嵌入。我们使用标准的可学习的一维位置嵌入，因为我们没有观察到使用更高级的二维感知位置嵌入能带来显著的性能提升（附录 D.4）。最终得到的嵌入向量序列将作为编码器的输入。

Transformer 编码器由交替的多头自注意力（MSA，参见附录A）层和多层感知器（MLP）模块组成（公式2、3）。每个模块之前应用层归一化（LN），每个模块之后应用残差连接。MLP包含两层，并具有 GELU 非线性激活函数。

Inductive bias。我们注意到，Vision Transformer 的图像特定归纳偏置远少于 CNN。在 CNN 中，局部性、二维邻域结构和平移等变性被嵌入到整个模型的每一层中。而在 ViT 中，只有 MLP 层是局部的和平移等变的，自注意力层是全局的。二维邻域结构的使用非常有限：仅在模型初始化时用于将图像切割成块，以及在微调阶段用于调整不同分辨率图像的位置嵌入（如下所述）。除此之外，初始化时的位置嵌入不包含任何关于图像块二维位置的信息，所有图像块之间的空间关系都需要从头开始学习。

Hybrid Architecture。除了原始图像块之外，输入序列还可以由卷积神经网络（CNN）的特征图构成。在该混合模型中，图像块嵌入投影 $\textbf E$（公式1）应用于从 CNN 特征图中提取的图像块。特殊情况下，图像块的空间大小可以为 1x1，这意味着输入序列是通过简单地展平特征图的空间维度并投影到Transformer维度而获得的。分类输入嵌入和位置嵌入的添加方式如上所述。

3.2 FINE-TUNING AND HIGHER RESOLUTION

通常，我们会在大型数据集上预训练 ViT，然后针对（较小的）下游任务进行微调。为此，我们会移除预训练的预测头，并添加一个初始值为零的 $D × K$ 前馈层，其中 $K$ 为下游类别数。通常，在比预训练更高的分辨率下进行微调是有益的。当输入更高分辨率的图像时，我们保持图像块大小不变，从而获得更大的有效序列长度。Vision Transformer 可以处理任意长度的序列（在内存限制范围内），但是，预训练的位置嵌入可能不再有意义。因此，我们会根据预训练的位置嵌入在原始图像中的位置对其进行二维插值。需要注意的是，分辨率调整和图像块提取是唯一需要手动将关于图像二维结构的归纳偏置注入 Vision Transformer 的步骤。

4.EXPERIMENTS