ViT-5:面向 2020 年代中期的 Vision Transformer

  • 关联论文：2602.08071
  • 作者：spark
  • 更新：2026-07-09

一句话结论

ViT-5 在不破坏「Attention-FFN」骨架的前提下，把过去五年里 Language Model / 视觉骨干网络的成熟技巧（LayerScale、RMSNorm、2D RoPE、QK-Norm、Register Token 等）系统地搬回 Vision Transformer，得到一套「即插即用」的新一代纯 ViT 骨干——ImageNet-1k 分类 84.2% top-1，在可比算力下稳超 DeiT-III（83.8%），并把 SiT 扩散框架的 FID 从 2.06 拉到 1.84。

解决的真问题

Vision Transformer 自 2020 年提出以来，「骨骼」几乎没变；真正在演化的，是 Language Model 侧的 Attention-FFN 模块细节（Norm、激活、位置编码、缩放、可学习 token）。但大量视觉团队还在用 2020 年的 DeiT / 原始 ViT 配方，导致：

  1. 骨干陈旧：大量 CV 工程仍在跑 vanilla ViT，结果既拼不过 DeiT-III，也吃不下最新的多模态预训练配方。
  2. 迁移性差：原始 ViT 的特征图里有 high-norm artifact，导致密集预测（检测、分割、深度）和生成建模（DiT、SiT）任务时表现不稳。
  3. 缺乏模块化升级清单：社区不缺单点改进（ConvNeXt 把卷积请回来、CoAt 让 Attention 与卷积杂交），缺一份「不动骨骼、只换零件」的纯 ViT 现代化清单。

ViT-5 的目标就是把这些「散落」的单点改进整理成一份有数据支撑、能直接 drop-in的现代化配方。

核心方法

ViT-5 的设计哲学可以浓缩成一句话：保留 Attention + FFN 的主结构，只在五个组件上做替换 / 增加。论文按组件给出 ablation 与选择依据：

1. 归一化：LayerNorm → RMSNorm

# 原 ViT
x = self.ln_1(x);  x = x + self.attn(x)   # Pre-LN + LayerNorm

# ViT-5
x = self.rms_1(x); x = x + self.attn(x)   # RMSNorm

理由：RMSNorm 去掉了均值中心化，等价于把「方差归一 + 可学习缩放」简化成「可学习缩放」，参数与算力都更省，且在大量 LLM 实践中已经验证稳定。

2. 激活 / 缩放：保留 GELU + 引入 LayerScale

  • LayerScale：每个残差分支乘一个可学习的对角矩阵 diag(l_1) … diag(l_L)（初始化为 ε ≈ 1e-5 或 1e-6），让深层在训练初期不至于被激活炸掉。
  • 抛弃 SwiGLU：论文明确说「有意拒绝 SwiGLU」，因为 SwiGLU 的 gating 已经在多个 Transformer 中稳定提升指标，但与 LayerScale 叠加时会过度调控，Loss 早期出抖。
  • QK-Norm：对 Q 和 K 做 L2 归一后再做点积，缓解注意力 logits 爆炸，与 LayerScale 互补。

3. 位置编码：APE + 2D RoPE 联用

# 输入 token 序列
tokens = patch_embed(x)         # B, N, D
tokens = tokens + ape_tokens    # 绝对位置
tokens = apply_rope_2d(tokens, grid_h, grid_w)  # 2D 旋转载码
  • APE（Absolute Positional Embedding）：保留少量可学习的绝对位置 token，给 ViT-5 在分类任务上稳定的归纳偏置。
  • 2D RoPE：把「高度」「宽度」两个轴各自的 rotary 频率独立编码，让模型显式学到空间相对关系，显著改善下游生成 / 检测 / 分割需要的位置感。

4. 可学习 token：Cls Token 之外追加 Register Token

