[Transformer]ViT:An Image Is Worth 16X16 Words: Transformers for Image Recognition at Scale

ViT:An Image is worth 16*16 words:Transformers for image recognition at scale

Abstract

Transformer已经成为自然语言处理任务中的一种标准模型，但将其用于计算机视觉任务还存在一些限制。在计算机视觉领域，注意力要么与卷积神经网络一起使用，或者用于替换卷积网络中的某些组件，本文认为attention无需依赖CNN，直接对图像序列应用attention就可以达到可观的效果。本文的ViT在大型数据集上训练后，迁移到多个中小型图像识别任务中均取得了较好的识别效果，同时所需的计算资源也进一步减少。

Section I Introduction

Transformer这种基于自我注意的架构已经广泛用于VLP领域，主要的方法就是现在一个大规模语料库上进行预训练，然后迁移到小型特定任务数据集上进行微调。鉴于Transformer的计算效率及可伸缩性，有的网络参数已经超过100B，随着数据集的不断扩充，Transformer的性能一直在提升，没有饱和的迹象。
然而在计算机视觉领域，CNN仍占主导地位。

受启发于Transformer，越来越多的研究者致力于将self-attention与CNN结合，有的甚至尝试完全将卷积层替换掉。后者虽然在理论上证明了有效性，但因为其使用的是特殊模式的注意力计算，无法在硬件上进行高效缩放；因此在大规模的视觉任务中最先进的仍然是经典的ResNet系列框架。


本文尝试直接用Transformer处理图像，不做过多的改动和调整；因此本文将输入图像切分成patch，将所有patch经过线性映射后的结果作为输入送入Transformer；输入的patch嵌入后的结果就作为token参与后续计算。



最初ViT在ImageNet等数据集上训练时没有进行正则化，此时训练得到的模型精度比ResNet略低几个百分点。这一结果可以预见：因为Transformer很难学习CNN的一些inductive bias，如平移不变性、locality局部性等，因此如果训练数据不充足泛化性能就不会很好。
但是如果在大型数据集上进行训练，则大规模数据集中包含的信息就可以胜过inductive bias，ViT就能达到优异的识别性能，最好的能在ImageNet,ImageNet-ReaL,CIFAR-100,VTAB上的精度分别为88.55%，90.72%，94,。55%，77.63%。

Section II Related Work

Transformer于2017年由Vaswani等人提出用于机器翻译任务，并已经成为诸多NLP任务中最先进的方法；通常都需要在大型语料库上进行预训练，然后在对特定任务进行微调。
原始的图像注意力的计算需要计算每个像素点与其他像素点之间的相关性；因为时间复杂度是O（n^2），所以不能扩展到实际的输入大小。
为了将Transformer用于图像处理类任务，许多学者做了诸多优化、近似的方法。如Parmar等在每一次query只查询临近的局部区域而不是计算全局的像素；这种局部的通过点积计算self-attention的方式可以完全取代卷积；还有尝试使用近似的方法，称之为稀疏Transformer，通过对global attention进行近似从而将Transformer用于图像；此外还有一种缩放注意力的方式是基于不同大小的block来计算注意力，极端情况下仅沿着一个轴计算。
上述peer work均限制注意力机制在计算机视觉任务中十分有前景，但还需要复杂的设计才能最终在硬件加速器上进行工程实现。

与本文采取相似思路的是Cordonnier等人的工作，他们将图像切分成2x2大小的patch然后计算注意力；本文则是基于这种思路进一步证明了，大规模的预训练可以使得Transformer性能与CNN相媲美。

如何将CNN与self-attention结合也吸引了诸多学者的研究兴趣，有通过增强特征图谱用于分类任务，或者将CNN的输出再经过self-attention的处理用于目标检测等。


近期发布的iGPT模型则是将Transformer用于一些基础的像素级别的视觉任务，如降低图像分辨率、颜色空间等。首先会以无监督的方式训练一个生成模型，然后将生成的结果通过微调或者线性探测的方式用于分类，在ImageNet上能达到72%的精度。
本文则是探索Transformer在比ImageNet更大规模数据集上进行图像识别的效果。

