《An Image is Worth 16x16 Words》完整版翻译

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

1. 论文地址：https://arxiv.org/abs/2010.11929
2. TensorFlow代码：https://github.com/google-research/vision_transformer
3. Pytorch代码：https://github.com/lucidrains/vit-pytorch

---

Abstract

  虽然Transformer结构已经成为NLP领域的标准，但在计算机视觉领域的应用还非常有限。在视觉领域，注意力要么是结合卷积网络来使用，要么是用于替换卷积网络的某些组件，同时保持整体结构不变。我们将证明，在图像分类任务上，对于CNN的依赖不是必要的，直接将Transformer应用到图像块序列上也可以有非常好的性能。当ViT在大型数据集上进行预训练，并迁移到中小型图像识别数据集（ImageNet、CIFAR-100，VTAB等）上时，性能大幅超过现SOTA的卷积网络，并且可以大幅降低训练所需的计算资源。

1. Introduction

  基于自注意力的结构，特别是Transformer，已经成为了NLP方法的标准选择。主流方法是先在大型文本语料库上进行预训练，然后在较小的特定任务数据集上微调。得益于Transformer的计算效率和可扩展性，训练超过1000亿参数的超大模型也成为可能。随着模型大小和数据集规模的增长，目前仍然没有性能饱和的迹象。
  但是在CV领域，卷积结构仍占主导地位。受NLP成功的启发，有多项工作尝试将卷积结构和自注意力相结合，有的则替换整个卷积结构。后一种方法虽然理论上有效，但是由于使用了特定的注意力模式，目前并不能在现在硬件加速器上进行有效的扩展。因此，经典的ResNet结构在大规模图像识别任务上仍处于领先地位。
  受到Transformer在NLP成功的启发，我们实验直接将标准Transformer用于图像，并进行尽可能少的修改。为此，我们将图像分解成图像块，然后将这些图像块的线性嵌入式序列作为Transformer的输入。图像块等同于NLP中tokens（words）的概念。我们使用监督方式在图像分类任务上训练模型。
  在不使用强正则化的情况下，我们在中等大小的数据集比如ImageNet上训练ViT，模型精度相比同等参数规模的ResNet低了几个百分点。这种看似令人沮丧的结果是符合预期的：Transformer缺少一些CNN的固有性质，比如平移不变性和局部参数共享，因此当在不充分数据集下训练时，泛化性能不是很好。
  但是，当模型在更大的数据集（1.4~30亿）上训练时，情况会有所改变，我们发现大规模训练胜过CNN的归纳偏置（Inductive Bias）。我们的Vision Transformer(ViT)当在足够大的数据集上进行预训练后，然后迁移至拥有较少数据集的任务上时，可以获得极好的结果。当在公开ImageNet-21K或JFT-300M室内数据集上进行预训练，ViT在多个图像识别基准上接近甚至超越当前的SOTA。特别的，最好的模型在ImageNet上达到了88.55%的准确率，在ImageNet-ReaL达到了90.72%，在CIFAR-100上达到了94.55%，在VTAB系列的19个任务上达到了77.63%。

2. Related Work

  Transformer由Vaswani等人于2017年提出用于机器翻译，并从此成为许多NLP任务的SOTA方法。大型基于Transformer的模型通常在大型语料库上预训练，然后在手头任务上进行微调：BERT使用去噪自监督预训练任务，GPT一系列工作则使用语言模型作为其预训练任务。
  对图像应用自注意力的原始方法需要每个像素关注自身以外的所有其他像素。由于像素数量的二次方开销，无法扩展到实际输出尺寸。因此，为了将Transformer应用到图像处理任务中，过去有人尝试几种近似方法。Parmar等人将自注意力仅用于每个查询像素的局部邻域内，而不是全局区域；局部多头点积自注意力模块可以完全替代卷积。另外，稀疏Transformer对全局自注意力使用可扩展近似，以便适用于图像。另一种扩展注意力的方法是将其用在不同尺寸的block上，极端情况下只沿单个轴。许多这些定制化的自注意力结构在CV任务上呈现了令人信服的结果，但是想要在硬件加速器上实现这些方法还需要复杂的工程。
  和我们工作最相关的是Cordonnier等人的工作，他们从输入图像中提取$2\times 2$大小的图像块，然后在顶层应用自注意力，这个模型和ViT非常相似，但是我们的工作进一步阐述了大规模预训练可以使质朴的Transformer和SOTA的CNN媲美。而且，Cordonnier等人使用很小的$2\times 2$大小的图像块，使得模型仅适用于低分辨率图像，而我们的方法同时可以处理中等分辨率的图像。
  同样有很多工作聚焦于将CNN和自注意力结合，比如为图像分类增强特征图，或对CNN的输出进一步使用自注意力（目标检测、视频处理、图像分类等）。
  另一个和ViT接近的模型是图像GPT，它首先对图像的分辨率和侧才空间进行缩减，然后对图像像素使用Transformer。GPT采用无监督方式训练作为生成模型，然后结果可以通过微调或线性探索提高分类性能，最终在ImageNet上达到了72%的准确率。
  在比标准ImageNet更大规模的数据集上进行图像识别的论文持续增加，我们方法也是其中之一。使用额外的数据源使得在标准基准上达到SOTA的结果。此外，Sun等人研究了CNN性能如何随数据集规模变化，Kolesnikov和Djolonga等人则对大规模数据集比如ImageNet-21k和JFT-300M的CNN迁移学习进行了实证探索。我们也关注这两个数据集，但是使用Transformer训练而不是先前方法所使用的ResNet模型。

