Vision Transformer (ViT)

视觉 Transformer (ViT) 是一种将语言建模中的 Transformer架构应用于图像的神经网络。它会将图像切分成固定大小的图像块，把每个图像块转换为一个词元嵌入，然后通过自注意力处理整个序列，因此从第一层开始，每个图像块都能利用其他所有图像块的信息。Google 研究人员在《An Image is Worth 16x16 Words》中提出了这一架构，并证明只要使用足够的数据进行预训练，不含任何卷积层的纯 Transformer 就能在图像分类任务上达到最先进的卷积神经网络 (CNNs)水平。如今，ViT 还被用作视觉语言模型、分割基础模型和自监督主干网络的图像编码器。

视觉 Transformer 的工作原理#

ViT 会先将图像转换为序列，再像处理文本一样处理该序列。以下步骤以 224×224 图像上的 ViT-Base/16 为例：

1. 图像块嵌入：图像会被切分成 16×16 像素的图像块，形成一个 14×14 网格，共 196 个图像块。每个图像块会被展平并乘以一个可学习矩阵，转换成 768 维的嵌入向量；通常用步幅为 16 的 16×16 卷积实现这一过程。
2. 类别词元和位置嵌入：序列开头会添加一个可学习的类别词元，并为每个词元添加可学习的位置嵌入。注意力本身会忽略顺序，因此这些嵌入是模型了解每个图像块来源位置的唯一信息。
3. Transformer 编码器：一组相同的模块会交替执行多头自注意力和两层 MLP；ViT-Base 包含 12 个模块。多头自注意力会混合所有词元的信息，而使用 GELU 激活函数的 MLP 则会单独转换每个词元。每个子层之前都会执行层归一化，外围则有残差连接。
4. 分类头：类别词元的最终输出会输入线性分类器。分割等密集预测任务则使用图像块词元，并将其重塑为 14×14 网格。

下面的 PyTorch 代码会为一张图像构建词元序列，并将其传入单个编码器模块。

import torch
from torch import nn

# ViT-Base/16 设置：16x16 图像块、768 维词元、12 个注意力头
image = torch.randn(1, 3, 224, 224)
patch_embed = nn.Conv2d(3, 768, kernel_size=16, stride=16)  # 每个图像块进行一次线性投影

# 将图像切分为 14x14 网格的图像块，并展平为 196 个词元
tokens = patch_embed(image).flatten(2).transpose(1, 2)

# 在序列前添加可学习的类别词元，并添加可学习的位置嵌入
cls_token = nn.Parameter(torch.zeros(1, 1, 768))
pos_embed = nn.Parameter(torch.zeros(1, 197, 768))
tokens = torch.cat([cls_token, tokens], dim=1) + pos_embed

# 一个前置归一化编码器模块，其中每个词元都会关注其他所有词元
block = nn.TransformerEncoderLayer(
    d_model=768, nhead=12, dim_feedforward=3072, activation="gelu", batch_first=True, norm_first=True
)
print(block(tokens).shape)

它会输出 torch.Size([1, 197, 768])：196 个图像块词元加上类别词元，每个词元仍为 768 维，可输入 ViT-Base 的下一个模块（共 12 个模块）。

模型规模与变体#

ViT 的名称包含模型规模和图像块尺寸：ViT-B/16 是拥有 8,600 万个参数、使用 16×16 图像块的 Base 模型；Large 和 Huge 模型分别拥有 3.07 亿和 6.32 亿个参数。更小的图像块会产生更多词元，从而提升细节表现，但计算开销更高。一些变体调整了 ViT 的训练方式或结构：

• DeiT：一种采用强力数据增强和蒸馏词元的训练方案，仅使用 ImageNet-1k 就能训练出具有竞争力的 ViT，无需额外的预训练数据。
• Swin Transformer：在移动后的局部窗口内计算注意力，并在不同阶段之间合并图像块，形成类似 CNN 的层次结构，适用于检测和分割任务。
• 掩码自编码器 (MAE) 和 DINOv2：这两种自监督方法会使用未标注图像预训练 ViT：MAE 通过重建被遮蔽的图像块进行训练，DINOv2 则通过匹配教师网络的输出进行训练。
• CLIP 和 SigLIP：通过训练 ViT 图像编码器匹配图像及其说明文字，广泛用作多模态模型的视觉主干。

应用场景#

• 视觉语言模型：ViT 编码器会将图像转换为一系列词元，由大型语言模型与文本一起读取，从而实现图像描述和视觉问答。
• 密集预测：ViT 主干网络生成的图像块词元可用于图像分割和深度估计；这类任务需要让每个输出像素都能利用图像远处的上下文信息。
• 视频理解：视频 ViT 会将片段切分为空间-时间图像块，而 V-JEPA 等自监督视频模型则从未标注视频中学习运动和外观特征，用于动作识别等任务。

视觉 Transformer 与 CNN 的比较#

CNN 在每一层中都内置了两种假设：局部性，因为每个滤波器只观察一个小邻域；平移等变性，因为每个位置都使用相同的滤波器。除图像块网格外，ViT 必须从数据中学习空间结构。在原始论文中，仅使用 ImageNet（130 万张图像）训练的 ViT，得分比同等规模的 ResNet低几个百分点；而在包含 3 亿张图像的 JFT-300M 数据集上进行预训练后，ViT-H/14 的 ImageNet top-1 准确率达到 88.55%。另一个区别是计算成本。自注意力会比较每个词元与其他所有词元，因此图像边长翻倍会使词元数量变为 4 倍，注意力计算成本则变为 16 倍。Flash attention可减少这项计算的内存流量，而窗口注意力则会减少实际比较次数。Ultralytics 可运行 Meta 的分割一切模型，其 sam_b.pt 检查点使用 ViT-B 主干网络对每张 1024×1024 图像进行编码：16×16 图像块构成 64×64 网格，共 4,096 个词元；因此其 12 个模块中有 8 个仅在 14×14 窗口内执行注意力，另有 4 个使用全局注意力。