原文:Dosovitskiy A, Beyer L, Kolesnikov A, et al. An image is worth 16x16 words: Transformers for image recognition at scale[J]. arXiv preprint arXiv:2010.11929, 2020.
源码1:https://github.com/google-research/vision_transformer
源码2:https://github.com/rwightman/pytorch-image-models
虽然Transformer架构已经成为NLP任务事实上的标准,但其在计算机视觉中的应用仍然有限。在计算机视觉中,注意力机制要么与卷积网络结合使用,要么用来替换卷积网络的某些组件,同时保持其整体架构不变。我们的研究表明,这种对CNN的依赖是不必要的,直接应用于图像patches序列的纯Transformer可以很好地执行图像分类任务。当在大量数据上进行预训练并迁移到多个中小型图像识别基准任务时,我们提出的Vision Transformer(ViT)相比最先进的卷积网络取得了优异结果,同时需要的计算资源也更少。
1. 论文故事
基于自注意力机制的架构,例如Transformer,已成为自然语言处理(NLP)中的首选模型。这种方法先在大型文本语料库上进行预训练,然后在较小的任务特定的数据集上进行微调。由于Transformer具有良好的的计算效率和可扩展性,可以训练超过100B参数的超大模型,并且随着模型和数据集的增大,仍然没有表现出性能饱和的迹象。
然而,在计算机视觉中,卷积架构仍然占据主导地位。受Transformer在NLP领域成功的启发,许多工作尝试将CNN式架构与自注意力机制相结合,有些工作甚至完全取代了卷积。后一种模型虽然在理论上有效,但由于使用了专门的注意力模式,尚未在现代硬件加速器上有效地扩展。因此,在大规模图像识别任务中,经典的ResNet式架构仍然是最先进的。
受Transformer在NLP领域成功的启发,我们尝试将标准Transformer直接应用于图像,并尽可能地减少修改。为此,我们将图像切分为多个patches,并将这些patches的线性嵌入序列输入到Transformer中。图像patches的处理方式与NLP中的tokens(单词)相同。我们以有监督的方式训练模型进行图像分类。
当在中等规模的数据集(如ImageNet)上进行训练时,如果没有很强的正则化,这些模型的准确率会比同等规模的ResNet低几个百分点。这个看似令人沮丧的结果是可以预料的:Transformer缺乏CNN固有的一些归纳偏置,因此在数据量不足的情况下不能很好地进行泛化。
然而,如果模型在更大的数据集上进行训练,情况就会发生变化。我们发现,大规模训练胜过归纳偏置。当使用足够多的数据预训练ViT,然后将其迁移到数据量较少的任务上时,ViT取得了优异的结果。在ImageNet-21k数据集或JFT-300M数据集上进行预训练时,ViT在多个图像识别基准任务上接近或超过了SOAT水平。特别是,ViT在ImageNet上的准确率达到了88.55%,在ImageNet-ReaL上达到90.72%,在CIFAR-100上达到94.55%,在VTAB(共19个任务)上达到77.63%。
2. 模型方法
我们在设计模型时尽可能地遵循原始的Transformer。这种做法的一个优点是,几乎可以开箱即用地使用NLP Transformer架构。
如图1所示,我们将图像切分成固定大小的patches,对他们进行线性嵌入,并且添加位置嵌入,然后将生成的向量序列输入标准Transformer编码器中。图像patches的处理方式与NLP中的tokens(单词)相同。为了进行图像分类,我们在序列前面额外添加一个可学习的“分类token”。
图1:ViT模型的架构图。
表1:ViT模型的变体。
3. 实验结果
表2:不同模型的图像分类结果的比较。在JFT-300M数据集上预训练的ViT模型在所有数据集上都优于基于ResNet的基线,同时预训练所需的计算资源也大大减少。在较小的ImageNet-21k数据集上预训练的ViT也表现良好。
图2:不同模型在VTAB、Natural、Specialized、Structured任务上的性能比较。
图3:使用ImageNet、ImageNet-21k、JFT-300M数据集预训练模型,然后将其迁移到ImageNet上的实验结果。当在小数据集上进行预训练时,ViT模型的性能逊于BiT ResNets(阴影区域),但在较大数据集上进行预训练时,ViT模型表现出色。
图4:小样本线性评估结果。ResNet在较小的预训练数据集上表现更好,但比ViT更早稳定,ViT在较大的预训练数据集上表现更好。
图5:ViT、ResNet、Hybrid模型的性能随预训练计算量的变化情况。在相同的计算量下,ViT的性能通常优于ResNet。
图6:注意力图的示例。
图7:左:ViT-L/32 RGB值的初始线性嵌入filters。中:ViT-L/32位置嵌入的相似度。右:不同Head的平均注意力距离随网络深度的变化情况。
表3:预训练的超参数。
表4:微调的超参数。
表5:在ImageNet、ImageNet-21k、JFT-300M上进行预训练时,ViT在各种数据集上的Top1准确率。
表6:模型的可扩展性实验结果。
表7:使用Adam和SGD预训练的ResNet模型的微调结果。
图8:ViT模型不同维度的可扩展性实验结果。
图9:类别token和全局平均池化分类器的比较。两者效果相似,但需要不同的学习率。
表8:ViT-B/16模型的位置嵌入消融实验的结果。
图10:使用不同超参数训练的模型的位置嵌入相似度。
图11:不同Head的平均注意力距离随网络深度的变化情况。
图12:ViT模型的速度与同量级的ResNet相当,且内存效率明显更高。
图13:基于轴向注意力的模型的性能。
图14:更多注意力图的示例(随机选取)。
表9:ViT在VTAB-1k不同数据集上的结果。
4. 总结讨论
我们探索了Transformer在图像识别中的直接应用。与以往在计算机视觉中使用自注意力机制不同,除了最初的patches提取步骤之外,我们不会在架构中引入特定于图像的归纳偏置。相反,我们将图像视为一系列patches,并利用NLP中使用的标准Transformer编码器对其进行处理。这种简单可扩展的策略结合大规模数据集的预训练,效果出奇地好。因此,Vision Transformer在许多图像分类数据集上都达到或超过了SOTA水平,同时预训练成本相对较低。
虽然这些初步的结果令人鼓舞,但仍然存在许多挑战。一个挑战是将ViT应用于其他计算机视觉任务(如检测和分割任务)。我们的结果,以及Carion等人的结果,预示了这种方法的前景。另一个挑战是继续探索自监督的预训练方法。我们的初步实验结果表明,自监督预训练的效果有所改善,但与大规模有监督预训练相比仍存在较大差距。最后,进一步扩展ViT仍有可能带来模型性能的提高。
多模态人工智能
为人类文明进步而努力奋斗^_^↑
欢迎关注“多模态人工智能”公众号,一起进步^_^↑