VLM经典论文阅读:【ViT】An image is worth 16X16 words: transformers for image recognition at scale
一、论文简介
1.1 ViT论文基本信息
题目:An image is worth 16X16 words: transformers for image recognition at scale
单位:OpenAI
会议:ICLR 2021
论文摘要:尽管 Transformer 结构已成自然语言处理任务的普遍标准,但其在计算机视觉领域的应用仍较为有限。在视觉任务中,注意力机制通常与卷积网络结合使用,或者替代卷积网络的部分组件,同时保持其总体结构不变。本文指出,这种对卷积神经网络的依赖并非必要,直接将纯粹的 Transformer 应用于图像块序列即可在图像分类任务上表现出色。当在大规模数据集上进行预训练并迁移到多个中小型图像识别基准测试(如 ImageNet、CIFAR-100、VTAB 等)时,Vision Transformer(ViT)能够以明显较少的训练计算资源,达到与当前最先进卷积网络相当的效果。
图1:模型概述。本文将输入图像分割为固定尺寸的图像块,对每个图像块进行线性嵌入,加入位置嵌入后,将生成的向量序列输入标准的 Transformer 编码器。为了完成分类任务,本文采用常规做法,在序列中添加一个可学习的“分类标记”作为额外输入。Transformer 编码器的设计参考了 Vaswani 等(2017)的方案。
二、ViT论文方法论
在模型设计方面,本文尽可能严格地遵循原始 Transformer(Vaswani 等,2017)。这种有意保持简洁的架构的一个好处是,可以直接利用可扩展的 NLP Transformer 架构及其高效实现,几乎不需要额外修改。
2.1 Vision Transformer(ViT)
本文在图 1 中展示了模型的整体架构。标准 Transformer 的输入是一维的 token 嵌入序列。为了处理二维图像,本文将图像 \( x \in \mathbb{R}^{H \times W \times C} \) 重塑为一系列二维图像块,并将其展平为序列 \( x_p \in \mathbb{R}^{N \times (P^2 \cdot C)} \) 。其中,(H,W) 为原始图像分辨率,C 为通道数,(P,P) 为每个图像块的分辨率,而 \( N = HW/P^2 \) 为图像块数量,同时也是 Transformer 的实际输入序列长度。Transformer 在所有层中使用固定维度 D 的隐含表示,因此本文将每个展平的 patch 通过一个可训练的线性映射转换到 D 维(公式 1),并称其为 patch 嵌入。
类似于 BERT 的 [class] 标记,本文在嵌入后的 patch 序列前添加了一个可学习的嵌入(记为 \( z_0^0 = x_{class} \) ),该嵌入在 Transformer 编码器输出端的表示( \( z_0^L \) )被用作图像表示 y(公式 4)。在预训练和微调阶段,本文均在 \( z_0^L \) 上附加一个分类头:预训练阶段使用带有单个隐藏层的 MLP,微调阶段则使用单层线性分类头。
为了保留位置信息,本文在 patch 嵌入中添加了位置嵌入。研究采用了标准的可学习一维位置嵌入,因为使用更复杂的位置嵌入并未显著提升性能(附录 D.4)。最终生成的嵌入向量序列作为编码器的输入。
Transformer 编码器(Vaswani 等,2017)由多个头部的自我注意层(MSA,见附录 A)与多层感知机(MLP)模块交替叠加而成(公式 2、3)。每个模块前执行 LayerNorm(LN),每个模块后则添加残差连接,该结构布局沿用了 Wang 等(2019)和 Baevski 与 Auli(2019)的设计。
MLP 包含两层,每层都有 GELU 非线性激活函数。
\( z_0 = [x_{\text{class}}; x_p^1 E; x_p^2 E; \cdots ; x_p^N E] + E_{\text{pos}}, \qquad E \in \mathbb{R}^{(P^2 \cdot C) \times D},\ E_{\text{pos}} \in \mathbb{R}^{(N+1)\times D} \) (1)
\( z'_{\ell} = \text{MSA}(\text{LN}(z_{\ell-1})) + z_{\ell-1}, \qquad \ell = 1 \ldots L \) (2)
\( z_{\ell} = \text{MLP}(\text{LN}(z'_{\ell})) + z'_{\ell}, \qquad \ell = 1 \ldots L \) (3)
\( y = \text{LN}(z^0_L) \) (4)
关于归纳偏置 (Inductive bias),文章提到 Vision Transformer 比起卷积神经网络(CNN)拥有较少的图像特异性归纳偏置。在 CNN 架构中,局部性、二维邻域结构及平移不变性在模型各层中被内建。然而,在 Vision Transformer 中,仅 MLP 层具备局部性和平移不变性,而自我注意层则是全局的。二维邻域结构的应用极为有限:一是在模型初始阶段将图像分割成多个小块,二是在微调过程中为适应不同分辨率的图像调整位置嵌入(如后文所述)。此外,位置嵌入在初始化时并不包含有关小块二维位置的信息,所有小块间的空间关系都需要从头开始学习。
对于混合架构 (Hybrid Architecture),文章建议输入序列可以由卷积神经网络(LeCun 等,1989)提取的特征图组成,而非直接采用原始图像的小块。在这种混合模型中,小块嵌入投影 E(公式 1)用于从卷积神经网络特征图中提取的小块。作为特殊情况,小块的空间尺寸可以设为 1×1,此时输入序列通过简单展开特征图的空间维度并映射至 Transformer 的维度来获取。分类输入嵌入与位置嵌入的添加方法与上述相同。
2.2 微调与更高分辨率(Fine-tuning and higher resolution)
一般而言,文章在大型数据集上预先训练 ViT,并在较小的下游任务上进行微调。为此,文章移除了预训练的预测头,并连接一个零初始化的 D×K 前馈层,其中 K 表示下游任务的类别数。与 Touvron 等(2019)和 Kolesnikov 等(2020)的研究结果一致,以高于预训练的分辨率进行微调通常更为有利。当输入更高分辨率的图像时,文章保持小块大小不变,导致有效的序列长度增加。Vision Transformer 能够处理任意长度的序列(受限于内存),但预训练的位置嵌入可能不再适用。因此,文章基于小块在原始图像中的位置,对预训练的位置嵌入进行了二维插值。值得注意的是,这种分辨率调整和小块提取,是文章在人工设计中引入二维图像结构归纳偏置的唯一环节。
三、论文实验 ???? 
文章考察了 ResNet、视觉变换器(ViT)以及组合模型的表示学习效能。为了掌握各模型的数据需求,研究在不同规模的数据集上实施了预训练,并在多种标准任务中进行了评估。就预训练的运算成本而言,ViT 展现了卓越的表现,以较低的预训练成本在大多数识别基准上达到了或接近了最先进水平。最终,研究开展了一个小型的自我监督实验,结果显示自我监督版 ViT 具有巨大的发展潜力。
3.1 设置(Setup)
在数据集(Datasets)方面,为了探讨模型的可延展性,研究采用了包含 1k 类别、130 万张影像的 ILSVRC-2012 ImageNet 数据集(以下简称 ImageNet),其扩展集 ImageNet-21k(21k 类别、1400 万张影像,Deng 等人,2009),以及 JFT(Sun 等人,2017),该数据集包含 18k 类别、3.03 亿张高清影像。研究根据 Kolesnikov 等人(2020)的方法,对预训练数据集进行了去重处理,以防与各下游任务的测试集重复。研究将基于这些数据集训练的模型迁移到多个基准任务中进行评估,涵盖:使用原始验证集标签与清理后的 ReaL 标签的 ImageNet(Beyer 等人,2020)、CIFAR-10/100(Krizhevsky,2009)、Oxford-IIIT Pets(Parkhi 等人,2012)以及 Oxford Flowers-102(Nilsback 和 Zisserman,2008)。这些数据集的预处理流程遵循了 Kolesnikov 等人(2020)的规定。
在模型变体(Model Variants)方面,研究依据 BERT(Devlin 等人,2019)的配置构建 ViT,不同的配置在表 1 中有所总结。“基础” 与 “大型” 模型直接采用了 BERT 的设定,并引入了更大规模的 “巨大” 模型。后文中研究使用缩写来标识模型规模与输入块尺寸,比如 ViT-L/16 表示使用 16×16 块尺寸的 “大型” 模型。值得注意的是,Transformer 的序列长度与块尺寸的平方呈反比关系,因此块尺寸越小,模型的计算成本越高。
表1 视觉变换器模型各变体的具体信息
对于基线 CNN,研究选择了 ResNet(He 等人,2016),但将批归一化(Ioffe 和 Szegedy,2015)替换为组归一化(Wu 和 He,2018),并应用了标准化卷积(Qiao 等人,2019)。这些调整能增强迁移性能(Kolesnikov 等人,2020),研究将这种改良后的模型命名为 “ResNet (BiT)”。对于混合架构,研究将中间特征图以单个“像素”的块尺寸输入 ViT。为了探究不同序列长度的影响,研究采用了两种配置:(i) 使用常规 ResNet-50 第 4 阶段的输出;或 (ii) 删除第 4 阶段,将相同数量的层加入第 3 阶段(保持总层数不变),并利用扩展后的第 3 阶段输出。方案 (ii) 将产生长度增加四倍的序列,从而形成计算成本更高的 ViT 模型。
训练与微调。本文对所有模型,包括 ResNet,均采用了 Adam 优化器(Kingma 和 Ba,2015)进行训练,参数设定为 β1 = 0.9,β2 = 0.999,批量大小为 4096,并施加较高的权重衰减 0.1,这已被证实对所有模型的迁移表现有益(附录 D.1 显示,在本文设定下,Adam 对 ResNet 的效果略胜过 SGD,与常见做法相比)。训练过程中采用了线性学习率预热和衰减策略,具体内容参见附录 B.1。在微调阶段,本文使用带动量的 SGD,批量大小为 512,适用于所有模型,详见附录 B.1.1。针对表 2 中的 ImageNet 结果,微调时采用了更高分辨率:ViT-L/16 为 512,ViT-H/14 为 518,并采用了 Polyak 和 Juditsky (1992) 平均方法,因子为 0.9999(参见 Ramachandran 等人,2019; Wang 等人,2020b)。
评估标准。本文在下游数据集上汇报结果,方式为少量样本(few-shot)或微调(fine-tuning)精确度。微调精确度体现了每个模型在其特定数据集上微调后的效能。少量样本精确度通过解决正则化最小二乘回归问题获取,该问题将训练图像子集的(固定的)表达映射到{?1, 1}^K目标向量。这种方法允许以封闭形式精确求解。尽管本文主要聚焦于微调表现,但在微调成本过高时,有时会利用线性少量样本精确度进行快速评估。
3.2 与最先进方法的对比(Comparison to SOTA)
首先,本文将最大的模型——ViT-H/14和ViT-L/16——与文献中的顶尖CNN进行对比。首个对比对象是大型转移(BiT)(Kolesnikov等人,2020),它采用大规模ResNet进行监督迁移学习。第二个对比对象是噪声学生(Noisy Student)(Xie等人,2020),该模型为大规模EfficientNet,通过在ImageNet和JFT300M上的半监督学习训练而成,去除了标签。目前,噪声学生在ImageNet上处于领先水平,而BiT-L在本文报道的其他数据集上表现最佳。所有模型均在TPUv3硬件上训练,文中报告了预训练每个模型所需的TPUv3核天数,即用于训练的TPUv3核心数量(每片芯片2个)乘以训练天数。
表2展示了结果。较小的ViT-L/16模型在JFT-300M上预训练后,在所有任务上均优于同样在该数据集上预训练的BiT-L,同时所需的计算资源明显较少。更大的ViT-H/14模型进一步提高了性能,特别是在更具挑战性的数据集上——ImageNet、CIFAR-100以及VTAB套件。有趣的是,该模型的预训练计算量仍明显低于之前的最先进方法。然而,本文指出,预训练效率不仅可能受架构选择的影响,还可能受到训练计划、优化器、权重衰减等因素的影响。本文在第4.4节提供了不同架构在性能与计算量之间对比的研究。最后,在公开的ImageNet-21k数据集上预训练的ViT-L/16模型在多数数据集上表现出色,同时预训练所需资源更少:使用标准云TPUv3(8核)大约可在30天内完成训练。
表2 与流行图像分类基准上的最先进方法对比。报告了准确率的平均值和标准差,取自三次微调运行的平均值。在JFT-300M数据集上预训练的Vision Transformer模型在所有数据集上均优于基于ResNet的基线,同时预训练所需的计算资源显著较少。在较小的公开ImageNet-21k数据集上预训练的ViT模型也表现出色。?Touvron等人(2020)报告的88.5%结果略有提高。
图2将VTAB任务分为各自的类别,并与该基准上的先前最先进方法进行对比:BiT、VIVI——在ImageNet和YouTube上联合训练的ResNet(Tschannen等人,2020),以及S4L——在ImageNet上进行监督和半监督学习的方法(Zhai等人,2019a)。ViT-H/14在自然类(Natural)和结构类(Structured)任务上优于BiT-R152x4及其他方法。在专业类(Specialized)任务上,排名前两位的模型性能相近。
图2 VTAB在自然类(Natural)、专业类(Specialized)和结构类(Structured)任务组中的性能分解。
3.3 预训练数据需求(Pre-Training Data Requirements)
Vision Transformer在大规模JFT-300M数据集上预训练时表现优异。鉴于相较于ResNet,ViT在视觉任务中具有较少的归纳偏差,那么数据集规模的重要性如何?本文进行了两组实验。
首先,在不同规模的数据集上预训练ViT模型:ImageNet、ImageNet-21k和JFT-300M。为了提升在小规模数据集上的表现,本文优化了三个基础正则化参数——权重衰减、dropout和标签平滑。图3展示了微调至ImageNet后的结果(其他数据集结果见表5)。在最小的数据集ImageNet上预训练时,ViT-Large模型的表现低于ViT-Base模型,即使进行了(适度)正则化。使用ImageNet-21k预训练时,其表现相仿。只有在JFT-300M上预训练时,才能充分发挥大模型的优势。图3还展示了不同规模BiT模型的性能范围。BiT CNN在ImageNet上优于ViT,但在更大规模数据集上,ViT实现了超越。

图3 在ImageNet上的迁移表现。当在小规模数据集上预训练时,大型ViT模型的表现逊色于BiT ResNet(阴影区域),但在大规模数据集上预训练时显著提升。类似地,随着数据集规模的扩大,较大的ViT变体超越了较小的模型。
其次,本文在JFT-300M数据集的随机子集上训练模型,子集规模为9M、30M和90M,以及完整的JFT-300M数据集。在较小的子集上不采用额外的正则化措施,并对所有设置使用统一的超参数。这样可以评估模型的内在特性,而非正则化的影响。然而,本文采用了早停策略,并报告训练过程中获得的最佳验证准确率。为节约计算资源,报告的是少样本线性准确率而非完全微调后的准确率。图4展示了结果。相较于计算成本相当的ResNet,Vision Transformer在小规模数据集上更易发生过拟合。例如,ViT-B/32略快于ResNet50;在9M子集上表现明显不佳,但在90M及以上子集上表现更优。ResNet152x2与ViT-L/16也呈现出相同趋势。这一结果进一步证实了卷积归纳偏置对小数据集有用的观点,但在大数据集上,直接从数据中学习相关模式既充分又有效。
总体来说,ImageNet上的少样本结果(图4)以及VTAB上的低数据结果(表2)展现了在极低数据量下的迁移具有良好的潜力。对ViT少样本特性的深入分析将是未来研究的一个有趣方向。
图4 在ImageNet上的线性少样本评估与预训练数据集规模的关系。ResNet在小规模预训练数据集上表现更佳,但性能迅速达到上限,而ViT在较大预训练数据集上表现更出色。ViT-b指的是将ViT-B的所有隐藏维度减半的模型。
3.4 规模研究(Scaling Study)
本文对不同模型进行了控制下的规模研究,通过评估其从JFT-300M上的迁移性能来对比。在此设定下,数据规模不会成为模型性能的障碍,本文评估了每个模型的性能与预训练成本的关系。模型组合包括:7个ResNet,分别是R50x1、R50x2、R101x1、R152x1、R152x2(预训练7个周期),以及R152x2和R200x3(预训练14个周期);6个Vision Transformer,分别是ViT-B/32、B/16、L/32、L/16(预训练7个周期),以及L/16和H/14(预训练14个周期);5个混合模型,分别是R50+ViT-B/32、B/16、L/32、L/16(预训练7个周期),以及R50+ViT-L/16(预训练14个周期)(对于混合模型,名称末尾的数字表示ResNet骨干中的总下采样比例,而非patch尺寸)。
图5展示了迁移性能与总预训练计算量的关系(计算成本详情见附录D.5),各模型的具体结果见附录表6。可以观察到几个规律。首先,Vision Transformer在性能/计算比率上优于ResNet,达到相同性能所需的计算量大约是ResNet的2?4倍(5个数据集的平均值)。其次,混合模型在小计算预算下略微优于ViT,但在较大模型上这种差异消失。这一结果有些出乎意料,因为人们可能预期卷积局部特征处理在任何规模下都能辅助ViT。第三,Vision Transformer在尝试的规模范围内似乎没有达到饱和状态,这为未来的模型扩展提供了动力。
图5 不同架构的性能与预训练计算量关系:Vision Transformer、ResNet及混合模型。Vision Transformer在相同计算预算下通常优于ResNet。混合模型在较小模型规模下优于纯Transformer,但在较大模型上这一差距消失。
3.5 Vision Transformer模型分析(Inspecting Vision Transformer)
为了理解Vision Transformer如何处理图像数据,本文分析了其内部表示。Vision Transformer的第一层将展平的图像patch线性映射到低维空间(公式1)。图7(左)显示了学习到的嵌入滤波器的主要成分,这些成分类似于表示每个patch细粒度结构的低维基函数。
映射之后,学习到的位置嵌入被添加到patch表示中。图7(中)显示模型通过位置嵌入的相似性编码图像中的距离,即相邻的patch具有更相似的位置嵌入。此外,还出现了行列结构;同一行或列的patch具有相似的嵌入。对于较大的网格,有时还能观察到正弦结构(附录D)。位置嵌入能够学习表示二维图像拓扑,这解释了为什么精心设计的二维感知嵌入变体未能提高性能(附录D.4)。
自注意力机制让ViT能够从底层就开始融合整张图像的数据。本文探讨了网络在多大程度上运用了这种能力。具体来说,根据注意力权重来计算信息融合的平均图像空间距离(图7,右),这个“注意力范围”类似于CNN中的感受野尺寸。研究发现,某些注意力头在最低层就能覆盖图像的广泛区域,这表明模型确实利用了全局信息融合的能力。而其他注意力头在底层中则始终保持着较短的注意力范围,这种高度局部化的注意力在先应用ResNet再接入Transformer的混合模型中不那么显著(图7,右),暗示其功能可能与CNN的早期卷积层类似。此外,随着网络层次加深,注意力范围也逐渐扩大。整体来看,模型关注的图像区域在语义上与分类任务相关(图6)。
图6 展示了从输出token到输入空间的注意力代表性例子。详情参阅附录D.7
图7 左侧展示了ViT-L/32对RGB值的初始线性嵌入过滤器。中间部分展示了ViT-L/32位置嵌入的相似性。每个小格子显示指定行列表面贴片的位置嵌入与其他所有表面贴片位置嵌入的余弦相似度。右侧展示了各个注意力头随网络深度变化的关注区域大小。每个点代表该层16个注意力头之一在所有图像上的平均注意力距离。详情参阅附录D.7
3.6 自监督学习(Self-supervision)
Transformer在自然语言处理任务上表现出色。不过,它的成功不仅仅归因于出色的可扩展性,还依赖于大规模的自监督预训练(Devlin et al., 2019;Radford et al., 2018)。本文还初步尝试了自监督的掩码patch预测,模仿BERT中采用的掩码语言建模任务。通过自监督预训练,较小的ViT-B/16模型在ImageNet上达到了79.9%的精度,相较于从零开始训练提高了大约2%,但仍比有监督预训练低大约4%。附录B.1.2提供了更多信息。对比式预训练(Chen et al., 2020b;He et al., 2020;Bachman et al., 2019;Hénaff et al., 2020)的探索留待未来工作。
结论
本文探讨了Transformer在图像识别中的直接应用。与之前在计算机视觉中使用自注意力的工作不同,除了最初的patch提取步骤外,本文没有在架构中引入图像特定的归纳偏置。相反,本文将图像视为patch序列,并通过NLP中常用的Transformer编码器进行处理。这种简洁且可扩展的方法在大规模数据集预训练下表现得异常出色。因此,Vision Transformer在多个图像分类数据集上达到了或超越了最先进的水平,同时预训练的成本相对较低。
尽管这些初步成果令人振奋,但仍面临诸多挑战。首先,是将ViT应用于其他计算机视觉任务,如目标检测和图像分割。本文的结果结合Carion et al. (2020)的研究表明这种方法具有潜力。其次,是继续探索自监督预训练方法。初步实验证明自监督预训练可以提高性能,但与大规模有监督预训练仍有较大差距。最后,进一步扩大ViT规模可能会进一步提升性能。
其他待补充。


雷达卡


京公网安备 11010802022788号







