Skip to content

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

相关链接

0 Abstract

  • Vision Transformer（ViT）首次将纯 Transformer 架构直接应用于图像分类任务，无需依赖卷积神经网络（CNN）。
  • 通过将图像分割为固定大小的图像块（patches）并转换为序列输入，ViT 在大规模数据集（如 ImageNet-21k、JFT-300M）上进行预训练后，在下游多个图像识别基准（如 ImageNet、CIFAR-100、VTAB 等）上取得了与当前最优卷积网络相当甚至更优的性能，且训练所需的计算资源显著减少。

1 Introduction

  • Transformer 在自然语言处理（NLP）中已成为主流架构，但在计算机视觉中仍以 CNN 为主导。
  • 本文提出 ViT，直接将标准 Transformer 应用于图像块序列，无需引入视觉特有的归纳偏置（如平移等变性、局部性）。
  • 实验表明，在小规模数据集上 ViT 表现不如 CNN，但在大规模数据集预训练后，ViT 表现优异，说明大规模训练可以超越归纳偏置。
  • Transformer 在 NLP 中的成功应用（如 BERT、GPT）。
  • 视觉中自注意力的应用：局部注意力、稀疏注意力、与 CNN 结合等。
  • 图像 Transformer 的早期尝试（如 iGPT、Cordonnier 等人的工作），但多数未在大规模数据上验证纯 Transformer 的竞争力。
  • 大规模图像识别中的数据扩展研究（如 BiT、Noisy Student）。

3 Method

alt text

如图1所示，整个流程可分为四步：

  1. 分块：将输入图像分割成固定大小的、不重叠的图像块。
  2. 线性投影：将每个图像块展平，并通过一个可训练的线性层（全连接层）映射为一个固定维度的向量，称为“图像块嵌入”。
  3. 添加位置与分类信息：在得到的图像块嵌入序列中，添加一个可学习的“[class]分类标记”到序列开头，并为所有标记（包括分类标记）添加位置嵌入。
  4. Transformer编码：将处理后的序列输入到一个标准的Transformer编码器中。编码器输出的对应“[class]标记”的向量，即作为整个图像的表示，用于最终的分类预测。

alt text

3.1 Vision Transformer (ViT)

1. 从图像到序列

为了适配Transformer处理一维序列的架构，ViT对输入图像 \(\mathbf{x} \in \mathbb{R}^{H \times W \times C}\) 进行如下变换：

  • 将图像分割为 \(N = HW / P^2\) 个大小为 \(P \times P\) 的图像块。
  • 将每个图像块展平为一个维度为 \(P^2 \cdot C\) 的向量。
  • 通过一个可训练的线性投影矩阵 \(\mathbf{E} \in \mathbb{R}^{(P^2 \cdot C) \times D}\)，将所有图像块向量映射到统一的模型维度 \(D\)，得到图像块嵌入序列，即上图（1）式。

这样，一个二维图像就被转化为一个长度为 \(N\) 的一维嵌入向量序列，类似于NLP中的“词序列”。

2. 特殊标记与位置嵌入

  • [class]标记：借鉴BERT，在图像块嵌入序列前添加一个随机初始化的可学习向量 \(z_0^0=\mathbf{x}_{\mathrm{class}}\)。经过Transformer编码器处理后，这个标记对应的输出向量 \(\mathbf{z}_L^0\) 被用作整个图像的全局表示(上图式4)。无论是预训练还是微调，分类头都附加在这个向量上。
  • 位置嵌入：由于Transformer本身不具备感知序列顺序的能力，需要向图像块嵌入中添加位置嵌入 \(\mathbf{E}_{pos}\) 来注入空间位置信息。论文发现，简单的、可学习的一维位置嵌入（按图像块的光栅扫描顺序分配）已经足够有效，使用更复杂的二维感知位置嵌入并未带来显著性能提升。

3. Transformer编码器架构

ViT使用的编码器与原始Transformer几乎一致：

  • 由 \(L\) 个相同的层堆叠而成。
  • 每一层包含一个多头自注意力（MSA）模块和一个多层感知机（MLP）模块（式2、3）
  • 在每个模块之前应用层归一化（LN）。
  • 每个模块都使用残差连接。
  • MLP包含两个全连接层，使用GELU激活函数。