Register tokens 是 Darcet et al. (2024) 的思路（论文对这一灵感给出引用）：在 patch tokens 之外再加一组「register tokens」，提供 high-norm artifact 的「倾倒区」。ViT-5 把 register token 接在输入序列后面，减少主特征图噪声，提升下游 dense 任务和迁移能力。

5. 消除 bias

qkv projection、FFN 的 linear 层都关掉 bias 项。这是 LLM 侧已经普遍遵循的简化，ViT-5 实验表明去掉 bias 后训练更稳定、推理略快，且不损失精度。

整体伪代码

class ViT5Block(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.rms1 = RMSNorm(dim)
        self.attn = Attention(dim, qkv_bias=False, qk_norm=True)
        self.ls1  = LayerScale(dim, init=1e-5)
        self.rms2 = RMSNorm(dim)
        self.mlp  = MLP(dim, hidden=4*dim, act=GELU, bias=False)
        self.ls2  = LayerScale(dim, init=1e-5)

    def forward(self, x):
        x = x + self.ls1(self.attn(self.rms1(x)))   # QK-Norm inside attn
        x = x + self.ls2(self.mlp(self.rms2(x)))
        return x

Stem 用普通 patch embed；外层用 tokens + APE + 2D RoPE + [register]*k，下游仍然是 cls_token 分类 head 或 patch_token 重建 / DiT conditioning。

关键实验与数据

论文给出的是逐组件消融加端到端对比，核心数字如下（出自 arxiv 摘要与公开推文复核）：

任务 骨干 指标 备注
ImageNet-1k 分类 ViT-5-Base 84.2% top-1 可比算力，超越 DeiT-III-Base 83.8%
ImageNet-1k 分类 vanilla ViT-Base ~82% 区间 原文未明确具体数字
扩散生成 FID（SiT backbone） ViT-5 1.84 vanilla ViT 同框架下为 2.06
扩散生成 FID（SiT backbone） vanilla ViT 2.06 给定 SiT 框架
  • 论文还报告了表示学习质量与空间推理行为的改善（具体指标未在摘要级别给出，按原文标注「原文未明确」）。
  • 跨任务迁移性：在多个 dense 与生成任务上一致提升，意味着这套组件组合不是「分类定制」。
  • 训练与推理差异不大，开源仓库 github.com/wangf3014/ViT-5 提供预训练权重。

备注：84.2% 与 DeiT-III 0.4 点的提升来自相同算力预算下的对照，作者使用 controlled compute 而不是同一 epoch 数。

亮点与局限

亮点

  1. 「不破坏骨骼」的兼容性：保持了 Attention-FFN 主干，任何已经在跑 vanilla ViT 的下游任务可以无痛迁移。
  2. 可解释的组件级改进：每条改动都对应一个被 LLM 实践验证过的细节，有据可循。
  3. 分类 + 生成统一提升：84.2% top-1 + SiT FID 1.84，证明改进不是某类任务偶发。
  4. 排除负向选择：明确说不引入 SwiGLU，给出 ablation 假设（与 LayerScale 过度叠加），是少有的诚实比较。
  5. 开源代码 + 权重：方便复现与对比。

局限

  1. 规模天花板未明：摘要未给出 Large / Huge 配置数字，对超大模型的扩展性原文未明确。
  2. 密集预测的定量数字有限：摘要提到「空间推理改善」，但 COCO / ADE20K 等具体数值未在摘要中给出（原文 PDF 未下载，对应数字标注「原文未明确」）。
  3. 混合卷积方案的对比缺失：ConvNeXt-V2、CoAt、MaxViT 这类结构层面的工作未在本摘要中量化对比。
  4. 缺乏多模态预训练配方：ViT-5 主要跑 ImageNet 单模态监督预训练，对 CLIP / LLaVA 这类 vision-language 流水线的兼容性原文未明确。
  5. APE + RoPE 双重位置编码有信息冗余：两次叠加的代价与收益对比，论文没给出独立 ablation 解释。

对工程落地的启发

  • 不重训 ViT？也想升级？ 直接用 ViT-5-Base/... 替换 vanilla ViT-Base 的初始化，微调成本低。
  • 生成 / 重建管线：如果项目用 SiT / DiT / MAR 跑图像生成，优先评估 ViT-5 作为 backbone，FID 1.84 vs 2.06 是无需任何代码改动的纯收益。
  • dense 任务：检测 / 分割 / 深度估计通常被 ViT 的 high-norm artifact 困扰，register token 是关键收益点。
  • 小团队 / 边缘部署：放弃 SwiGLU、剔除 bias、换 RMSNorm——这些改动都能省 FLOPs，不必等大模型蒸馏。

与同方向工作的关系

工作 关系
DeiT-III 直接 baseline，被 ViT-5 以「不增算力 + 84.2 vs 83.8」超越
ConvNeXt-V2 / CoAt 走「结构改造」，ViT-5 走「组件替换」，二者是互补路线
DINOv2 / DINOv3 走「自监督 + 大规模」路线，ViT-5 提供的是模块清单，可与 DINOv2 风格预训练叠加
FlexiViT / SigLIP 走「分辨率 / 多模态」路线，ViT-5 的组件理论上可正交地拼到这些骨干
PE / SigLIP-Vision 一旦 SigLIP 系列开源，关注 ViT-5 注册 token 在 vision-language 预训练里的复用

适合谁读

  • 多模态 / 视觉团队工程师：想要一份「不用动骨骼、只换零件」的现代化 ViT 清单。
  • 生成模型研究者：SiT / DiT / FLUX / MAR 类项目在选 backbone，ViT-5 是值得评估的默认候选。
  • Dense 预测工程师：检测 / 分割 / 深度估计里做 ViT 替换的团队，register token 是可直接复用的 trick。
  • 算法研究者：做架构 ablation、写 down-scaling / scaling law 的，希望看到「组件级 ablation」的人。
  • 不建议：只关心纯文本 LLM、不碰视觉的人；以及正在评估 结构级改动（CoAt、MaxViT、ConvNeXt）的人——ViT-5 不在这条路线上。

参考链接

  • arXiv 摘要：https://arxiv.org/abs/2602.08071
  • 实验性 HTML：https://arxiv.org/html/2602.08071v1
  • 代码：https://github.com/wangf3014/ViT-5
  • 提交日期：2026-02-08

工程落地与核查（Jay）

事实核查

核查项 原文表述 核查结果 存疑等级
arXiv 2602.08071 ViT-5 论文存在性 ✅ arXiv HTML 确认，标题 "Vision Transformers for The Mid-2020s" —
84.2% vs DeiT-III 83.8% ImageNet-1k top-1 ✅ 摘要原文确认，"84.2% top-1 accuracy under comparable compute, exceeding DeiT-III-Base at 83.8%" —
SiT FID 1.84 vs 2.06 扩散生成 FID 对比 ✅ 摘要原文确认，"1.84 FID versus 2.06 with a vanilla ViT backbone" —
github.com/wangf3014/ViT-5 代码仓库存在性 ✅ GitHub API 确认仓库存在 —
Darcet et al. (2024) Register Token Register Token 来源 ⚠️ 2024 年 Register Token 论文存在（DiT 相关的 Darcet 等人），引用合理，但原文未给完整标题或链接 低
Large/Huge 配置数字 规模扩展性 ⚠️ 原文摘要确未给出 Large/Huge 数字，解读已标注"原文未明确" 低
COCO/ADE20K 具体数值 密集预测数字 ⚠️ 原文摘要未给具体数字，解读已标注"原文未明确" 低

⚠️ APE + RoPE 双重位置编码存疑：解读指出论文未给出独立 ablation 解释双重叠加的代价与收益，这是一个真实未解决问题，不影响核心结论，但工程师在叠加其他位置编码方案时需谨慎。

可读性精修建议

  1. 表格中的 "~82% 区间" 应注明为"vanilla ViT-B 参考区间"而非精确数字；建议改为 "~82%（原文未给出 vanilla ViT-B 精确数字，此为参考区间）"。
  2. "SwiGLU 过度调控" 在代码注释中可补充 LayerScale × SwiGLU 叠加导致 Loss 震荡的具体表现（如训练早期 Loss 曲线毛刺化），方便工程判断是否在自己场景复现。

工程落地指南

适用场景

  • 图像分类 backbone 升级（ImageNet 1k / 22k）
  • 扩散模型 backbone 替换（SiT / DiT / MAR）
  • 检测 / 分割 / 深度任务骨干迁移

实际系统怎么用

路径 A：直接替换（最小改动）

# 安装（若有 pip 包）
pip install vit5-base  # 假设发布时提供官方包

# 微调 vs 零样本
# ViT-5-Base → 直接替换 timm 的 vit_base_patch16_224
import timm
model = timm.create_model('vit_base_patch16_224', pretrained=False)
# 加载 ViT-5 权重（github.com/wangf3014/ViT-5 的预训练权重）

路径 B：从源码迁移组件（推荐深度定制）

# 在已有 timm ViT 模型基础上，换以下 5 个组件（不碰 Attention-FFN 主干）
# ① RMSNorm → timm 需自定义或用 RMSNorm 实现
# ② LayerScale → 在每个 block 的 attention 和 mlp 输出残差加可学习对角矩阵
# ③ QK-Norm → 在 Attention forward 里对 Q, K 做 nn.functional.normalize
# ④ 2D RoPE → 用 timm 的 RoPE 或第三方实现（注意 grid_h, grid_w 维度）
# ⑤ Register Token → 在 sequence 末尾追加 n_register 个可学习 token

路径 C：扩散 backbone（SiT/DiT）

# 假设已有 SiT 框架，将 backbone 从 vanilla ViT 换成 ViT-5
# 仅需替换初始化权重，无需改 SiT 主循环
backbone = ViT5Backbone.from_pretrained("wangf3014/ViT-5")

主要坑点

  1. LayerScale 初始化敏感：论文用 ε ≈ 1e-5 或 1e-6 初始化，训练初期 Loss 震荡与此强相关；从零训练时务必保留，从头调大初始化值会导致深层失活。
  2. QK-Norm 算力开销：对 Q/K 各做一次 normalize 会额外引入 kernel 调度开销，在边缘设备（手机 / Jetson）上需测 FPS 损耗再做决定。
  3. Register Token 数量需调：论文未给出最优 k 值（几号数量）；建议跑 ablation 确认，一般 4-8 个足够。
  4. APE + RoPE 双重编码可能冗余：下游任务若已是纯相对位置模型（如 MAE 自监督），APE 的绝对位置归纳偏置可能是冗余的；建议 ablation 比对。
  5. github.com/wangf3014/ViT-5 仓库 stars=0：无社区验证，实际部署前建议自行跑 ImageNet 分类基准确认精度。
  6. 多模态兼容性未验证：在 CLIP / LLaVA pipeline 里替换 backbone，CLIP 对比学习目标的分布可能不兼容 ViT-5 组件（如 RMSNorm 对比 LayerNorm 的数值稳定性差异），需独立验证。

最小可跑命令集

# 评估（ImageNet-1k）
git clone https://github.com/wangf3014/ViT-5
cd ViT-5
pip install -e .  # 或 python setup.py install

# 微调示例（需要 ImageNet 数据）
python run_finetune.py --model vit5_base --data-dir /path/to/imagenet --epochs 100

# 扩散 backbone 替换（以 SiT 为例，假设 SiT 已安装）
python -c "from vit5 import ViT5Backbone; m = ViT5Backbone.from_pretrained(); print('loaded')"