【论文笔记】ViT:Transformer 能否实现 CV 与 NLP 的大一统?

737 阅读6分钟

这是我参与2022首次更文挑战的第19天,活动详情查看:2022首次更文挑战

本文发表于 ICLR 2021,第一作者 Alexey Dosovitskiy 来自 Google Research。
原文链接:An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

Preface

这篇文章就是大名鼎鼎的 Vision Transformer (ViT) 的出处,自 2021 年提出以后,ViT 成功刷榜各类视觉任务,其间也衍生出一些 ViT 的变体,比如目标检测任务的 SOTA 模型 Swin Transformer,该工作获 ICCV 2021 最佳论文奖。

【论文笔记】Swin Transformer:一种优异的计算机视觉骨干网络 - 掘金 (juejin.cn)

Motivation

Transformer 已经在 NLP 领域大放异彩,虽然 CV 中也由 Transformer 的应用,但是现有的工作总是离不开 CNN,注意力要么与 CNN 结合,要么用于替换 CNN 的某些组件。作者认为对 CNN 的依赖是不必要的,用纯 Transformer 也能很好地胜任图片分类任务。但是纯 Transformer 应用在视觉中也不是第一次了,所以作者又从可扩展性的角度考虑,希望能直接将一个标准的 Transformer 应用于 CV,以维持 Transformer 良好的可扩展性。

Related Work

也有人尝试在视觉任务中用 Transformer 替代 CNN。由于图片输入相对于文本尺寸太大,因此直接把二维图片打成 patch 作为 Transformer 输入是不现实的,所以过去的工作尝试了几种方法来处理这个问题:

  • Parmar 等人只在每个 query 像素的局部邻域应用了自注意力;
  • Child 等人尝试了一种稀疏 Transformer ,采用全局自注意力的可缩放近似。

但是这些方法过于复杂,不便于硬件加速,因此很难训练大模型。所以在大规模的图像识别上,传统的残差网络还是效果最好的。

与作者的工作最接近的是 Cordonnier 等人在 ICLR 2020 上发表的文章,相比于本文,用的是 2x2 的小 patch,在 CIFAR-10 上训练,图片尺寸为 32x32,但是该工作用的数据集图片尺寸太小,模型输入的 patch 尺寸也太小,不适合现代的视觉任务。本文的工作虽然跟它思路一样,但是得到了很多很有意义的结论,证明了标准 Transformer 经过大规模的预训练能够达到比 ResNet 更好的效果。

此外,2020 年有个 image GPT 的工作,也应用了 Transformer,通过大规模无监督学习,经过微调或者用来作为特征提取器,在 ImageNet 上最高也只能达到 72% 的准确率。

Model

将一张图片划分成 16x16 的小 patch,把这些小 patch 通过全连接层转变成一个个的向量,再把这些向量当成 NLP 任务里的单词,在序列首部加上 位置嵌入,送入 Transformer Encoder。作者尽可能遵循最初的 Transformer 结构,以保证模型的可扩展性。模型结构如下:

image.png

其中,位置嵌入是一个很值得关注的地方。作者尝试了几种不同的位置嵌入(1D/2D/相对),发现效果都差不多,作者可视化位置编码发现即便是使用1D位置编码,在自注意过程中模型已经学习到了2D的位置信息。

Pre-training

由于没有使用 CNN 中的先验知识(归纳偏置),ViT 需要大量数据进行预训练才能达到比较好的效果。作者在 JFT-300M 数据集上进行了预训练,并且在消融实验中展示了模型表现与数据集大小、计算预算等的关系。

image.png

Fine-tuning

在大型数据集上预训练的 ViT 用于下游任务时,图片尺寸可能发生变化。保持 patch 的大小不变,那么 ViT 输入序列的长度就会发生变化。作者直接 patch 在原始图像中的位置执行预训练位置嵌入的2D插值(直接用 pytorch 官方已有的实现)。这里也是 ViT 中唯一用到归纳偏置的地方。

Experiment

Compared with Baseline

image.png

ViT-Huge 在所有数据集上都取得了 SOTA 性能,虽然高得不多,但是训练成本比其他模型低很多,充分体现了 ViT 的优越性。

Ablation Study

作者做了大量的消融实验。

image.png

在小型数据集上训练的 ViT 性能远不如 ResNet,但是数据集提升到 ImageNet-21k 的时候,两者性能相当;进一步增大数据集,换成 JFT-300M 时,ViT 的性能已经普遍媲美甚至优于 ResNet。而且,图上并没有显示出性能随数据集规模增大而趋于饱和的迹象。这也说明了利用了先验知识的 CNN 在数据量不那么大的时候是有优势的。

上图的右边是作者用从 JFT-300M 中抽取的子集来预训练模型,这样做可以保持几个不同规模的数据集的分布比较一致(因为是从一个数据集中抽取的),这样受数据集影响比较小,更能体现模型本身的特性。

image.png

作者还研究了 ViT 和 ResNet 的表现与算力的关系,体现了 ViT 更容易训练的优点。

在少量计算资源的情况下,加入了卷积区域特征的 ViT (即图中的 Hybrid )体现出了优势,但在更大的模型中这种优势逐渐消失。这有点不符合直觉,按理说卷积区域特征应该在各种尺寸的模型中都能辅助 ViT。

Self-supervision

NLP 任务中 Transformer 的成功不仅仅因为可扩展性,还因为它对大规模自监督学习的支持。作者希望 ViT 也能支持自监督学习,为此做了一个初步的实验,模拟 BERT 中的 MLM,掩蔽 patch 并通过周围的图像块预测(Masked Patch Prediction),结果只有79.9%的准确率,这也是作者希望能进一步探究的方向。

Summary

作者直接将标准 Transformer 用于图像识别,没有做任何针对图像任务的修改,也没有利用图像领域惯用的先验知识。这种方法经过大规模预训练之后效果很好,超越了图像分类任务的 baseline,而且相对来讲更容易训练。

作者还展望了未来,一方面希望将 ViT 用在其它视觉任务中;另一方面,作者希望继续探索自监督学习,因为当前的自监督学习与大规模监督学习相比差距还很大。

作者团队后来又做了一个 ViT-G,进一步扩大了模型,也取得了更好地效果。原文链接:Scaling Vision Transformers