3. Method

  在模型设计方面，我们尽可能和原始Transformer保持一致。这种有意设计简单化的优点在于对NLP的Transformer有良好的扩展性，几乎可以实现“开箱即用”。

3.1 Vision Transformer(ViT)

  模型整体结构如图1所示。标准Transformer接收1维token embeddings序列作为输入。为了处理2维图像，我们将$\mathrm{x}\in {\mathbb{R}}^{H\times W\times C}$的图像reshape成拉平的2D图像块序列${\mathrm{x}}_p\in {\mathbb{R}}^{N\times (P^2\cdot C)}$，其中$(H,W)$是原始图像的分辨率，$C$是图像通道数，$(P,P)$是每个图像块的分辨率，$N=HW/P^2$为图像块的数量，也是Transformer有效输出序列的长度。Transformer的所有层均使用固定维度$D$的隐藏向量，所以我们使用一个可训练的线性层将图像块映射到$D$维，如式(1)所示。我们将该线性映射的输出称为patch embeddings。
  类似BERT方法中的 [class] token，我们在patch embeddings之前增加了一个可学习的embedding：${\mathrm{z}}_0^0={\mathrm{x}}_{class}$，它在Transformer编码器输出端的状态${\mathrm{z}}_L^0$作为图像的表征，如式(4)所示。分类头在预训练阶段由带有一个隐藏层的MLP实现，微调阶段时由一个线性层实现。
  position embeddings加在patch embeddings之后，用于保留位置信息。由于使用更加高级的2维 position embeddings并未带来明显的性能提升，这里我们仅使用标准的可学习的1维position embeddings。最终加入position embeddings得到的向量序列作为Transformer编码器的输入。
  Transformer编码器由包含多头自注意力（MSA）、MLP块的层组成。每个块之前应用一个Layernorm（LN）层，每个块之后有一个残差连接。MLP包含两个带有GELU的非线性层。
$${\mathrm{z}}_0=[{\mathrm{x}}_{class};x_p^1\text{E};x_p^2\text{E};\dots ;x_p^N\text{E}]+{\text{E}}_{pos}，\text{E}\in {\mathbb{R}}^{({\mathrm{P}}^2\cdot \mathrm{C})\times \mathrm{D}},\text{E}\in {\mathbb{R}}^{(\mathrm{N}+1)\times \mathrm{D}}\text{\hspace{1em}(1)}$$
$${\mathrm{z}}_{\ell }^{\prime }=\mathrm{M}\mathrm{S}\mathrm{A}(\mathrm{L}\mathrm{N}({\mathrm{z}}_{\mathrm{\ell }-1}))+{\mathrm{z}}_{\mathrm{\ell }-1},\ell =1\dots L\text{\hspace{1em}(2)}$$
$${\mathrm{z}}_{\ell }=\mathrm{M}\mathrm{L}\mathrm{P}(\mathrm{L}\mathrm{N}({\mathrm{z}}_{\mathrm{\ell }}^{\prime }))+{\mathrm{z}}_{\mathrm{\ell }}^{\prime },\ell =1\dots L\text{\hspace{1em}(3)}$$
$$\mathrm{y}=\mathrm{L}\mathrm{N}({\mathrm{z}}_L^0)\text{\hspace{1em}(4)}$$

