【论文精读】AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

前言

Transformer架构应用在CV的开山之作，证明了预训练+微调范式的Transformer同样适用于CV领域，这种跨界的联动，才是新工作的源源不断的动力。本文参考了沐神的B站视频讲解，清晰易懂，感兴趣的朋友可以点击链接观看。

---

Paper: https://arxiv.org/pdf/2010.11929.pdf
Code: https://github.com/google-research/vision_transformer

Abstract

虽然Transformer已经成为了NLP领域的标准，但是Transformer做CV任务还是受到很大的限制，只能和卷积网络结合使用，无法修改卷积框架。本文证明了Transformer在大规模数据集上进行预训练，也能在图像分类任务上达到SOTA，只不过还是需要很多的计算资源。

1 INTRODUCTION

Transformer已经是NLP领域模型的选择，它的核心思想在大规模数据集上进行预训练，在小规模任务特定数据集上微调。由于Transformer框架计算高效和可扩展性，模型可以扩展到超过100B规模。随着模型和数据集的增长，模型的性能仍未见饱和。
但是在计算机视觉领域，卷积架构仍然处于主导地位，虽然受到NLP启发，部分工作尝试将CNN和注意力机制结合或者全部替换成自注意力，后者虽然理论可行，但是无法在GPU上有效扩展。
本文受到Transformer在NLP领域扩展成功的启发，将Transformer直接应用到图像上，几乎不做额外的改动。为了实现这个过程，作者将图像分解为图像块，然后将这些图像块的embedding顺序输入到Transformer中。作者以监督学习的方式训练图像分类模型。
但是在训练时作者发现，ViT模型缺乏了CNN固有的归纳偏置（局部性和平移不变性），导致在训练数量不足的情况下达不到同等大小ResNet的性能。因此，作者尝试了在更大的数据集上预训练（14M-300M）发现，大规模的训练会优于CNN的归纳偏置，如在ImageNet-21k数据集或JFT-300M数据集上进行预训练，ViT可以达到很多图像识别benchmark的SOTA。

这里要探讨一下Transformer应用在视觉领域的难点：

1. 如何把2d的图片转为1d的元素集合。
2. 虽然可以简单把2d图像拉直，但是输入token个数太多计算复杂度高不可攀（224×224）。

此外，文中所提及训练数量不足的情况，即在ImageNet这样的数据集上都远远不够（作者称之为中型数据集），可见预训练还是相当吃资源的。

2 RELATED WORK

Transformer自提出以来已经在很多NLP任务中处于SOTA地位，典型的代表如使用去噪自监督预训练的BERT，使用语言模型作为预训练任务的GPT。
原始的自注意机制在图像上的应用需要将所有像素作为token，一般的图像都有224×224个像素，这将是巨大的开销。为了解决这个问题，一些工作只对某个像素的局部区域应用自注意力，这种局部多头点积自注意力块可以完全取代卷积。另一个方向的工作应用稀疏Transformer，采用可扩展的全局注意力的近似，以作用于整个图像。此外还有轴注意力的工作，将2d图像拆分成两个1d的向量，从高度和宽度两个方向应用自注意力。这些工作都展示了有希望的结果，但是无法在GPU上得到加速。
和本文最相近的工作是将图分解为大小为2×2的图像块（本文是16×16），虽然思想一致，但是本文更突出的贡献在于大规模预训练。
此外还有很多卷积和自注意结合的工作，例如将CNN的输出输入到自注意力等。最近一个和本文相关的工作是iGPT，它在降低图像分辨率和颜色空间后输入到Transformer中，该模型以无监督方式进行训练，然后对得到的表示在特定任务上进行微调。但在ImageNet上只有72%的准确率。
本文的工作还涉及数据集对训练的影响，前人有研究CNN性能如何随着数据集大小的变化，以及对ImageNet-21k 和 JFT-300M 等大规模数据集的 CNN 迁移学习进行了实证探索。本文也重点关注这两个数据集在Transformer上的表现。

3 METHOD

模型设计中，作者尽可能遵循原始的Transformer框架，这样的好处是可以利用NLP中Transformer的可扩展性和高效实现。

作者尽可能最最小的改动，目的就是为了证明Transformer可以直接用在CV领域，无需CV领域的知识，这为Transformer向多模态进军提供了无限可能。

3.1 VISION TRANSFORMER (VIT)

