用于大规模图像识别的 Transformer

虽然卷积神经网络(CNN)自 20 世纪 80 年代以来就已用于计算机视觉领域,但直到 2012 年AlexNet 的表现远超当时最先进的图像识别方法时,它才处于领先地位。有两个因素促成了这一突破:(i) 像ImageNet这样的训练集的可用性,以及 (ii) 商品化 GPU 硬件的使用,这为训练提供了更多的计算能力。因此,自 2012 年以来,CNN 已成为视觉任务的首选模型。

使用 CNN 的好处是,它们无需手工设计视觉特征,而是直接从数据“端到端”学习执行任务。然而,虽然 CNN 避免了手工提取特征,但架构本身是专门为图像设计的,计算量很大。展望下一代可扩展视觉模型,人们可能会问,这种领域特定的设计是否必要,或者我们是否可以成功利用更多与领域无关且计算效率更高的架构来实现最先进的结果。

作为朝这个方向迈出的第一步,我们提出了Vision Transformer (ViT),这是一种尽可能接近最初为基于文本的任务设计的Transformer架构的视觉模型。ViT 将输入图像表示为图像块序列,类似于将 Transformers 应用于文本时使用的词嵌入序列,并直接预测图像的类标签。在对足够的数据进行训练时,ViT 表现出色,其性能优于同类最先进的 CNN,而计算资源却少了四倍。为了促进该领域的进一步研究,我们已经开源了代码和模型。

视觉转换器

原文 Transformer 将单词序列作为输入,然后将其用于分类、翻译或其他 NLP 任务。对于 ViT,我们尽可能少地修改 Transformer 设计,使其直接对图像而不是单词进行操作,并观察模型可以自行学习多少有关图像结构的知识。

ViT 将图像划分为方形块网格。通过连接块中所有像素的通道,然后将其线性投影到所需的输入维度,将每个块展平为单个向量。由于 Transformers 与输入元素的结构无关,因此我们向每个块添加可学习的位置嵌入,这允许模型了解图像的结构。先验地,ViT 不知道图像中块的相对位置,甚至不知道图像具有 2D 结构 — 它必须从训练数据中学习此类相关信息,并在位置嵌入中编码结构信息。

扩大规模

我们首先在 ImageNet 上训练 ViT,它取得了 77.9% 的 top-1 准确率。虽然这对于首次尝试来说还不错,但与最先进的水平相差甚远——目前在 ImageNet 上训练的最佳 CNN 在没有额外数据的情况下达到 85.8%。尽管采取了缓解策略(例如正则化),但由于缺乏关于图像的内置知识,ViT 仍会过度拟合 ImageNet 任务。

为了研究数据集大小对模型性能的影响,我们在ImageNet-21k(1400 万张图像,21000 个类别)和JFT(3 亿张图像,18000 个类别)上训练 ViT,并将结果与在相同数据集上训练的最先进的 CNN Big Transfer (BiT) 进行比较。如前所述,在 ImageNet(100 万张图像)上训练时,ViT 的表现明显差于 CNN 等效模型(BiT)。然而,在 ImageNet-21k(1400 万张图像)上,性能相当,而在 JFT(3 亿张图像)上,ViT 现在的表现优于 BiT。

最后,我们研究了模型训练中涉及的计算量的影响。为此,我们在 JFT 上训练了几种不同的 ViT 模型和 CNN。这些模型涵盖了各种模型大小和训练时长。因此,它们需要不同数量的计算来进行训练。我们观察到,对于给定的计算量,ViT 的性能优于同等的 CNN。

1729583903035.jpg

左图: ViT 在不同数据集上进行预训练时的性能。右图: ViT 实现了良好的性能/计算平衡。

高性能大规模图像识别

我们的数据表明:(1) 经过充分训练后,ViT 可以表现得非常好;(2) ViT 在较小和较大的计算规模下都能实现出色的性能/计算平衡。因此,为了查看性能改进是否能延续到更大的规模,我们训练了一个 600M 参数的 ViT 模型。

这种大型 ViT 模型在多个流行基准上都取得了最佳性能,包括 ImageNet 上 88.55% 的 top-1 准确率和 CIFAR-10 上 99.50% 的准确率。ViT 在ImageNet 评估集“ImageNet-Real”的清理版本上也表现良好,达到了 90.72% 的 top-1 准确率。最后,即使使用很少的训练数据点,ViT 也能很好地完成各种任务。例如,在VTAB-1k 套件(19 个任务,每个任务有 1,000 个数据点)中,ViT 达到了 77.63%,远远领先于单一模型的最先进水平 (SOTA)(76.3%),甚至与多个模型集合达到的SOTA (77.6%)相当。最重要的是,与之前的 SOTA CNN 相比,这些结果是使用更少的计算资源获得的,例如,比预先训练的 BiT 模型少 4 倍。

1729583890992.jpg

Vision Transformer 在流行基准测试中的表现与最先进的 CNN 相当甚至更胜一筹。左图:流行的图像分类任务(ImageNet,包括新的验证标签ReaL和CIFAR、Pets和Flowers)。右图: VTAB 分类套件中 19 个任务的平均值。

可视化

为了直观地了解模型的学习内容,我们将其内部的一些工作原理可视化。首先,我们查看位置嵌入(模型学习编码补丁相对位置的参数),发现 ViT 能够重现直观的图像结构。每个位置嵌入与同一行和列中的其他位置嵌入最相似,这表明模型已经恢复了原始图像的网格结构。其次,我们检查每个变压器块中一个元素与另一个元素之间的平均空间距离。在较高层(深度为 10-20)仅使用全局特征(即较大的注意距离),但较低层(深度为 0-5)捕获全局和局部特征,这由平均注意距离的较大范围表示。相比之下,CNN 的较低层中只存在局部特征。这些实验表明,ViT 可以学习硬编码到 CNN 中的特征(例如对网格结构的感知),但也可以自由学习更通用的模式,例如较低层的局部和全局特征的混合,这有助于泛化。

1729583879481.jpg

左图: ViT 通过其位置嵌入学习图像块的网格状结构。右图: ViT 的较低层包含全局和局部特征,较高层仅包含全局特征。

概括

虽然 CNN 彻底改变了计算机视觉,但我们的结果表明,为成像任务量身定制的模型可能是不必要的,甚至是次优的。随着数据集规模的不断增加,以及无监督和半监督方法的不断发展,开发能够在这些数据集上更有效地训练的新视觉架构变得越来越重要。我们相信 ViT 是迈向通用、可扩展架构的初步步骤,可以解决许多视觉任务,甚至来自许多领域的任务,我们对未来的发展感到兴奋。

我们的工作的预印本 以及代码和模型均已公开。

致谢

我们要感谢柏林、苏黎世和阿姆斯特丹的合著者:Alexey Dosovitskiy、Lucas Beyer、Alexander Kolesnikov、Xiaohua Zhai、Thomas Unterthiner、Mostafa Dehghani、Matthias Minderer、Georg Heigold、Sylvain Gelly 和 Jakob Uszkoreit。我们要感谢 Andreas Steiner 在基础设施和开源方面提供的关键帮助,感谢 Joan Puigcerver 和 Maxim Neumann 在大规模训练基础设施方面的工作,感谢 Dmitry Lepikhin、Aravindh Mahendran、Daniel Keysers、Mario Lučić、Noam Shazeer 和 Colin Raffel 的有益讨论。最后,我们要感谢 Tom Small 在本篇文章中创建了 Visual Transformer 动画。

版权声明

本文仅代表作者观点,不代表本站立场。
本文系作者授权发表,未经许可,不得转载。

评论