AI 棱镜站

主题

Vision Transformer 架构演进 ​

Vision Transformer（ViT）是计算机视觉领域的里程碑式工作，首次证明了纯 Transformer 架构在图像任务上可以超越传统 CNN，开启了视觉模型架构的新范式。

Vision Transformer架构示意

ViT 核心思想 ​

ViT 的核心创新在于将 NLP 领域的 Transformer 直接应用于图像：

  1. 图像分块（Patch Embedding）：将图像切分为固定大小的 patch（如 16×16），每个 patch 线性投影为向量
  2. 位置编码（Position Embedding）：为每个 patch 添加可学习的位置信息
  3. 分类 Token（[CLS] Token）： prepend 一个特殊 token 用于全局分类
  4. Transformer Encoder：标准的多头自注意力 + FFN 结构
python
import torch
import torch.nn as nn

class PatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        super().__init__()
        self.num_patches = (img_size // patch_size) ** 2
        self.projection = nn.Conv2d(
            in_channels, embed_dim,
            kernel_size=patch_size, stride=patch_size
        )

    def forward(self, x):
        # x: [B, C, H, W] -> [B, embed_dim, H/P, W/P]
        x = self.projection(x)
        # -> [B, num_patches, embed_dim]
        x = x.flatten(2).transpose(1, 2)
        return x

class VisionTransformer(nn.Module):
    def __init__(self, img_size=224, patch_size=16, embed_dim=768,
                 depth=12, num_heads=12):
        super().__init__()
        self.patch_embed = PatchEmbedding(img_size, patch_size, 3, embed_dim)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(
            torch.zeros(1, self.patch_embed.num_patches + 1, embed_dim)
        )
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(embed_dim, num_heads,
                                       dim_feedforward=embed_dim*4,
                                       dropout=0.1),
            num_layers=depth
        )
        self.head = nn.Linear(embed_dim, 1000)

    def forward(self, x):
        B = x.shape[0]
        x = self.patch_embed(x)
        cls_tokens = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls_tokens, x], dim=1)
        x = x + self.pos_embed
        x = self.encoder(x)
        return self.head(x[:, 0])

ViT 的关键发现 ​

ViT 论文的核心实验结论包括：

  • 数据规模至关重要：在小数据集（如 ImageNet-1K）上 ViT 表现不如 ResNet，但在大规模数据集（JFT-300M）预训练后显著超越 CNN
  • 模型规模与性能正相关：更大的模型（更多层、更高维度）在大数据上收益更大
  • 位置编码可学习：可学习的位置编码比正弦编码效果更好
  • 全局注意力优于局部：从第一层开始就使用全局注意力，而非逐步扩大感受野

预训练策略

ViT 的成功很大程度上依赖于大规模预训练。原始论文使用 JFT-300M（3 亿图像）进行预训练，然后在目标数据集上微调。这一范式后来被 MAE 等自监督方法进一步改进。

ViT 的局限性 ​

原始 ViT 存在几个明显不足：

  • 计算复杂度：自注意力的 O(n²) 复杂度限制了高分辨率输入
  • 缺乏归纳偏置：不像 CNN 具有平移不变性和局部性，需要更多数据学习
  • 固定分辨率：位置编码固定，难以处理不同分辨率的输入
  • 小目标检测弱：全局注意力对细粒度局部特征不够敏感

后续演进方向 ​

ViT 的提出催生了大量后续工作，主要演进方向包括：

方向代表工作核心改进
高效注意力Swin Transformer移动窗口局部注意力
数据高效训练DeiT蒸馏策略减少数据需求
多尺度特征PVT / ViTDet金字塔结构适配下游任务
自监督预训练MAE掩码自编码器预训练
卷积融合CoAtNet卷积与注意力混合架构

相关资源 ​

最近更新