模型的概述如上图所示。为了解决输入的问题，作者将图像$\mathrm{x}\in {\mathbb{R}}^{H\times W\times C}$分解成图像块${\mathrm{x}}_{\mathrm{p}}\in {\mathbb{R}}^{N\times (P^2\cdot C)}$序列。其中$(H,W)$是原始图像分辨率，$C$是通道大小。$(P,P)$是图像块的分辨率，$N=HW/P^2$是图像块的个数，也作为输入Transformer有效序列的长度。Transformer在其所有层中使用恒定大小的隐向量长度$D$，因此图像块铺平后还需要通过一个线性层映射成$D$维度的embedding。
类似于BERT的特殊token [class]，作者也设计了一个可学习的embedding添加到序列embedding中。此外为了加入位置信息，作者使用标准的可学习的1D位置embedding。生成的embedding向量序列作为输入加入到encoder中。

在附录部分，作者同样选择了2D位置信息进行了实验，发现结果并没有明显的差异，分析其原因可能是图像块和像素数量存在着明显的差异，图像块只有196个，而像素块个数有224×224个，因此图像块的位置远远没有像素块那样敏感，你可以理解为拼图游戏，面对只有一百多块拼图和超过五万块拼图，显然前者要容易的多得多。

Transformer的编码器部分由多头注意力层和MLP层交替组成，每层之前还有层归一化和残差，整个模型的计算过程如下：
$$\begin{array}{rlrl}{\mathbf{z}}_0& =\left[{\mathbf{x}}_{\text{class }};{\mathbf{x}}_p^1\mathbf{E};{\mathbf{x}}_p^2\mathbf{E};\cdots ;{\mathbf{x}}_p^N\mathbf{E}\right]+{\mathbf{E}}_{\text{pos }},& & \mathbf{E}\in {\mathbb{R}}^{\left(P^2\cdot C\right)\times D},{\mathbf{E}}_{\text{pos }}\in {\mathbb{R}}^{(N+1)\times D}\\ {\mathbf{z}}_{\ell }^{\prime }& =\mathrm{MSA}\left(\mathrm{LN}\left({\mathbf{z}}_{\ell -1}\right)\right)+{\mathbf{z}}_{\ell -1},& & \ell =1\dots L\\ {\mathbf{z}}_{\ell }& =\mathrm{MLP}\left(\mathrm{LN}\left({\mathbf{z}}_{\ell }^{\prime }\right)\right)+{\mathbf{z}}_{\ell }^{\prime },& & \ell =1\dots L\\ \mathbf{y}& =\mathrm{LN}\left({\mathbf{z}}_L^0\right)& & \end{array}$$

Inductive bias

作者注意到ViT的图像特定归纳偏置比CNN少得多，只有MLP层是具有局部性和平移不变性的，而自注意层是全局的。图片的2D信息只有在初始时切割为图像块时才使用到。因此所有图像块之间的空间信息必须从头开始学，这也从一定程度上解释了为什么需要在大规模数据集上预训练才能有很好的微调效果。

Hybrid Architecture

作为原始图像块的平替，输入序列可以通过CNN得到特征图，铺平后和打成图像块的序列大小维度一致，因此直接输入到Transformer中进行分类即可，这是一种可行的混合模型架构。

3.2 FINE-TUNING AND HIGHER RESOLUTION

通常来说，ViT在大规模数据集上预训练，在较小的数据量的下游任务上微调。有研究表明在更高分辨率下进行微调是有益的。但是更高的分辨率意味着更长的有效序列，预训练中的位置编码可能不再有意义，因此作者进行了简单的2D插值，这只是一个临时的解决方案，面对从序列长度差距过大的场景（如256×256的图像换成512×512的图像），会对结果产生影响。

4 EXPERIMENTS

作者在不同规模数据集上评估了ResNet、ViT和混合模型的表征学习能力。当考虑计算成本时，ViT的表现十分出色，以更低的预训练成本在大多数benchmark上达到了先进水平。最后作者还尝试了自监督预训练，并认为该模式具有一定的前景。

4.1 SETUP

Datasets

略，详见原文。

Model Variants

仿照BERT的配置，不同规模ViT如上表所示。ViT-L/16表示Large规模，输入的图像块大小为16×16。需要注意的是，模型的序列长度和图像块大小的平方成反比，图像块越小，计算上就需要更大的开销。
baseline的CNN模型采用ResNet，但是将批归一化替换为组归一化，使用标准化卷积，这样的改进可以提升迁移能力。对于混合模型，作者将中间特征图输入到ViT中。

Training & Fine-tuning

略。

Metrics

作者通过小样本或者微调精度来报告模型在下游任务上的结果。虽然主要关注微调的性能，但是小样本精度可以快速动态评估模型的性能，微调的成本太高。

4.2 COMPARISON TO STATE OF THE ART

