Transformer论文阅读：ViT算法笔记

标题：An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
会议：ICLR2021
论文地址：https://openreview.net/forum?id=YicbFdNTTy

Abstract

虽然Transformer结构已经成为自然语言处理任务的事实标准，但其在计算机视觉中的应用仍然有限。在视觉中，注意力要么与卷积网络结合使用，要么用于替换卷积网络的某些组件同时保持它们的整体结构。我们表明，这种对CNN的依赖不是必要的，直接用于图像块序列的纯Transformer可以在图像分类任务中表现得很好。当在大量数据上进行预训练，再迁移到多个中小型图像识别基准（ImageNet、CIFAR- 00、VTAB等）时，视觉变换器（Vision Transformer，ViT）与SOTA的卷积网络相比取得了优异的结果，并且需要的训练计算资源更少。

1 Introduction

基于自注意力机制的结构，尤其是Transformers，已经成为自然语言处理（NLP）的首选模型。主流的方法是在大型文本语料库上进行预训练，然后在较小的特定任务数据集上进行微调。得益于Transformers的计算效率和可扩展性，训练参数超过100B的规模空前的模型已经成为可能。随着模型和数据集的增长，目前仍然没有性能饱和的迹象。
然而，在计算机视觉中，卷积结构仍然占据主导地位。受NLP成功的启发，许多工作尝试将CNN结构与自注意力相结合，其中有些完全取代了卷积。后者模型虽然在理论上是有效的，但由于使用了专门的注意力模式，在现代硬件加速器上还没能有效地推广。因此，在大规模图像识别中，经典的ResNet式结构仍然是SOTA的。
受NLP中Transformer大规模成功的启发，我们以尽可能少的改动，直接将标准的Transformer用在图片上进行实验。为此，我们将图像拆分为小块（patch），并提供这些patch的线性嵌入序列作为Transformer的输入。图片块像NLP程序里一样被当作tokens（words）。我们以监督的方式训练图像分类模型。
当在中等大小的数据集（如ImageNet）上进行训练时，如果没有强大的正则化，这些模型的精度会比同等大小的ResNets低几个百分点。这种看似令人沮丧的结果可能是意料之中的：Transformers缺少CNN固有的一些归纳偏差，如平移等变性和局部性，因此在训练数据量不足时不能很好地泛化。
然而，如果在更大的数据集（14M~300M张图像）上训练模型，情况会发生变化。我们发现大规模训练要比归纳偏差更重要。我们的视觉变换器（Vision Transformer，ViT）在足够大规模的预训练后，再迁移到数据点较少的任务时取得了优异的效果。当在公开的ImageNet-21k数据集或内部的JFT-300M数据集上预训练后，ViT在多个图像识别基准上接近或超过了当前的SOTA。

2 Related Work

Transformers由Vaswani等人提出用于机器翻译，此后成为许多NLP任务中SOTA的方法。基于大型Transformer的模型往往在大型语料库上进行预训练，然后针对手头的任务进行微调：BERT使用一个去噪自监督预训练任务，而GPT系列使用语言建模作为其预训练任务。
自注意力在图像上的简单应用需要每个像素关注所有其它像素。由于像素数平方级的代价，这并不能扩展到实际的输入尺寸。因此，为了将Transformers用于图像处理的环境，过去已经尝试了几种近似方法。Parmar等人对每个查询像素只在局部而不是全局邻域使用自注意力。这样的局部多头点积自注意力块可以完全替代卷积。在不同的工作方向中，Sparse Transformers采用可扩展的全局自注意力近似以便适用于图像。改变注意力的另一种方法是将其用于不同大小的块，在极端情况下只沿单个坐标轴。很多这些特殊的注意力结构在计算机视觉任务上展示出有前景的结果，但需要复杂的工程才能在硬件加速器上有效地实现。
与我们最相关的是Cordonnier等人的模型，该模型从输入图像中提取大小为2×2的patch，并在顶部使用完全自注意力。该模型与ViT非常相似，但我们的工作进一步证明了大规模预训练使得原始的Transformers与SOTA的CNN具有竞争力（甚至更好）。此外，Cordonnier等人使用2×2像素的小块尺寸，这使得模型只适用于小分辨率的图像，而我们也处理中等分辨率的图像。
卷积神经网络（CNNs）与自注意力的结合也引起了广泛的兴趣，例如通过增强特征图进行图像分类，或者通过自注意力进一步处理CNN的输出，如用于目标检测、视频处理、图像分类、无监督目标发现或统一的文本视觉任务。
最近另一个相关的模型是图像GPT（iGPT），它在减少图像分辨率和颜色空间后，将Transformers用于图像像素。该模型以无监督的方式作为生成模型进行训练，然后可以对生成的表示进行微调或线性搜索，以提高分类性能，在ImageNet上达到72%的最高精度。
我们的工作增加了在比标准ImageNet数据集更大规模上探索图像识别的论文集。使用额外的数据源可以在标准基准上取得SOTA的结果。此外，Sun等人研究了CNN性能如何随数据集大小变化；Kolesnikov等人和Djolonga等人在ImageNet-21k和JFT-300M等大规模数据集上进行CNN迁移学习的实证研究。我们也关注后两个数据集，但训练的是Transformers而不是之前工作中使用的基于ResNet的模型。