归纳偏置
  我们注意到Vision Transformer相比CNN具有更少的针对图像的归纳偏置。在CNN中，局部参数共享、2维邻域结构和平移不变性贯穿于整个模型的每一层。在ViT中，只有MLP具有局部参数共享和平移不变性，而且自注意力层是全局的。二维邻域结构使用的非常少，仅在模型最开始裁剪图像块时和微调阶段针对不同分辨率图像调整position embedding有用到。除此以外，position embeddings在初始化时不携带任何2维位置信息，图像块之间的所有空间关系都是从头学习得到的。
混合结构
  作为原始图像块的替代，输入序列可以从CNN的特征图构建。在这种混合模型中，patch embedding映射$\text{E}$应用于从CNN特征图中提取的patch。在极端情况下，patches尺寸可以是$1\times 1$，即输入序列可以是特征图经过简单拉伸并映射到Transformer维度得到。分类输入embedding和position embeddings按照上述相同的方式加入。

3.2 Fine-Tuning and Higher Resolution

  通常我们在大型数据集上对ViT进行预训练，然后在下游小型数据集上进行微调。为此，我们移除了预训练的预测头并增加一个零初始化的$D\times K$的前向层，其中$K$是预测的类别数量。通常使用比预训练更高分辨率的图像微调是有益的。当输入更高分辨率图像时，我们保持图像块尺寸不变，这会导致更长的有效序列长度。Vision Transformer在内存允许的情况下可以处理任意长度的序列，但是预训练得到的position embeddings就不再有意义了。因此我们根据原始图像的位置对预训练的position embeddings进行2D插值。请注意，这种分辨率调整和块提取操作是Vision Transformer中唯一一处手动引入图像2维结构的点。

4. Experiment

  我们评估了ResNet、Vision Transformer(ViT)和混合结构的表达学习能力。为了理解每个模型的数据需求，我们在不同大小的数据集上进行了预训练，并在许多基准任务上进行了评估。当考虑预训练的计算量时，ViT表现得非常出色，以更少的预训练开销达到了SOTA水平。最后，我们进行了一个自监督的小实验，表明了自监督ViT在未来是具有潜力的。

4.1 Setup

数据集
  为了探索模型的可扩展性，我们使用ILSVRC-2012 ImageNet（1000类别，1300万张图像）、ImageNet-21k（2.1万类别，1.4亿万张图像）以及JFT（1.8万类别，30.3亿张图像）数据集。我们依照Kolesnikov等人，参照下游任务的测试集对预训练集进行去重。我们将在这些数据集上训练的模型迁移到一些基准任务上：原始验证集标签上和清理过的ReaL标签上的ImageNet，Oxford-IIIT Pets，以及Oxford Flowers-102。对于这些数据集，预处理遵循Kolenikov等人的方法。
  我们同样在具有19个分类任务的VTAB数据集上进行了评估。VTAB每个任务使用1000张训练图像，评估有限数据到各种任务的迁移能力。这些任务分为3个组：自然图像任务——和上述Pets、CIFAR相似，特定图像任务——医学和卫星图像，结构化图像任务——需要理解几何，比如定位。

模型变体
  如表1所示，我们在BERT所使用的模型结构基础上确定ViT配置。“Base”和“Large”直接取自BERT，“Huge”是我们增加的更大模型。在下文中，我们使用简明的注释来表示模型尺寸和输入图像块尺寸：比如ViT-L/16表示输入块尺寸为$16\times 16$的“Large”模型。注意Transformer的序列长度和图像块尺寸的平方成反比，因此图像块尺寸更小的模型计算量更大。
  对于CNN基线，我们使用ResNet，但是将Batch Norm层替换成Group Norm层，然后使用标准化卷积。这些改动可以提升迁移的性能，我们用ResNet(BiT)表示修改后的模型。对于混合模型，我们将中间层的特征图送给ViT，块尺寸为1个像素。为了实验不同长度的序列，我们(i)使用常规ResNet50中stage 4的输出(ii)移除stage4，使用stage3中同样数量的层进行替换，然后取这个扩展stage 3的输出。选项(ii)可以得到4倍长度的序列，因此对应的ViT模型计算开销更大。
