AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE 论文精读

题目:AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

题目:一图抵 16×16 个词:面向大规模图像识别的 Transformer 架构

论文地址:ICLR 2021 2010
代码地址:https://github.com/google-research/vision_transformer

AT SCALE 大规模

IMAGE RECOGNITION 图像识别

AN IMAGE IS WORTH 16X16 WORDS 一图抵 16×16 个词

目录

  • [题目:AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE](#题目:AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE)
  • [题目:一图抵 16×16 个词:面向大规模图像识别的 Transformer 架构](#题目:一图抵 16×16 个词:面向大规模图像识别的 Transformer 架构)
    • [ABSTRACT 摘要](#ABSTRACT 摘要)
    • [1 INTRODUCTION 引言](#1 INTRODUCTION 引言)
    • [2 RELATED WORK 相关工作](#2 RELATED WORK 相关工作)
    • [3 METHOD 方法](#3 METHOD 方法)
      • [3.1 VISION TRANSFORMER (VIT) 视觉Transformer(ViT)](#3.1 VISION TRANSFORMER (VIT) 视觉Transformer(ViT))
        • [3.1.1 ViT基础架构](#3.1.1 ViT基础架构)
          • [3.1.1.1 一维词元➡二维图像](#3.1.1.1 一维词元➡二维图像)
          • [3.1.1.2 VIT引入BERT](#3.1.1.2 VIT引入BERT)
          • [3.1.1.3 VIT引入位置嵌入](#3.1.1.3 VIT引入位置嵌入)
          • [3.1.1.4 VIT的Transformer编码器架构](#3.1.1.4 VIT的Transformer编码器架构)
        • [3.1.2 Inductive bias 归纳偏置](#3.1.2 Inductive bias 归纳偏置)
        • [3.1.3 Hybrid Architecture 混合架构](#3.1.3 Hybrid Architecture 混合架构)
      • [3.2 FINE-TUNING AND HIGHER RESOLUTION 微调与高分辨率适配](#3.2 FINE-TUNING AND HIGHER RESOLUTION 微调与高分辨率适配)
        • [3.2.1 ViT 的预训练 + 微调两阶段范式](#3.2.1 ViT 的预训练 + 微调两阶段范式)
        • [3.2.2 预训练和微调的图像分辨率不一致](#3.2.2 预训练和微调的图像分辨率不一致)
    • [4 EXPERIMENTS 实验](#4 EXPERIMENTS 实验)
      • [4.1 SETUP 实验设置](#4.1 SETUP 实验设置)
        • [4.1.1 Datasets 数据集](#4.1.1 Datasets 数据集)
        • [4.1.2 Model Variants 模型变体](#4.1.2 Model Variants 模型变体)
        • [4.1.3 Training & Fine-tuning 训练与微调](#4.1.3 Training & Fine-tuning 训练与微调)
        • [4.1.4 Metrics 评价指标](#4.1.4 Metrics 评价指标)
      • [4.2 COMPARISON TO STATE OF THE ART 与当前最优方法的对比](#4.2 COMPARISON TO STATE OF THE ART 与当前最优方法的对比)
      • [4.3 PRE-TRAINING DATA REQUIREMENTS 预训练的数据需求](#4.3 PRE-TRAINING DATA REQUIREMENTS 预训练的数据需求)
      • [4.4 SCALING STUDY 缩放性研究](#4.4 SCALING STUDY 缩放性研究)
      • [4.5 INSPECTING VISION TRANSFORMER 视觉Transformer的机制探查](#4.5 INSPECTING VISION TRANSFORMER 视觉Transformer的机制探查)
      • [4.6 SELF-SUPERVISION 自监督学习](#4.6 SELF-SUPERVISION 自监督学习)
    • [5 CONCLUSION 结论](#5 CONCLUSION 结论)

ABSTRACT 摘要

尽管Transformer架构已成为自然语言处理任务的事实标准,但其在计算机视觉领域的应用仍较为有限。在视觉领域中,注意力机制要么与卷积网络结合使用,要么在保留卷积网络整体结构的前提下替换其中部分组件。我们的研究表明,这种对卷积神经网络(CNN)的依赖并非必要,直接将纯Transformer应用于图像块序列,也能在图像分类任务上取得出色表现。

在大量数据上完成预训练后,将视觉Transformer(ViT)迁移至多个中等或小规模图像识别基准测试(如ImageNet、CIFAR-100、VTAB等),其结果不仅能与当前最优的卷积网络相媲美,同时训练所需的计算资源也大幅减少。

1 INTRODUCTION 引言

基于自注意力的架构,尤其是Transformer(Vaswani等人,2017),已成为自然语言处理(NLP)领域的首选模型。主流方法是在大规模文本语料库上进行预训练,然后在较小的特定任务数据集上进行微调(Devlin等人,2019)。得益于Transformer的计算效率和可扩展性,如今已能训练出参数量超过1000亿的前所未有的大模型(Brown等人,2020;Lepikhin等人,2020)。随着模型和数据集规模的不断增长,其性能仍未出现饱和迹象。

然而在计算机视觉领域,卷积架构仍占据主导地位(LeCun等人,1989;Krizhevsky等人,2012;He等人,2016)。受NLP领域成功的启发,多项研究尝试将类CNN架构与自注意力机制相结合(Wang等人,2018;Carion等人,2020),部分研究则完全替换了卷积操作(Ramachandran等人,2019;Wang等人,2020a)。后一类模型虽在理论上高效,但因采用了特殊的注意力模式,尚未能在现代硬件加速器上实现有效缩放。因此,在大规模图像识别任务中,经典的类ResNet架构仍处于业界领先水平(Mahajan等人,2018;Xie等人,2020;Kolesnikov等人,2020)。

受NLP领域Transformer缩放成功的启发,我们尝试将标准Transformer直接应用于图像,并尽可能减少修改。为此,我们将图像拆分为图像块,并将这些图像块的线性嵌入序列作为Transformer的输入。图像块的处理方式与NLP应用中的令牌(单词)完全一致。我们以监督学习的方式在图像分类任务上训练该模型。

当在中等规模数据集(如ImageNet)上训练且未使用强正则化时,这类模型的准确率表现一般,比同等规模的ResNet低几个百分点。这一看似令人沮丧的结果或许在预料之中:Transformer缺少CNN固有的一些归纳偏置,例如平移等变性和局部性,因此在数据量不足时泛化能力较差。

不过,若在更大规模的数据集(1400万至3亿张图像)上训练模型,情况则会发生变化。我们发现,大规模训练的效果胜过了归纳偏置的作用。我们的视觉Transformer(ViT)在足够规模的预训练后,迁移到数据量较少的任务时能取得优异结果。在公开的ImageNet-21k数据集或内部JFT-300M数据集上进行预训练后,ViT在多个图像识别基准测试上达到或超越了业界最优水平。其中,表现最佳的模型在ImageNet上准确率达88.55%,在ImageNet-ReaL上达90.72%,在CIFAR-100上达94.55%,在包含19个任务的VTAB套件上达77.63%。

Transformer由Vaswani等人(2017)提出,最初用于机器翻译,此后已成为众多自然语言处理(NLP)任务的主流方法。基于Transformer的大型模型通常先在大规模语料库上进行预训练,再针对具体任务进行微调:BERT(Devlin等人,2019)采用去噪自监督预训练任务,而GPT系列工作则以语言建模作为预训练任务(Radford等人,2018、2019;Brown等人,2020)。

若将自注意力直接应用于图像,需要让每个像素都关注其他所有像素,其计算复杂度随像素数量呈二次增长,无法适配实际的输入尺寸。因此,过去在图像处理场景中应用自注意力时,曾尝试过多种近似方法。Parmar等人(2018)仅让每个查询像素在局部邻域内应用自注意力,而非全局注意力。这类局部多头点积自注意力模块可以完全替代卷积操作(Hu等人,2019;Ramachandran等人,2019;Zhao等人,2020)。在另一项工作中,稀疏Transformer(Child等人,2019)采用可扩展的近似方法处理全局自注意力,使其能适配图像任务。另一种缩放注意力的可行方式是分块、按不同尺寸应用注意力(Weissenborn等人,2019),极端情况下甚至仅沿单个轴计算注意力(Ho等人,2019;Wang等人,2020a)。尽管许多这类专用注意力架构在计算机视觉任务上展现出良好前景,但需要复杂的工程优化才能在硬件加速器上高效运行。

与我们工作最相关的是Cordonnier等人(2020)提出的模型,该模型从输入图像中提取2×2大小的图像块,并在其之上应用全自注意力。这一模型与ViT非常相似,但我们的工作进一步证明:大规模预训练能让普通Transformer与(甚至优于)当前最优的卷积神经网络(CNN)相抗衡。此外,Cordonnier等人(2020)采用的图像块尺寸仅为2×2像素,导致模型仅适用于低分辨率图像,而我们的模型也能处理中等分辨率图像。

学界也有大量研究关注将卷积神经网络(CNN)与自注意力形式相结合,例如:通过增强特征图实现图像分类(Bello等人,2019),或利用自注意力进一步处理CNN的输出,应用于目标检测(Hu等人,2018;Carion等人,2020)、视频处理(Wang等人,2018;Sun等人,2019)、图像分类(Wu等人,2020)、无监督目标发现(Locatello等人,2020)或统一文本-视觉任务(Chen等人,2020c;Lu等人,2019;Li等人,2019)等场景。

另一项近期相关工作是图像GPT(iGPT)(Chen等人,2020a),该模型通过降低图像分辨率和颜色空间,将Transformer应用于图像像素。模型以无监督生成式方式训练,得到的特征表示可后续用于微调或线性探测以实现分类,在ImageNet上最高准确率达72%。

我们的工作补充了越来越多的研究成果,这些研究在比标准ImageNet数据集更大的规模上探索图像识别。利用额外数据源,能够在标准基准测试上取得当前最优结果(Mahajan等人,2018;Touvron等人,2019;Xie等人,2020)。此外,Sun等人(2017)研究了CNN性能随数据集规模的变化规律,Kolesnikov等人(2020)、Djolonga等人(2020)则对从ImageNet-21k和JFT-300M等大规模数据集上学习的CNN迁移能力进行了实证探索。我们同样聚焦于后两个数据集,但训练的是Transformer,而非前人工作中使用的基于ResNet的模型。

3 METHOD 方法

在模型设计上,我们尽可能严格地遵循了原始Transformer架构(Vaswani等人,2017)的设计规范。这种刻意简化的设计方案的优势在于:自然语言处理(NLP)领域中各类可扩展的Transformer架构,以及对应的高效实现代码,几乎都可以开箱即用。

3.1 VISION TRANSFORMER (VIT) 视觉Transformer(ViT)

3.1.1 ViT基础架构
3.1.1.1 一维词元➡二维图像

模型整体结构如图1所示。标准Transformer的输入是一维的词元嵌入序列。

为了处理二维图像,我们将输入图像 x ∈ R H × W × C \boldsymbol{x} \in \mathbb{R}^{H \times W \times C} x∈RH×W×C重塑为一系列展平的二维图像块序列 x p ∈ R N × ( P 2 ⋅ C ) \boldsymbol{x}_p \in \mathbb{R}^{N \times (P^2 \cdot C)} xp∈RN×(P2⋅C)。

其中 ( H , W ) (H, W) (H,W)是原始图像的分辨率, C C C是图像通道数; ( P , P ) (P, P) (P,P)是每个图像块的分辨率, N = H W / P 2 N = HW/P^2 N=HW/P2是最终得到的图像块数量,该数值同时也作为Transformer的有效输入序列长度。

Transformer在所有层中均使用固定维度的隐向量 D D D,因此我们先将所有图像块展平,再通过一个可训练的线性投影将其映射到 D D D维空间(公式1)。我们将该投影的输出称为图像块嵌入(patch embeddings)

3.1.1.2 VIT引入BERT

与BERT中的 class 词元类似,我们在图像块嵌入序列的首位拼接一个可学习的嵌入( z 0 0 = x class \boldsymbol{z}0^0 = \boldsymbol{x}{\text{class}} z00=xclass),其在Transformer编码器输出端的状态( z L 0 \boldsymbol{z}_L^0 zL0)将作为最终的图像表征 y \boldsymbol{y} y(公式4)。

在预训练和微调阶段,我们都会在 z L 0 \boldsymbol{z}_L^0 zL0后接一个分类头:预训练阶段,该分类头由一个带单隐藏层的多层感知机(MLP)实现;微调阶段则仅使用单个线性层实现。

3.1.1.3 VIT引入位置嵌入

我们在图像块嵌入中加入位置嵌入,以保留图像的位置信息。本文使用标准的可学习一维位置嵌入,因为实验中我们并未发现使用更复杂的、感知二维结构的位置嵌入能带来显著的性能提升(见附录D.4)。最终得到的嵌入向量序列将作为编码器的输入。

3.1.1.4 VIT的Transformer编码器架构

Transformer编码器(Vaswani等人,2017)由交替堆叠的多头自注意力模块(MSA,见附录A)和多层感知机模块组成(公式2、3)。每个模块前都会先进行层归一化(LN)操作,每个模块后都会添加残差连接(Wang等人,2019;Baevski & Auli,2019)。

该多层感知机(MLP)包含两层网络,并使用GELU非线性激活函数。
z 0 = [ x class ; x p 1 E ; x p 2 E ; ⋯   ; x p N E ] + E p o s , E ∈ R ( P 2 ⋅ C ) × D , E p o s ∈ R ( N + 1 ) × D (1) \boldsymbol{z}0 = \left[\boldsymbol{x}{\text{class}}; \boldsymbol{x}p^1 \boldsymbol{E}; \boldsymbol{x}p^2 \boldsymbol{E}; \cdots; \boldsymbol{x}p^N \boldsymbol{E}\right] + \boldsymbol{E}{pos}, \quad \boldsymbol{E} \in \mathbb{R}^{(P^2 \cdot C) \times D}, \boldsymbol{E}{pos} \in \mathbb{R}^{(N+1) \times D} \tag{1} z0=[xclass;xp1E;xp2E;⋯;xpNE]+Epos,E∈R(P2⋅C)×D,Epos∈R(N+1)×D(1)
z ℓ ′ = MSA ( LN ( z ℓ − 1 ) ) + z ℓ − 1 , ℓ = 1 ... L (2) \boldsymbol{z}'
\ell = \text{MSA}(\text{LN}(\boldsymbol{z}{\ell-1})) + \boldsymbol{z}{\ell-1}, \quad \ell = 1 \dots L \tag{2} zℓ′=MSA(LN(zℓ−1))+zℓ−1,ℓ=1...L(2)
z ℓ = MLP ( LN ( z ℓ ′ ) ) + z ℓ ′ , ℓ = 1 ... L (3) \boldsymbol{z}\ell = \text{MLP}(\text{LN}(\boldsymbol{z}'\ell)) + \boldsymbol{z}'_\ell, \quad \ell = 1 \dots L \tag{3} zℓ=MLP(LN(zℓ′))+zℓ′,ℓ=1...L(3)
y = LN ( z L 0 ) (4) \boldsymbol{y} = \text{LN}(\boldsymbol{z}_L^0) \tag{4} y=LN(zL0)(4)

3.1.2 Inductive bias 归纳偏置

我们注意到,与卷积神经网络(CNN)相比,视觉Transformer的图像专属归纳偏置要少得多。

图像专属归纳偏置是什么?

图像里的物体是局部连在一起的(局部性

物体在图里挪位置,样子不变(平移等变性

像素是二维网格排列的(二维邻域结构

在CNN中,局部性、二维邻域结构和平移等变性,是原生嵌入到模型的每一层中的。

CNN:天生自带「图像常识」
局部性 :卷积核只看一小块区域,天然认为 "附近的像素才有关系"
二维邻域 :卷积是在 2D网格上滑窗,天生知道图像是平面结构
平移等变性:物体挪到别的位置,卷积还是能认出它,不会因为位置变了就认不出来

而在ViT中,只有多层感知机(MLP)层具备局部性和平移等变性,自注意力层则是全局的。

ViT:几乎没自带「图像常识」,全靠学
切分成图像块(Patch) :把 224×224 切成 16×16的小块,这一步勉强保留了一点局部邻域关系
微调时调整位置嵌入:换更高分辨率时,对位置嵌入做 2D 插值,这是手动补一点二维信息

模型对二维邻域结构的使用非常有限:仅在模型输入端将图像切分为图像块时,以及微调阶段为不同分辨率的图像调整位置嵌入时(如下文所述)会用到。除此之外,初始化阶段的位置嵌入不携带任何关于图像块二维位置的信息,图像块之间的所有空间关系都需要模型从头学习。

自注意力层是全局的 :它会把所有 Patch 都看一遍,不管离得近还是远,完全没有 "附近更相关" 的先验
初始化的位置嵌入不带任何 2D 信息 :模型一开始根本不知道哪个 Patch 在左上、哪个在右下

所有空间关系(比如 "猫的耳朵在头的上方"),都要模型自己从数据里一点点学出来

3.1.3 Hybrid Architecture 混合架构

混合架构 (Hybrid Architecture)是纯 ViT 的一种扩展方案 ,它不再直接用原始图像块作为 Transformer输入,而是:

先用 CNN 提取图像的特征图 ,再把特征图处理成序列,喂给 Transformer。

作为原始图像块的替代方案,模型的输入序列也可以由CNN的特征图构建而来(LeCun等人,1989)。

在该混合模型中,我们将公式1中的图像块嵌入投影矩阵 E \boldsymbol{E} E,作用于从CNN特征图中提取出的图像块。

1×1 Patch = 不需要 "切 Patch",直接把特征图的每个空间位置当作一个元素

作为一种特殊情况,图像块的空间尺寸可以设为1×1,这意味着只需将特征图的空间维度展平,再投影到Transformer的维度,即可得到输入序列。
BERT & PE

分类输入嵌入(class token)和位置嵌入(position token)的添加方式,与前文所述一致。

3.2 FINE-TUNING AND HIGHER RESOLUTION 微调与高分辨率适配

3.2.1 ViT 的预训练 + 微调两阶段范式

预训练和微调:分类任务

下游任务的类别数和预训练任务完全不匹配

所以需要把原先的 D ∗ K p r e t r a i n D * K_{pretrain} D∗Kpretrain

替换为 D ∗ K D * K D∗K前馈层

然后零初始化=新线性层的权重、偏置全部初始化为0

通常,我们在大规模数据集上对ViT进行预训练,再针对(规模更小的)下游任务进行微调。

为此,我们会移除预训练阶段的预测头,替换为一个零初始化的 D × K D \times K D×K前馈层,其中 K K K为下游任务的类别数。

3.2.2 预训练和微调的图像分辨率不一致

在比预训练更高的分辨率下进行微调,通常能带来性能增益(Touvron等人,2019;Kolesnikov等人,2020)。

Transformer 的序列长度维度是"动态"的 ,和模型权重无关

Transformer 本质上是一个 对序列长度"无感知" 的架构,只要每个元素的嵌入维度 D 不变,不管序列长度 N 是多少,都能直接跑。

一、为什么能处理任意长度的序列?

权重形状计算逻辑两方面看:

  1. 权重和序列长度无关
    Transformer 里所有可学习参数的形状,都不依赖输入序列长度 N
      - 线性投影层(Q/K/V 投影、MLP):权重形状是 [D, D],只和嵌入维度 D 有关
      - 层归一化(LayerNorm):只对 D 维做归一化,和序列长度无关
      - 位置嵌入:预训练时是 [N_pretrain+1, D],但如果新序列长度 N_new 变了,只要对位置嵌入做插值/扩展即可,模型结构本身不用改
    模型权重只关心"每个元素是几维",不关心"有多少个元素"
  2. 自注意力计算是"逐元素全局配对"
    自注意力的核心是计算:
    Attention ( Q , K , V ) = softmax ( Q K ⊤ D ) V \text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{D}}\right)V Attention(Q,K,V)=softmax(D QK⊤)V
      - Q , K , V Q,K,V Q,K,V 的形状是 [batch_size, N, D]
      - Q K ⊤ QK^\top QK⊤ 得到 [batch_size, N, N] 的注意力分数矩阵:不管 N 是 196 还是 576,矩阵乘法逻辑完全一样,只是矩阵大小变了
      - 只要显存放得下这个 N×N 的注意力矩阵,就能计算
    序列长度变多,只是让矩阵变大了,计算流程完全不变

二、「图像块变多了仍然可以用原先的 Transformer」是什么意思?

以 ViT 为例:

- 预训练时 :输入是 224×224 图像 → 切 16×16 Patch → N=196 个 Patch → 序列长度 196+1=197

- 微调高分辨率时 :输入是 384×384 图像 → 同样切 16×16 Patch → N=576 个 Patch → 序列长度 576+1=577

  1. Transformer 编码器的结构、权重完全不用改 :所有层的权重形状还是 [D,D],和序列长度无关
  2. 只需要处理位置嵌入 :预训练的位置嵌入是 [197, D],现在序列长度变成 577,只要对它做 2D 插值,扩展成 [577, D],就能继续用
  3. 前向/反向传播流程完全不变,只是 N×N 注意力矩阵从 197×197 变成了 577×577,显存占用变大,但模型结构本身兼容
    只要显存够,序列长度随便改,模型结构不用动。
    理论上可以处理任意长度,但实际受限于 GPU 显存能放下多大的 N×N 矩阵

Transformer 只要求「每个元素维度固定」,不限制「元素个数」,所以图像块变多(序列变长)时,模型结构不用改,只要显存够,就能直接跑。

当输入更高分辨率的图像时,我们保持图像块的尺寸不变,这会使最终的有效输入序列长度变长。视觉Transformer可以处理任意长度的序列(仅受显存上限限制),但预训练得到的位置嵌入可能不再具备对应位置的语义。因此,我们根据预训练位置嵌入在原始图像中的对应位置,对其执行二维插值操作。

二维插值

①:位置嵌入 pos_emb 形状是 (N+1, D):

第 0 个元素 :class_token_pos_emb → 形状 (1,768)(保留不动,直接复用)

后面 N 个元素 :patch_pos_emb → 形状 (196, 768)(需要插值的部分)

②:预训练的196 个 Patch 对应 14×14 的 2D 网格,所以把 (196, 768) reshape 成 (14, 14, 768)

③:把14×14×768 的 2D 网格,插值到目标分辨率 24×24×768:

需要注意的是,这种分辨率调整与图像块提取操作,是我们向视觉Transformer中手动注入图像二维结构归纳偏置的唯二环节。

4 EXPERIMENTS 实验

我们对ResNet、视觉Transformer(ViT)以及混合架构的表征学习能力进行了评估。

为明确各模型的数据需求,我们在不同规模的数据集上完成预训练,并在多个基准任务上开展性能测评。

在考量模型预训练的计算开销时,ViT展现出了极具优势的性能表现:它以更低的预训练成本,在绝大多数图像识别基准任务上达到了当前最优水平。

最后,我们通过自监督学习完成了一组小型实验,结果表明,自监督ViT具备良好的发展前景。

4.1 SETUP 实验设置

4.1.1 Datasets 数据集

为探究模型的可扩展性,我们使用了包含1000个类别、130万张图像的ILSVRC-2012 ImageNet数据集(下文简称ImageNet)、它的超集ImageNet-21k(包含21000个类别、1400万张图像,Deng等人,2009),以及JFT数据集(Sun等人,2017,包含18000个类别、3.03亿张高分辨率图像)。

我们参照Kolesnikov等人(2020)的方法,针对下游任务的测试集,对预训练数据集进行了去重处理。

为我们将在这些数据集上训练完成的模型,迁移到多个基准任务上开展测评,包括:使用原始验证标签和优化后ReaL标签的ImageNet任务(Beyer等人,2020)、CIFAR-10/100数据集(Krizhevsky,2009)、牛津-IIIT宠物数据集(Parkhi等人,2012),以及牛津花卉-102数据集(Nilsback & Zisserman,2008)。针对这些数据集,我们均参照Kolesnikov等人(2020)的方法完成数据预处理。

VTAB任务

VTAB 全称是 Visual Task Adaptation Benchmark(视觉任务适配基准),是 Google 团队在 2019 年提出的标准化视觉测评基准,用来评估模型小样本迁移能力 、跨场景泛化能力的核心测试套件。
VTAB 最关键的规则 :所有 19 项任务,每个任务仅允许使用最多 1000 个标注训练样本,不允许使用额外数据、额外预训练,所有模型使用完全统一的微调设置。
评估 :模型在极低数据量 下的迁移适配能力

自然图像类 :通用自然图像分类任务,CIFAR-10/100等

专业领域类 :专业领域图像任务,医学影像、卫星遥感图像等

结构化类:要求模型理解图像的几何结构、位置关系的任务,比如目标定位、关键点检测、场景深度估计等

我们还在包含19项任务的VTAB分类基准套件上完成了模型测评(Zhai等人,2019b)。VTAB基准用于评估模型在多样化任务上的小数据迁移能力,每个任务仅使用1000个训练样本。该基准的任务分为三大类:自然图像类------包含上述宠物分类、CIFAR等常规视觉任务;专业领域类------包含医学影像、卫星遥感影像相关任务;结构化类下现在XIANZAI XN12112------包含目标定位等需要几何理解能力的任务。

4.1.2 Model Variants 模型变体

我们参考BERT(Devlin等人,2019)的配置设计了ViT的不同模型规格,详见表1。

三个变体:Base,Large,Huge

其中Vit-B/16 这里的16是指图像块尺寸为16*16

其中"Base(基础版)"和"Large(大型版)"模型直接沿用BERT的配置,同时我们新增了规模更大的"Huge(巨型版)"模型。下文我们将使用简写符号标注模型规模和输入图像块尺寸:例如,ViT-L/16代表输入图像块尺寸为16×16的"Large"模型变体。

需要注意的是,Transformer的输入序列长度与图像块尺寸的平方成反比,因此图像块尺寸越小的模型,计算开销越高。

批归一化vs组归一化

1.BN(Batch Normalization 批归一化)

沿着批次维度,对每个通道单独做归一化。对每一个特征通道,把整个批次里所有样本、该通道的所有像素值拉到一起,计算均值和方差,再做归一化。

优点 :大批次训练,统计的均值/方差非常准确。

缺点 :极度依赖batchsize,如果batchsize很小(比如≤8,甚至1),模型性能会暴跌。

迁移学习缺陷 :训练和推理行为不一致。

2.GN(Group Normalization 组归一化)

完全不看批次里的其他图像,只针对当前单张图,把它的通道分成几个组,每个组内的所有像素值算均值和方差,再做归一化。

优点 :与batchsize无关。

迁移学习优势 :训练和推理的行为完全一致。

小缺点:在超大batch训练时,效果略逊于BN(差距极小)。

作为基线的卷积神经网络,我们采用ResNet(He等人,2016)架构,但将其中的批归一化(Batch Normalization)层(Ioffe & Szegedy, 2015)替换为组归一化(Group Normalization)层(Wu & He, 2018),并使用了标准化卷积(standardized convolutions,Qiao等人,2019)。这些修改能够提升模型的迁移学习性能(Kolesnikov等人,2020),我们将修改后的模型命名为"ResNet (BiT)"。

对于混合架构,我们将CNN的中间特征图输入ViT,此时ViT的图像块尺寸为1个"像素"。

不同序列长度下开展实验
混合架构 里,Transformer 的输入序列长度 = CNN 最终输出特征图的「高度 × 宽度」(因为用了 1×1 的 Patch,特征图里每个空间位置就是一个 token)。
方案(i) : 取常规ResNet50 第4阶段的输出 。用完整的 ResNet50,最终输出 stage4 的特征图,尺寸是 7×7 ,得到49个token的短序列
方案(ii) :移除第 4 阶段,把对应层补到第 3 阶段,取扩展后的 stage3 输出 。移除stage4,补层到stage3,所以现在stage3一共6+3=9层。原stage3的特征图为14*14,不变。所以一共14×14=196个token序列,刚好是方案(i)的四倍。

为了在不同序列长度下开展实验,我们采用两种方案:(i)取常规ResNet50 第4阶段 的输出;(ii)移除ResNet50的第4阶段,将对应数量的层补充到第3阶段中(保持总层数不变),取扩展后的第3阶段的输出。方案(ii)会使输入序列长度变为原来的4倍,对应的ViT模型计算开销也更高。

4.1.3 Training & Fine-tuning 训练与微调

预训练

我们使用Adam优化器(Kingma & Ba, 2015)训练所有模型(包括ResNet),其中β₁=0.9,β₂=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分辨率,同时使用了系数为0.9999的Polyak & Juditsky (1992) 权重平均策略(Ramachandran等人,2019;Wang等人,2020b)。

4.1.4 Metrics 评价指标

我们通过少样本分类准确率或微调分类准确率,汇报模型在下游数据集上的表现。其中,微调准确率反映了各模型在对应下游数据集上完成微调后的性能;少样本准确率则通过求解带正则化的最小二乘回归问题得到,该回归任务将部分训练图像的(冻结的)表征映射到 { − 1 , 1 } K \{-1, 1\}^K {−1,1}K的目标向量中。该公式可让我们通过闭式解得到精确结果。尽管我们主要聚焦于模型的微调性能,但在微调计算开销过高的场景下,我们也会使用线性少样本准确率完成快速的即时测评。

4.2 COMPARISON TO STATE OF THE ART 与当前最优方法的对比

对比实验设置

我们首先将规模最大的模型------ViT-H/14与ViT-L/16,与现有文献中的当前最优CNN模型进行对比。

第一个对比基线是Big Transfer(BiT)(Kolesnikov等人,2020),该方法基于大型ResNet开展有监督迁移学习。第二个对比基线是Noisy Student(Xie等人,2020),这是一个大型EfficientNet模型,在ImageNet和去除标签的JFT-300M数据集上通过半监督学习训练得到。目前,Noisy Student是ImageNet数据集上的当前最优方法,而BiT-L是本文其余汇报数据集上的当前最优方法。

所有模型均在TPUv3硬件上完成训练,我们同时汇报了各模型预训练所需的TPUv3核·天数量,即训练所用的TPUv3核心数(每块芯片含2个核心)与训练天数(以天为单位)的乘积。
对比实验结果

表2展示了实验结果。在JFT-300M上预训练的、规模更小的ViT-L/16模型,在所有任务上的表现均优于在相同数据集上预训练的BiT-L模型,同时训练所需的计算资源大幅降低。规模更大的ViT-H/14模型实现了性能的进一步提升,尤其在ImageNet、CIFAR-100和VTAB基准套件这类更具挑战性的数据集上表现突出。值得关注的是,该模型的预训练计算量仍显著低于此前的当前最优方法。

不过我们注意到,预训练效率不仅受架构选择的影响,还会受训练策略、优化器、权重衰减等其他参数的影响。我们在4.4节中针对不同架构,开展了性能与计算量的对照实验。

最后,在公开数据集ImageNet-21k上预训练的ViT-L/16模型,在大多数数据集上也取得了优异的表现,同时预训练所需资源更少:使用标准的8核云TPUv3,仅需约30天即可完成训练。
VTAB基准任务上的结果

图2将VTAB基准的任务按对应类别拆分,并与该基准上此前的SOTA方法进行了对比,包括:BiT、VIVI(在ImageNet和Youtube数据集上联合训练的ResNet模型,Tschannen等人,2020),以及S4L(在ImageNet上开展有监督+半监督学习的方法,Zhai等人,2019a)。ViT-H/14在自然图像类和结构化类任务上,性能超越了BiT-R152x4及其他方法;在专业领域类任务上,排名前二的模型性能相近。

4.3 PRE-TRAINING DATA REQUIREMENTS 预训练的数据需求

视觉Transformer在大规模JFT-300M数据集上完成预训练时表现优异。相较于ResNet,视觉Transformer针对视觉任务的归纳偏置更少,那么数据集规模对其而言究竟有多关键?我们开展了两组实验进行验证。

首先,我们在规模递增的数据集上对ViT模型进行预训练,分别为ImageNet、ImageNet-21k和JFT-300M。为提升模型在小规模数据集上的表现,我们对三个基础正则化参数进行了优化,分别是权重衰减、随机失活(dropout)和标签平滑。图3展示了在ImageNet上微调后的结果(其他数据集的结果见表5)²。

小规模数据集上 :大模型性能表现不及基准模型
中规模数据集上 :模型性能表现相近
大规模数据集上:大模型全部性能优势

在规模最小的ImageNet数据集上预训练时,尽管采用了(适度的)正则化策略,ViT-Large模型的表现仍不及ViT-Base模型;在ImageNet-21k上完成预训练后,两者的表现相近;只有在JFT-300M数据集上预训练时,我们才能看到大模型的全部性能优势。图3同时展示了不同规模BiT模型的性能区间:基于CNN的BiT模型在ImageNet上的表现优于ViT,但随着数据集规模扩大,ViT实现了反超。
同一数据集的不同规模子集

其次,我们分别在900万、3000万、9000万规模的JFT-300M随机子集,以及完整的JFT-300M数据集上训练模型。

我们未在小规模子集上额外添加正则化,所有实验设置均使用相同的超参数。通过这种方式,我们评估的是模型的固有属性,而非正则化带来的影响。不过我们采用了早停策略,并汇报训练过程中取得的最优验证准确率。为节省计算资源,我们汇报的是少样本线性分类准确率,而非完整微调后的准确率,实验结果见图4。

少样本线性分类准确率

线性探测(Linear Probing)+ 少样本设定

完全冻结预训练的模型主干 :Transformer 的所有权重一丝不动,不做任何修改,只保留它预训练学到的特征提取能力;

少样本设定 :下游任务只取极少量的标注样本(比如每个类别 1/5/10 张图),不用完整训练集;

只学一个线性分类头 :把少样本图片输入冻结的主干,提取特征后,只训练一个最简单的线性层(分类头),而且这个线性层不用梯度下降迭代,直接用「带正则化的最小二乘闭式解」一步算出最优权重;

最后在测试集上算出的准确率,就是少样本线性分类准确率。

在小规模数据集上,计算开销相近的情况下,视觉Transformer比ResNet 更容易过拟合 。例如,ViT-B/32的推理速度略快于ResNet50,它在900万样本的子集上表现远差于ResNet50,但在9000万及以上样本的子集上表现更优;ResNet152x2与ViT-L/16的对比也符合这一规律。该结果进一步印证了我们的直观判断:卷积带来的归纳偏置在小规模数据集上十分有效,但对于大规模数据集,直接从数据中学习相关模式就已足够,甚至能带来更优的效果。

总体而言,ImageNet上的少样本实验结果(图4),以及VTAB基准上的小数据实验结果(表2)表明,ViT在极低数据量的迁移学习场景中具备良好的潜力。对ViT少样本特性的进一步分析,是未来极具价值的研究方向。

4.4 SCALING STUDY 缩放性研究

我们通过评估模型在JFT-300M数据集上的迁移性能,针对不同模型开展了对照缩放性研究。在该实验设置下,数据规模不会成为模型性能的瓶颈,我们以此评估各模型的性能与预训练计算开销的对应关系。

实验的模型集合包括:7个ResNet模型,分别为预训练7个轮次的R50x1、R50x2、R101x1、R152x1、R152x2,以及预训练14个轮次的R152x2和R200x3;6个视觉Transformer模型,分别为预训练7个轮次的ViT-B/32、B/16、L/32、L/16,以及预训练14个轮次的L/16和H/14;5个混合架构模型,分别为预训练7个轮次的R50+ViT-B/32、B/16、L/32、L/16,以及预训练14个轮次的R50+ViT-L/16(对于混合架构,模型名称末尾的数字不代表图像块尺寸,而是ResNet骨干网络的总下采样率)。

图5展示了迁移性能与总预训练计算量的对应关系(计算开销的详细信息见附录D.5),各模型的详细结果见附录表6。我们可以观察到几个规律:第一,在性能-计算量的权衡上,视觉Transformer全面优于ResNet。为达到相同的性能(5个数据集的平均结果),ViT所需的计算量仅为ResNet的1/2~1/4。第二,在计算预算较小时,混合架构的表现略优于ViT,但随着模型规模扩大,这一差距便会消失。这一结果有些出乎意料,因为人们通常认为,无论模型规模大小,卷积的局部特征处理都能为ViT提供辅助。第三,在本次实验的测试范围内,视觉Transformer的性能尚未出现饱和,这为后续的模型缩放研究提供了方向。

4.5 INSPECTING VISION TRANSFORMER 视觉Transformer的机制探查

为理解视觉Transformer对图像数据的处理方式,我们对其内部表征进行了分析。视觉Transformer的第一层会将展平后的图像块线性投影到低维空间中(公式1)。图7(左)展示了学习到的嵌入滤波器的核心主成分,这些主成分与用于表征每个图像块内精细结构的低维基函数高度契合。

完成投影后,我们将可学习的位置嵌入添加到图像块表征中。图7(中)表明,模型通过位置嵌入学习到了对图像内的空间距离进行编码,即距离越近的图像块,其位置嵌入的相似度越高。此外,模型还学习到了行列结构:同一行/列的图像块具有相似的嵌入。最后,对于更大的网格,有时还能观察到正弦式的结构(附录D)。位置嵌入能够自主学习表征二维图像的拓扑结构,这也解释了为什么人工设计的二维感知嵌入变体无法带来性能提升(附录D.4)。

注意力距离 ,就是这个 token 在整合信息时,平均跨越了图像上多大的空间范围。

一个左上角的 Patch,注意力权重几乎全给了相邻的几个Patch ,那它的注意力距离就很小,只关注图像局部;

如果它的权重均匀分给了全图所有Patch ,那注意力距离就很大,能直接整合整幅图的全局信息。

具体怎么怎么算:

①先确定每个 Patch 在图像 2D 网格里的空间坐标,计算两个 Patch 之间的欧氏距离 (空间距离);

②用注意力权重 给所有 Patch 的空间距离做加权平均 ,就得到了这个 token 的注意力距离;

③一层 Transformer 所有 token 的平均注意力距离,就是这一层的注意力距离,完全对应 CNN 里「一层的感受野大小」。

自注意力机制使ViT 即使在网络的最底层,也能整合整幅图像的信息。我们探究了网络对这一能力的利用程度。具体而言,我们基于注意力权重,计算了信息整合所跨越的图像空间平均距离(图7,右),这一"注意力距离"与CNN中的感受野大小概念类似。

CNN的早期卷积层

感受野小(3*3),只做局部信息交互,提取图像里最底层的基础视觉特征。
Vit的小注意力头

只会给和自己相邻的Patch Token分配高注意力权重,计算范围局限在图像的极小局部里。
ViT 里小注意力距离的注意力头,和 CNN早期卷积层,干的是完全一样的活 ------ 只在图像的极小局部范围内做信息交互,提取边缘、纹理、角点这类基础视觉特征,它起到了 CNN早期卷积层的作用。

我们发现,部分注意力头在网络最底层就已经能够关注到图像的大部分区域,这表明模型确实利用了全局信息整合的能力。另一些注意力头在低层网络中始终保持很小的注意力距离。这种高度局部化的注意力,在Transformer前接入ResNet的混合架构模型中表现得更弱(图7,右),这说明它可能起到了与CNN中早期卷积层类似的作用。

此外,注意力距离会随着网络深度的增加而增大。整体来看,我们发现模型会关注对分类任务具有语义相关性的图像区域(图6)。

4.6 SELF-SUPERVISION 自监督学习

大规模的自监督预训练 vs 有监督预训练

ViT 这里模仿的是 BERT 的掩码任务
BERT 怎么做(NLP)

  • 遮住一句话里的几个词
  • 让模型猜被遮住的词是什么

ViT怎么做(视觉)

  • 把图片切成很多 patch(token)
  • 随机遮住一部分 patch
  • 让模型 猜被遮住的 patch 长什么样

Transformer在自然语言处理(NLP)任务中展现出了优异的性能。然而,其成功不仅源于出色的可扩展性,还得益于大规模的自监督预训练(Devlin等人,2019;Radford等人,2018)。

我们模仿BERT中使用的掩码语言建模任务,针对自监督学习的掩码图像块预测任务开展了初步探索。

经过自监督预训练后,我们规模较小的ViT-B/16模型在ImageNet上达到了79.9%的准确率,相较于从头训练的结果显著提升了2个百分点,但仍比有监督预训练的结果低4个百分点。

更多细节详见附录B.1.2。我们将对比学习预训练的相关探索(Chen等人,2020b;He等人,2020;Bachman等人,2019;Hénaff等人,2020)留作未来工作。

5 CONCLUSION 结论

我们探索了Transformer直接应用于图像识别的可行路径。与此前在计算机视觉领域使用自注意力的研究不同,除了初始的图像块提取步骤外,我们并未在架构中引入任何图像特定的归纳偏置。相反,我们将图像视为图像块序列,并采用自然语言处理中的标准Transformer编码器对其进行处理。这种简单、可扩展且具备策略性的方法,在大规模数据集上进行预训练后,表现出了惊人的良好效果。因此,视觉Transformer(ViT)在众多图像分类任务上达到甚至超越了当前最优水平,且预训练过程相对廉价。

尽管这些初始结果令人鼓舞,但仍有诸多挑战亟待解决。其一,需将ViT成功应用于检测、分割等其他计算机视觉任务,我们的初步实验结果与Caron等人(2020)的研究均表明了这一方法的潜力。其二,需继续探索自监督预训练方法,目前自监督预训练与大规模监督预训练之间仍存在显著的性能差距,而对ViT的进一步优化或许能缩小这一差距。最后,提升ViT的缩放能力有望进一步改善模型性能。
致谢

本研究在柏林、苏黎世和阿姆斯特丹完成。我们感谢谷歌的众多同事提供的帮助,尤其感谢Andreas Steiner在基础设施方面的关键支持,以及Juho Puigcerver和Maxim Neumann在大规模训练基础设施上提供的协助;同时感谢Dmitry Lepikhin、Aravind Mahendran、Daniel Keysers、Mario Lučić、Noam Shazeer、Ashish Vaswani和Colin Raffel提供的有益讨论。

相关推荐
码农的神经元2 小时前
Deep-HMM 融合 Transformer:序列分类的动态隐状态建模新范式
人工智能·深度学习·transformer
movigo7_dou3 小时前
SIFT的一些内容
论文阅读·图像处理·学习·计算机视觉
格林威3 小时前
Baumer相机金属冲压件毛刺高度测量:量化去毛刺效果的 5 个核心方法,附 OpenCV+Halcon 实战代码!
人工智能·opencv·计算机视觉·c#·视觉检测·机器视觉·工业相机
码农的神经元3 小时前
Transformer-SSM 混合模型 vs 传统 SVM:文本分类任务的深度对比实验
支持向量机·分类·transformer
再不会python就不礼貌了11 小时前
从工具到个人助理——AI Agent的原理、演进与安全风险
人工智能·安全·ai·大模型·transformer·ai编程
春日见14 小时前
自驾算法的日常工作?如何提升模型性能?
linux·人工智能·机器学习·计算机视觉·自动驾驶
我材不敲代码16 小时前
OpenCV实战:全自动答题卡识别与评分系统
人工智能·opencv·计算机视觉
霖大侠17 小时前
Wavelet Meets Adam: Compressing Gradients forMemory-Efficient Training
人工智能·深度学习·算法·机器学习·transformer
不熬夜的熬润之19 小时前
APCE-平均峰值相关能量
人工智能·算法·计算机视觉