3 Method

在模型设计中，我们尽可能地遵循原始的Transformer。这种故意简单设置的一个优点是，可扩展的NLP Transformer结构——以及它们的高效实现——几乎可以开箱即用。

3.1 Vision Transformer

模型的概述如图1所示。标准的Transformer接收一个1维的token嵌入序列作为输入。为了处理2维图像，我们将图像$\mathrm{x}\in {\mathbb{R}}^{H\times W\times C}$重构为一系列展平的2维图像块${\mathrm{x}}_p\in {\mathbb{R}}^{N\times (P^2\cdot C)}$，其中，$(H,W)$是原始图像的分辨率，$(P,P)$是每个图像块的分辨率，$N=HW/P^2$是产生的图像块数，它也是Transformer的有效输入序列长度。Transformer在其所有层中使用恒定的隐藏向量大小$D$，因此我们将图像块展平，并使用可训练的线性投影来将其映射到$D$维（公式1）。我们将此投影的输出称为块嵌入（patch embeddings）。
与BERT的[class] token类似，我们在嵌入块（${\mathrm{z}}_0^0={\mathrm{x}}_{\mathrm{class}}$）的序列上预置一个可学习的嵌入，其在Transformer编码器输出处的状态（${\mathrm{z}}_L^0$）作为图像表示$\mathrm{y}$（公式4）。在预训练和微调期间，一个分类头被附加到${\mathrm{z}}_L^0$。分类头由一个MLP实现，在预训练时有一个隐藏层，在微调时有一个线性层。
位置嵌入（position embeddings）被添加到块嵌入（patch embeddings）以保留位置信息。我们使用标准的可学习的1维位置嵌入，因为我们没有观察到使用更先进的2维感知位置嵌入带来的显著性能增益（附录D.3）。得到的嵌入向量序列作为编码器的输入。
Transformer编码器由交替的多头自注意力（multiheaded self-attention，MSA，见附录A）和MLP块(公式2和3)层组成。Layernorm（LN）用于每个块之前，残差连接用于每个块之后。
MLP包含两个具有GELU非线性的层。
$$\begin{array}{rlrlr}{\mathrm{z}}_0& =[{\mathrm{x}}_{\mathrm{class}};{\mathrm{x}}_p^1\mathbf{E};{\mathrm{x}}_p^2\mathbf{E};\cdot \cdot \cdot ;{\mathrm{x}}_p^N\mathbf{E}]+{\mathbf{E}}_{pos},& \mathbf{E}& \in {\mathbb{R}}^{(P^2\cdot C)\times D},{\mathbf{E}}_{pos}\in {\mathbb{R}}^{(N+1)\times D}& \\ {{\mathrm{z}}^{\prime }}_{\ell }& =\mathrm{MSA}(\mathrm{LN}({\mathrm{z}}_{\ell -1}))+{\mathrm{z}}_{\ell -1},& \ell & =\mathrm{1...}L& \\ {\mathrm{z}}_{\ell }& =\mathrm{MLP}(\mathrm{LN}({{\mathrm{z}}^{\prime }}_{\ell }))+{{\mathrm{z}}^{\prime }}_{\ell },& \ell & =\mathrm{1...}L& \\ \mathrm{y}& =\mathrm{LN}({\mathrm{z}}_L^0)& & & \end{array}$$归纳偏置。 我们注意到Vision Transformer比CNN具有的特定于图像的归纳偏置更少。在CNN中，局部性、2维邻域结构和平移不变性贯穿于整个模型的每一层。在ViT中，只有MLP层是局部和平移不变的，而自注意力层是全局的。2维邻域结构使用得非常少：仅在模型最开始裁剪图像块时和微调阶段针对不同分辨率图像调整位置嵌入（position embedding）有用到。除此之外，位置嵌入在初始化时没有任何关于块的2维位置信息，块之间的所有空间关系都必须从头开始学习。
混合结构。 作为原始图像块的替代，输入序列可以从CNN的特征图中形成。在这个混合模型中，块嵌入（patch embedding）投影$\mathrm{E}$（公式1）被用于从CNN特征图中提取的块。在极端情况下，块的空间尺寸可以是1×1，这意味着输入序列是通过简单地将特征图的空间维度展平并投影到Transformer维度得到的。分类输入嵌入和位置嵌入按照上述方式加入。