训练和微调
  我们训练所有包括ResNets在内的模型，都是用Adam优化器，${\beta }_1=0.9,{\beta }_2=0.999，\mathrm{w}\mathrm{e}\mathrm{i}\mathrm{g}\mathrm{h}\mathrm{t}\_\mathrm{d}\mathrm{e}\mathrm{c}\mathrm{a}\mathrm{y}=0.1,\mathrm{B}\mathrm{A}\mathrm{T}\mathrm{C}\mathrm{H}\_\mathrm{S}\mathrm{I}\mathrm{Z}\mathrm{E}=4096$，我们发现这对于所有模型的迁移都很有效（附录D.1表明，和一般经验相反，对于ResNets训练，Adam比SGD要稍好一些。）。我们使用一个线性学习率预热和衰减，细节见附录B.1。微调时我们使用带动量的SGD，$\mathrm{B}\mathrm{A}\mathrm{T}\mathrm{C}\mathrm{H}\_\mathrm{S}\mathrm{I}\mathrm{Z}\mathrm{E}=512$，见附录B.1.1。对于表2中的ImageNet结果，我们使用高分辨进行微调：ViT-L/16使用512，ViT-H/14使用518。
评价指标
  我们报告了下游数据集小样本和微调的准确率。微调准确率体现的是每个模型在对应数据集上微调后的性能。小样本准确率是通过求解将训练图像子集的表征映射到 { − 1 , 1 } K \{-1,1\}^K { −1,1}K目标向量的最小二乘回归问题得到。该公式使得我们能够以闭环的方式获取精确解。尽管我们主要关注微调性能，有时候微调开销过大时，我们也会使用线性小样本准确率来进行快速的动态评估。

4.2 Comparison to SOTA

  我们首先拿最大的模型ViT-H/14和ViT-L/16和SOTA文献中的CNN比较。第一个比较点是Big Transfer(BiT)，使用大型ResNets进行有监督的迁移学习；第二个点是Noisy Student，使用半监督的方式在去除标签的ImageNet和JFT-300M数据集上训练的EfficientNet。目前，Noisy Student是ImageNet上的SOTA，BiT-L是其他数据集上的SOTA。所有模型都使用TPUv3训练，我们报道了每个模型预训练TPUv3-core-days的数值，该数值等于用于训练的TPUv3核心数量（每块TPUv3有2个核心）乘以训练天数的乘积。
  表2是对比实验结果。同样在JFT-300M上进行预训练，更小的ViT-L/16在所有任务上表现都优于BiT-L，而且大幅减少了所需的训练资源。更大的ViT-H/14性能有进一步提升，尤其是在诸如ImageNet、CIFAR-100、VTAB这类难度更大的数据集上。有趣的是，ViT预训练的资源开销相比之前的SOTA方法也大大减少。然而，我们注意到预训练的效率不仅和模型结构选择相关，而且和训练策略、优化器、权重衰减等因素有关。我们在4.4节提进行了一个关于不同模型结构性能和训练量的控制实验。最后，在Image-21k上预训练的ViT-L/16在绝大多数数据集上都表现的很好，而且所需预训练资源更少：使用标准的8核TPUv3大约使用30天可以完成训练。

  图2展现了BiT、VIVI（在ImageNet和Youtube上联合训练的ResNet）和S4L（在ImageNet混合监督和半监督方式训练）在VTAB任务上的性能。在Natural和Structured任务分支上，ViT-H/14比BiT-R152x4和其他方法好，在Specialized分支上和BiT-R152x4性能接近。

4.3 Pre-Training Data Requirements

  Vision Transformer当使用JFT-300M数据集预训练时性能较好。相比ResNets，ViT使用的归纳偏置更少，那么数据集大小是关键吗？对此我们进行了两组实验。
  首先，我们在尺寸递增的不同数据集上对ViT进行预训练：ImageNet、ImageNet-21k，JFT-300M。为了提升较小数据集上的性能，我们优化3种基本正则化参数——权重衰减、Dropout以及标签平滑。图3是使用不用预训练数据集，然后在ImageNet上微调的结果（其他数据集的结果见表5）。当在较小的ImageNet数据集上进行预训练时，尽管进行了正则化，但ViT-Large模型性能仍不如ViT-Base。当在ImageNet-21k上预训练时，两个模型的性能接近。只有当使用JFT-300M数据集预训练时，更大模型的优势才得以体现。图3同样展现了不同大小的BiT模型的性能：ImageNet上BiT性能优于ViT，但其他两个更大的数据集则ViT更优。

  第二，我们在9M、30M、90M大小的JFT-300M的随机子集训练了我们的模型。对于较小的子集，我们没有用其他额外的正则化操作，所有实验都使用相同的超参数进行。我们通过这种方式评估模型的本质而不是正则化的影响。我们使用早停的方式获取了训练过程中验证集最高的准确率，如图4所示。

  总的来说，从JFT和VTAB上的几组实验结果来看，少样本迁移看起来很有希望。ViT的少样本特性在未来是一个令人激动的研究方向。