作者首先将ViT-H/14和ViT-L/16与当前领域SOTA模型BiT和Noisy Student进行比较，前者采用监督迁移学习，后者采用半监督学习。所有模型都在TPUv3上进行训练。实验结果如下表所示：

较小的在JFT-300M数据集上预训练的ViT-L/16模型性能已经优于BiT-L，同时所需的计算资源要少得多。更大的模型ViT-H/14进一步提升了性能。和现有的技术相比，ViT预训练所需要的计算量少得多。

上图是VTAB任务分解为各自的组后进行的实验，也表现出相同的结果。

4.3 PRE-TRAINING DATA REQUIREMENTS

本小节分析数据集的规模对模型性能的影响。

作者首先在越来越大的数据集ImageNet、ImageNet-21k和JFT-300M上进行预训练。为了让模型在较小的数据集上表现良好，作者应用了三个基本的正则化参数：权重衰减，dropout和标签平滑。左图展示了ImageNet上进行微调的结果（其他数据集上的结果见下表）。当在更小数据集上微调时，ViT-Large性能低于ViT-Base，在中等规模数据上，二者性能相似，在更大数据集上，才真正看到大型模型的优势。
左图灰色部分也展示了ResNet模型的性能范围，可以看到随着训练数据的增大，ViT模型从开始的差距到相近再到更好的性能。
不对模型做正则化的结果如右图所示，可以得出这样的结论：卷积归纳偏置对于较小的数据集有用，但相较于更大的数据集，直接从数据中学习相关的模式甚至更有益。

4.4 SCALING STUDY

上图展示了模型传输性能和预训练计算量的比较。可以看到，ViT在性能/计算量上主导了ResNet，ViT只用了ResNet计算量的1/2或者1/4就达到了相同性能。其次，混合模型的性能在计算预算较小的情况下优于ViT，但是在更大的计算量上这种优势就消失了，从理论上来说，卷积所带来的归纳偏置知识结合Transformer应该会一直提升模型性能。最后，ViT和ResNet在尝试的范围内都未到饱和，这激励了未来的工作。

4.5 INSPECTING VISION TRANSFORMER

为了了解ViT如何处理图像，作者分析其内部表示。左图展示了ViT第一层学习到的图像组成部分，可以看到它学习到了部分全局的纹理、颜色等特征。中间图表明模型学习到了每个图像块所处的位置信息。右图展示了多头注意力随着层数的变化整合全局信息的能力。具体来说，作者根据注意力权重计算图像空间中信息整合的注意力距离，注意力距离可以类比CNN中的感受野大小。可以看到，注意力距离随着网络深度的增加而增加，并且更加关注与分类语义相关的图像区域（下图所示）。

4.6 SELF-SUPERVISION

Transformer在NLP上的成功很大原因在于大规模的自监督预训练。作者同样尝试模仿BERT中的mask语言模型任务，对mask的图像块进行预测。通过自监督的预训练，作者发现在ImageNet上实现了79.9%的准确率，这比重头开始训练提升了2%，但是仍然落后监督预训练4%。

5 CONCLUSION

本文探索了Transformer在图像识别中的应用。与之前CV中的注意力工作不同，作者并没有将特定于图像的归纳偏置引入到架构中，而是将图像分解为图像块，以类似于文本token的形式将这些图像块输入到Transformer中。当与大规模数据集预训练相结合，这种简单可扩展的策略效果出奇好，并在许多图像分类数据集上达到或超过了最先进的水平，同时预训练的成本相对较低。

阅读总结

ViT作为将Transformer框架迁移到CV领域的成功范例和开山之作，其成功主要在于两点：

1. Transformer和自注意力机制的优越性。
2. 预训练和微调范式的优越性。

这篇工作并没有特别的创新点，唯一可以称作创新的工作在于作者将图片分解为多个图像块，铺平映射为embedding，以token序列的形式输入到Transformer模型中。其他工作完全是对Transformer编码器层的复用。ViT的成功打通了CV和NLP之间的壁垒，统一了CV和NLP的架构，为多模态的工作带来了振奋人心的消息。之后无论是视觉、文本、语音甚至是触觉，都可以考虑使用统一的Transformer架构来建模，大一统模型时代确实越来越近了。
这篇文章对日后的工作来说也有很多的帮助。首先，跨领域、跨模态的方法往往可以为当前领域带来更多的启发，方法不能局限于同一领域，简单的A+B工作是没有意义的。其次，如果方法简单，那么就要在实验上下功夫，探究模型究竟是如何在这样的方法下工作的。比如本文就对数据的规模，attention机制，位置编码学习，模型融合等多角度进行了实验，通过丰富的实验又得到了很多新的发现和观点。