3.2 Fine-Tuning and Higher Resolution

通常，我们在大型数据集上预训练ViT，然后微调到（更小的）下游任务。为此，我们去掉预训练的预测头，增加一个零初始化的$D\times K$的前馈层，其中$K$为下游任务预测的类别数量。通常使用比预训练更高分辨率的图像微调是有益的。当输入更高分辨率的图像时，我们保持块的尺寸不变，这将导致更大的有效序列长度。Vision Transformer可以处理任意的序列长度（直至内存限制），但是预训练的位置嵌入可能不再有意义。因此，我们对预训练的位置嵌入进行2维插值，根据它们在原始图像中的位置。值得注意的是，这种分辨率调整和块提取是关于图像2维结构的归纳偏置被手动添加到Vision Transformer的唯一一处。

4 Experiments

我们评估了ResNet、ViT和混合模型的表示学习能力。为了了解每个模型的数据需求，我们在不同大小的数据集上进行预训练，并在许多基准任务上进行了评估。当考虑预训练模型的计算成本时，ViT表现非常出色，以更低的预训练成本在大多数识别基准上达到了SOTA。最后，我们进行了一个使用自监督的小实验，并表明自监督的ViT在未来是有前景的。

4.1 Setup

数据集。 介绍了实验所用的数据集。
模型变体。 我们在BERT所使用的模型结构基础上确定ViT的配置，如表1所示。“Base”和“Large”模型是直接取自BERT，“Huge”是我们增加的更大的模型。下面我们使用简短的符号表示模型大小和输入块尺寸：例如，ViT-L/16表示输入块尺寸为16×16的“Large”模型变体。注意，Transformer的序列长度与块尺寸的平方成反比，因此块尺寸更小的模型在计算上更加高昂。

对于CNN基线，我们使用ResNet，但将Batch Normalization层替换为Group Normalization层，并使用标准化卷积。这些改动可以提升迁移的性能，我们用“ResNet(BiT)”表示修改后的模型。对于混合模型，我们将中间层的特征图送给ViT，块尺寸为1个像素。为了实验不同长度的序列，我们要么①使用常规ResNet50中stage4的输出，要么②移除stage4，在stage3中放置相同的层数（保持总层数不变），然后取这个扩展的stage3的输出。选项②会导致4倍长的序列长度，对应的ViT模型计算开销更大。
训练&微调。 介绍了训练和微调的实现细节。
评价指标。 我们报告了下游数据集少样本和微调的准确率结果。微调准确率体现的是每个模型在对应数据集上微调后的性能。小样本准确率是通过求解一个将训练图像子集的（冻结）表示映射到{−1,1}K\{-1,1\}^K{−1,1}K个目标向量的正则化最小二乘回归问题得到。这个公式使得我们能够以闭环的方式获取精确解。尽管我们主要关注微调性能，但当有时候微调成本太高时，我们会使用线性少样本准确率来进行快速的在线评估。

4.2 Comparison to State of the Art

和SOTA的对比实验，具体的实验结论可以参照原文。

4.3 Pre-Training Data Requirements

在大型的JFT-300M数据集上预训练时，Vision Transformer表现较好。相比ResNets，ViT使用的视觉归纳偏置更少，那么数据集大小是关键吗？对此我们进行了两组实验。
首先，我们在越来越大的数据集上预训练ViT模型：ImageNet、ImageNet-21k和JFT300M。为了提高在较小数据集上的性能，我们优化了三个基本的正则化参数——权重衰减、Dropout和标签平滑。图3显示了在ImageNet上微调后的结果（其它数据集上的结果如表5所示）（注意，ImageNet预训练模型也进行了微调，但仍在ImageNet上进行。这是因为微调过程中分辨率的提高改善了性能）。当在最小的ImageNet数据集上进行预训练时，尽管进行了（普通的）正则化，但ViT-Large模型性能仍不如ViT-Base。当在ImageNet-21k上预训练时，两个模型的性能接近。只有当使用JFT-300M数据集预训练时，我们才能看到更大模型的全部优势。图3还展现了不同大小的BiT模型的性能：BiT CNNs在ImageNet上的表现优于ViT，但在更大的数据集上，ViT超过了BiT。