4.4 Scaling Study

  我们通过比较JFT-300M数据集上的迁移性能对不同模型进行了实验。这组实验中，数据集大小不是模型的性能瓶颈，我们评估每个模型性能和预训练开销的关系。实验模型包含：预训练7个epoch的7ResNets，R50x1，R50x2，R101x1，R152x1，R152x2；加上预训练14个epoch的R152x2，R200x3；预训练7个epoch的ViT-B/32，B/15，L/32，L/16，加上预训练14个epoch的R50+ViT-L/16（为了测试混合结构，模型后面的数字不是表示图像块大小，而是ResNet骨干中的总下采样率）。
  实验结果如图5所示，每个模型的详细结果见表6。可以发现一些规律。第一，ViT在性能/计算开销的均衡方面全面由于ResNet，可以减少2~4倍的训练量达到相同的性能水平（平均5个数据集）。第二，混合模型在计算预算较少情况下，性能优于ViT，但是这个差距会随着模型增大而消失。这个结果有些令人惊讶，因为我们也许可以期望通过局部卷积特征来辅助任意尺寸的ViT。第三，ViT在尝试范围内未出现性能饱和的现象，推动未来的扩展工作。

4.5 Inspecting Vision Transformer

  为了理解ViT如何处理图像数据，我们分析了模型的内部表征。ViT的第一层将展平的图像块线性映射到一个较低纬度的空间（式1）。图7左展现了学习得到的嵌入滤波器的主要作用，似乎是每个图像块中精细结构的低维基础函数。
  线性映射之后，图像块特征附加了一个通过学习得到的position embedding。图7中展现了模型将图像中的距离使用position embedding相似性进行编码，即靠的更近的图像块倾向于有更相似的position embedding。而且，同一行或同一列的图像块也有相似的position embedding。最后，在更大的网格中有时会出现明显的正弦结构。position embedding可以学习到表达2D图像的拓扑结构，这也解释了为什么手工设计的2D position embedding性能没有提升。

  自监督使得ViT即使是在最底层也能整合整张图像的信息。我们对网络自监督能力的使用程度进行了研究。具体来说，我们根据注意力权重计算图像空间中信息整合的平均距离，如图7右所示。这个“注意力距离”和CNN中的接收域尺寸相似。我们发现有些注意力头在网络底层就已经整合了图像的绝大部分区域，表明ViT的确使用全局信息整合的能力。另外的注意力头则只注意到图像的一小部分。这种高度集中的注意力在混合模型（ViT前加一个ResNet）中更少见，表明局部注意力的作用和CNN中前几层卷积层的作用相近。而且，注意力距离随着网络的加深而增大。从全局上看，我们发现模型更关注与和分类语义相关的图像区域，如图6和14所示。

4.6 Self-Supervision

  Transformer在NLP任务中展现了优异的性能。然而他们大多数的成功并不仅仅是因为Transformer优异的扩展性，同样归功于大规模的自监督预训练。我们还模仿BERT中使用的masked language modeling任务，对自监督的masked patch prediction做了初步探索。在自监督预训练下，我们较小的ViT-B/16模型在ImageNet上达到了79.9%的准确率，相比从头训练提升了2%，但是相比有监督的预训练仍然落后4%。附录B.1.2包含更多的细节。我们将对比预训练的探索留给未来工作。

5. Conclusion

  我们探索了直接将Transformer应用于图像识别。和前人在计算机视觉中使用自监督的工作不同，除了初始图像块提取以外，我们不引入针对图像的归纳偏置到网络结构中。作为代替，我们将一张图像视为一个图像块序列，然后使用NLP的标准Transformer编码器处理。这种简单但可扩展的策略结合在大型数据集上进行预训练效果出奇的好。因此，ViT在许多图像范磊数据集上逼近甚至超越了SOTA，而且其预训练开销相比之下降低了很多。
  尽管这些初始的结果很令人激动，但仍面临着许多挑战。其中一个是将ViT应用到其他计算机视觉任务中，比如检测和分割。我们以及Carion的结果表明这种方法的可信度。另一个挑战是持续探索自监督的预训练方法。我们最初的实验表明自监督预训练带来的提升，但是自监督和大规模有监督之间仍有巨大的鸿沟。最后，ViT的进一步扩展可能会带来新的性能提升。

6. 其他

  最后附上ViT-B/16的网络结构图。我自己在ImageNet-1k上从头训练了一下，2080的单卡，batch_size=24，一个epoch跑6个小时，最后跑了3个epoch准确率才提到20%左右，收敛速度是在有点慢，但看loss曲线还是在下降的，难怪论文里实验里训了300个epoch。