Section III Method

本文的架构力求完全复刻原始的Transformer架构，原始架构中简洁的结构便于放缩、实现起来也更加简单，几乎可以开箱即用。

Part 1 ViT-Vision Transformer

Fig 1展示了ViT的基本结构，原始的Transformer接受的是token嵌入后的1D输入，为了处理2D图像，本文将图像reshape成一系列展平的2Dpatch组成的序列。




Step 1:将原始的HWC的图像 切成N个PPCpatch；并将patch通过线性投影成一维向量P2C




Step 2:
为了仿照原始Transformer中的positional embedding，本文预先append了一个可训练的1D vector用于嵌入位置信息，参见Eq 4.
随后patch+positional embedding联合作为encoder的输入。




Step 3:Transformer Encoder由多层多头注意力模块（MSA）和全连接层（MLP）模块堆叠而成，模块之间经过LN层归一化，并且模块之间使用了残差连接。
Inductive bias

本文发现ViT中包含比CNN更多的图像相关的归纳偏执。CNN 中每一层会学习到局部特征、平移不变性等；而ViT中MLP模块可以有效学习局部特征和平移不变性，而MSA模块则可以有效学习全局关联。因为送入encoder的是以图像patch的形式，不懈怠任何二维位置信息，因此这些位置信息、空间关系都需要从头学起。


Hybrid Architecture



除了将patch作为输入，还可以将特征图作为输入。在混合模型中，输入的patch是混合了从CNN提取的特征图谱经过嵌入之后的结果。比如输入的patch维度是1*1大小，则是完全取空间维度展平的结果作为输入。

Part 2 Fine-tuning and Higher Resolution

通常会先在大型数据集上对Transformer进行预训练，然后在特定任务的小型数据集上进行微调。
因此，本文移除了预训练中做预测的部分（prediction head），将其替换为初始化为0的D*K的前馈层，K则是特定任务的分类种类。 通常在微调时使用更高分辨率效果更好；此时保持输入patch不变，则会导致整个输入序列变长；但ViT好处就是可以处理任意长度的输入序列。






此外，本文会在预训练后得到的位置嵌入信息中通过二维插值调整分辨率。切分patch和调整分辨率 这是在ViT中唯二引入人工inductive bias的地方。

Section IV Experiments

本文评估了ResNet，ViT和混合模型的表征学习能力。训练都是先在大型数据集上进行预训练然后在benchmark上进行对比。
实验结果是ViT以较低的训练成本在大多数benchmark上达到了SOTA。
本文还测试了自监督ViT的性能，展示了ViT十分具有前景。

Part 1 Set up

为了探究模型的可扩展性，本文基于的数据集包括：

ImagheNet 1000类 1.3M


ImageNet-21k 21k类 14M



JFT 18k类 303M高分辨率图像



在上述数据集进行预训练后再以下数据集上进行测试：




ImageNet





ReaL






CIFAR-10/100







Oxford-IIII Pets







Oxford Flowers-102

实验模型
本文采用的Transformer基础模型主要保持和BERT中一致，具体层、每层的详细规模参见Table 1。


细节：输入序列长度与patch大小的平方 成反比，因此patch越小 其实计算成本更大。

本文采用的CNN模型则是ResNet，但将BN替换为GN，使用的卷积是Qiao等人提出的标准化卷积。

本文采用的混合模型：将CNN提取出的中间层feature map作为ViT的输入，并且输入不同分辨率的feature map作为输入，从而使得ViT的输入序列长度是变化的。
训练及微调
Adam优化器
batch-size 4096
评价指标
微调后的精度
少样本的精度

Part 2 Comparison Results

Table 2展示了不同规模的ViT与BiT-L,Moist Student网络在不同数据集上的对比结果。
可以看到在JFT预训练后的Transformer在所有数据集上的效果均优于ResNet，同时预训练所需的计算成本是最低的；在小型数据集上ViT也表现良好。