然后，我们在9M、30M和90M的随机子集以及完整的JFT300M数据集上训练我们的模型。我们不对较小的子集执行额外的正则化，并在所有设置中使用相同的超参数。这样，我们评估的是内在的模型性质，而不是正则化的效果。然而，我们使用了早停，并报告了在训练过程中达到的最佳验证精度。为了节省计算量，我们报告了少样本线性精度，而不是完全微调精度。结果如图4所示。结果强化了一个直觉，即卷积归纳偏差对于较小的数据集是有用的，但对于较大的数据集，直接从数据中学习相关模式是足够的，甚至是有益的。具体的实验结论可以参照原文。
总体而言，ImageNet上的少样本结果（图4），以及VTAB上的低数据结果（表2）看起来有望实现非常低数据的迁移。进一步分析ViT的少样本特性是未来工作的一个令人激动的方向。

4.4 Scaling Study

我们通过评估JFT-300M的迁移性能，对不同模型进行了控制扩展研究。在这种设定下，数据集大小不是模型的性能瓶颈，我们评估每个模型性能和预训练开销的关系。具体的实验设置可以参照原文。


图5包含了迁移性能与总预训练计算，每个模型的详细结果见表6。可以观察到一些模式。首先，Vision Transformers在性能/计算权衡上与ResNets相比占绝对优势。ViT使用大约2-4倍更少的计算量来达到相同的性能（平均超过5个数据集）。其次，在较小的计算预算下，混合模型略优于ViT，但对于较大的模型，这种差异消失了。这个结果有点令人惊讶，因为我们可能期望通过局部卷积特征处理来辅助任意尺寸的ViT。第三，Vision Transformer在尝试范围内未出现性能饱和的现象，刺激未来的扩展努力。

4.5 Inspecting Vision Transformer

为了开始理解Vision Transformer如何处理图像数据，我们分析了它的内部表示。Vision Transformer的第一层将展平的块线性投影到一个较低纬度的空间（公式1）。图7左显示了学习得到的嵌入滤波器的顶端主要组成。这些组成类似于可靠的基函数，用于低维表示每个块内的精细结构。



在投影之后，一个学习到的位置嵌入被添加到块表示中。图7中表明模型将图像中的距离使用位置嵌入相似性进行编码，即距离更近的块倾向于具有更相似的位置嵌入。而且，出现了row-column结构，同一行/列的块有相似的位置嵌入。最后，在较大的网格中有时会出现明显的正弦结构。位置嵌入可以学习到表示2维图像的拓扑结构，这解释了为什么手工设计的2维感知嵌入变体没有带来改进。
自监督使得ViT即使是在最底层也能整合整张图像的信息。我们对网络自监督能力的使用程度进行了研究。具体来说，我们根据注意力权重计算图像空间中信息整合的平均距离（图7右）。这个“注意力距离”和CNN中的感受野大小相似。我们发现，有些头在网络最底层就已经注意到了图像的绝大部分区域，表明模型的确使用了全局信息整合的能力。其它的注意力头在低层具有一致的小注意力距离。这种高度局部化的注意力在Transformer之前使用ResNet的混合模型中不太明显（图7右），这表明它可能与CNN中的前几个卷积层具有类似的功能。而且，注意力距离随着网络深度的增加而增加。从全局来看，我们发现模型关注的是与分类语义相关的图像区域（图6）。

4.6 Self-Supervision

Transformer在NLP任务上表现出令人印象深刻的性能。然而，它们的成功很大程度上不仅源于其出色的可扩展性，还源于大规模的自监督预训练。我们还模仿BERT中使用的masked language modeling任务，对自监督的masked patch prediction做了初步探索。在自监督预训练的情况下，我们较小的ViT-B/16模型在ImageNet上达到了79.9%的准确率，相比从头训练提升了2%，但仍然落后于有监督预训练4%。

5 Conclusion

我们探索了Transformers在图像识别中的直接应用。与以前在计算机视觉中使用自注意力的工作不同，除了最初的块提取步骤外，我们没有将特定于图像的归纳偏置引入到结构中。取而代之的是，我们将一幅图像理解为一系列的图像块，并通过NLP中使用的标准Transformer编码器进行处理。这种简单但可扩展的策略在与大型数据集上的预训练相结合时效果出奇地好。因此，在许多图像分类数据集上，Vision Transformer匹配或超过了SOTA，同时相对便宜的预训练。
虽然这些初步成果令人鼓舞，但仍存在许多挑战。一种是将ViT应用于其它计算机视觉任务，如检测和分割。我们的结果，加上Carion等人的结果，表明了这种方法的前景。另一个挑战是继续探索自监督预训练方法。我们的初步实验表明自监督预训练有所改进，但自监督预训练与大规模监督预训练之间仍有较大差距。最后，ViT的进一步扩展可能会导致性能的提高。