GELU函数

  • GELU是一种平滑、概率化的激活函数，它通过结合输入的高斯分布特性来决定激活强度，在Transformer架构中表现出优秀的训练稳定性和模型性能，已成为现代深度学习模型（尤其是基于注意力的模型）中非线性层的标准选择之一。
  • \(\text{GELU}(x)=x\cdot \Phi(x)\)，\(\Phi(x)\)是标准正泰分布的累计分布函数（CDF）

4. 归纳偏置的对比

CNN有如下归纳偏置（先验知识）：

  • 局部性：图像中有意义的特征（如边缘、纹理、角点）都存在于局部的小范围邻域内，远处的像素通常不直接相关
    • 结构体现：使用小尺寸卷积核，每次只看图像的小尺寸区域
  • 平移不变性：图像中物体的特征（如一只猫的耳朵），无论它出现在画面的左上角还是右下角，都应该被同一种特征检测器识别出来。
    • 结构体现：使用权重共享的卷积核。同一个卷积核会滑动扫描整张图像，无论到哪里都用相同的参数去检测相同的特征。

ViT具有极弱的图像特定归纳偏置。

  • 自注意力层是全局的，模型必须从头学习图像块之间的空间关系。
  • 位置嵌入在初始化时并不包含任何2D位置信息。
  • 图像结构知识仅在两个阶段被手动引入：
    1. 在模型开始时，通过将图像切割成网格状图像块。
    2. 在微调调整分辨率时，对位置嵌入进行2D插值。

5. 混合架构

  • 作为纯ViT的替代方案，输入序列也可以来自CNN的中间特征图，而非原始图像块。
  • 在这种混合模型中，CNN（如ResNet）首先对图像进行特征提取，然后将特征图切割成“块”并进行线性投影，再输入Transformer。当“块”大小为1x1时，相当于直接将CNN特征图的空间维度展平作为序列。

3.2 微调与更高分辨率

ViT的典型使用范式是大规模预训练 + 下游任务微调。

  1. 微调流程：

    • 移除预训练时使用的分类头（通常是MLP）。
    • 附加一个零初始化的、维度为 \(D \times K\) 的线性层作为新的分类头，其中 \(K\) 是下游任务的类别数。

  2. 高分辨率微调：

    • 论文发现，在比预训练时更高的分辨率下进行微调通常能提升性能（这也是CNN中的常见技巧）。
    • 保持图像块大小 \(P\) 不变，更高的分辨率会导致图像块数量 \(N\) 增加，即序列长度变长。
    • 核心挑战：预训练好的位置嵌入是基于原始序列长度 \(N\) 的，无法直接对应到新的、更长的序列。
    • 解决方案：根据图像块在原始2D图像中的坐标，对预训练的位置嵌入进行2D插值。这是ViT在处理不同分辨率图像时，唯一需要手动引入图像2D结构知识的步骤。

4 Experiments

  • 数据集：ImageNet、ImageNet-21k、JFT-300M 等。
  • 模型配置：ViT-Base、Large、Huge，以及不同图像块大小（如 16×16）。
  • 主要结论：
  • 在大规模数据预训练后，ViT 在多个基准上超越或比肩 ResNet，且训练计算成本更低。
  • ViT 对数据规模敏感，小数据下易过拟合，大数据下表现更优。
  • 随模型规模增大，ViT 性能持续提升，未出现饱和现象。

alt text

5 Conclusion

ViT 展示了纯 Transformer 在图像识别中的强大潜力，尤其在大规模预训练下性能显著。未来研究方向包括：

  • 将 ViT 扩展至检测、分割等视觉任务。
  • 探索更高效的自监督预训练方法。
  • 进一步扩展模型规模以提升性能。

附录要点

  • 自注意力机制：详细说明 MSA 的结构与计算方式。
  • 位置编码消融实验：不同位置编码方式对性能影响不大。
  • 计算效率分析：ViT 在内存效率上优于 ResNet，尤其适合大分辨率输入。
  • 注意力可视化：浅层部分头关注局部，深层头关注全局，注意力距离随深度增加。
  • 自监督预训练初步实验：掩码图像块预测任务可提升性能，但仍落后于有监督预训练。

Comments