此外本文还发现预训练的效率不仅会被具体网络结构影响，还与其他参数有关：如训练策略、优化器、衰减率等。

在ImageNet-21k数据集上预训练的ViT-L/16在绝大多数数据集上表现优异，但所需的计算资源更少：在8核TPUv3上预训练30天。

Part3 对预训练数据的要求

ViT在JFT数据集上预训练后效果很棒，但数据集规模的影响到底有多大，本文也进行了实验。
实验分别在不同规模数据集上进行训练，分别是ImageNet,ImageNet-21k,JFT-300M.



为了提升在较小规模数据集上的性能，本文使用了3类基本的正则化：权重衰减、随机失活、标签平滑。






Fig 3展示了最终在ImageNet上的训练结果。可以看到预训练数据集较小，BiT性能更佳；但随着预训练数据集的增大，ViT的性能超过了BiT.、
此外，本文还测试了使用不同规模的预训练子集，即JFT-9M,30M,90M及整个JFT-300M。但没有进行任何正则化，使用同样超参数的网络。对比结果参见Fig4.
可以看到Bit-B/32比ResNet略快但是在9M时性能较差，90M时性能更好。
这种实验结果再一次证明，CNN中的归纳偏执对小规模数据集更有用；当数据集规模足够大，就可以直接从数据中学习各种相关模式了，甚至学习效果更好。

Part 4 缩放实验

本文还测试了不同训练成本对模型精度的影响。ViT不同后缀代表patch的大小，对比结果参见Fig 5.可以看到ViT性能通常优于混合模型和Resnet。

Part 5 深入分析ViT

为了进一步理解ViT是如何处理图像数据的，本文分析了其内部的数据表示。
Fig 7展示了原始输入经过嵌入以及线性映射之后的结果，左上方显示了主要学习到的内容，主要代表了每一个patch中的细节特征。



Fig 7中间显示的是不同patch之间的注意力，可以看到相邻的patch之间其拥有更加相似的positional embedding结果；而同一行/列的patch也拥有相似的embedding。
self-attention使得ViT即使在较低层次都可以有效捕获整张图像的信息；通过计算这种信息被聚合的平均距离，发现这种“平均距离”类似于CNN中感受野的概念。



此外，本文还发现有些head在较浅的层次就已经出现，这表明模型确实使用了注意力来聚合全局信息；而其他一些head在较低层次的注意力距离一直很小，表现出的是局部注意力，表明可能与CNN中的早期卷积层具有类似的功能。


 从全局来看，ViT倾向于关注语义上属于同一类的区域，参见Fig6.
Part 6 Self-supervision
Transformer在NLP领域取得了优异的性能，但是其成功的原因除了其优秀灵活的特征提取能力，还得益于其在大规模数据集上自监督预训练的结果。


 
本文同样测试了使用带mask的patch进行自监督训练的结果，模拟了BERT中使用mask language的实验。通过自监督预训练，在ImageNet上Vit-B/16精度为79.9%，比从头训练精度获得了2%的提升。我们下一步则是探索对比预训练。

Section V Conclusion

本文直接将Transformer应用于图像识别任务，提出了ViT模型。与之前研究在计算机视觉中引入slef-attention机制的不同之处在于，本文没有引入任何图像相关的归纳偏执，而是将图片处理为一系列patch序列作为标准Transformer encoder的输入。
这种方式结合在大规模数据集上进行预训练后取得了优异的效果。因此ViT在预训练成本较低的前提下在诸多图像分类数据集上取得了相媲美甚至更优秀的识别结果。

尽管初步的应用取得了振奋人心的效果，但仍然存在许多挑战。一是如何将ViT应用于其他视觉任务，如检测和分割；另一方面是在自监督的预训练方法上做进一步探索。本文的实验初步证明了使用自监督预训练会带来一定的优化，但仍然与在大规模数据集上预训练的结果存在较大的差距。
最后，ViT也将会在如何优化自身性能上做进一步的探索。