An Image Is Worth 16X16 Words: Transformer For Image Recognition At Scale

论文背景

Transformer在NLP相关任务中已经取得了非常好的效果，也成为了NLP任务的首选模型。这个方法主要是通过在大型文本语料库上进行预训练，然后针对较小的任务特定数据集进行微调。Transformer的优点是良好的计算效率和可扩展性，这个模型现在可以训练出参数超过100B的模型，性能也没有饱和的迹象。但是在计算机视觉的领域，卷积架构仍然占主导地位，在大规模的图像识别任务上，经典的类ResNet架构仍然是最先进的。
在本文之前已经有研究将标准的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% 的准确率。

重要方法

如下图所示，本文的架构基本遵循原版Transformer架构，方便研究人员开箱可用。标准的Transformer接收一维序列的token作为输入，因此为了处理二维的图像，作者将图像$$x\in R^{H\ast W\ast C}$$压缩成为一系列图像块$$x_p\in R^{N\ast (P^2\ast C)}$$, 然后在所有层都使用一个可训练的线性映射把每一个patch向量映射到模型维度D上，得到如下图所示粉色区域，也就是patch embedding。与标准的Transformer一样，作者也在patch embedding上加position embedding，把他两的和一起输入到Encoder里面。这里作者比较了1维和2维的position embedding，发现对结果影响不大，作者的解释是分成patch之后，2维的position embedding并没有比1维position embedding带有的信息多多少。
与BERT类似，作者在一系列patch embedding之前添加了一个可学习的embedding，在encoder里面这个作为图像表征。预训练和微调阶段，classification head都依附于这个KaTeX parse error: Expected group after '_' at position 2: Z_̲.加这个原因是因为最后要输入到最后的分类头里面去，但如果输入后面的每一块都不太对，因此单独多加一块作为最后分类的依据，也就是整个图像的表征

实验结果

首先是与SOTA方法对比，可以看到在正确率上有所提升，而在训练时间上的提升更大。

下图显示了ViT在小型数据集上不如ResNet，但在随着数据量的增大，他的效果会越好越好。该结果强化了这一认知：卷积归纳偏置对于规模较小的数据集较为有用，但对于较大的数据集而言，学习相关模式就足够了，甚至更加有效。

在可迁移能力上，ViT的效果也比ResNet要好

感想与总结

这一篇工作可以说是挖开了一个新坑，即把Transformer以一种真正脱离CNN的方式用在了计算机视觉的问题中。而且这篇文章还留下了很多值得讨论的地方，比如它只是把Transformer用在了图像分类这个CV中最基本的问题中（用了一个分类头就解决了），那么在目标检测，分割甚至是图像生成等任务中，我们又该如何使用Transformer呢？
另外在appendix里面，作者还提出了他们基于Axial Attention（即每一轴上都上一个attention）进行的一个尝试，他们的结果如下图所示，可以说是正确率上又上了新高，但随之带来的计算量也是大增。Axial Attention其实也是非常naive的做法，怎么想到一个更好的在多个轴上都实现attention，且不会带来太大的计算量的方法也是我们下一步可以解决的点。