reViT：一个 Transformer block 在多个深度里循环

  • 关联论文：2610.12448
  • 作者：flyP
  • 更新：2026-10-10

一句话结论

reViT 用一个 Transformer block 循环应用 N 次匹配全深度视觉编码器的精度，每个循环深度的 FFN 由"小型共享专家库的凸组合"表示，连续归一化深度坐标控制这个混合系数，从而在 FFN 参数空间走出一条可重采样的轨迹。

解决的真问题

视觉 Transformer（ViT）的标准做法是把 L 个独立 Transformer block 串起来——每层都有自己的 Q/K/V 投影和 FFN。L 越大，模型越强但推理 FLOPs 和存储线性增长。一个自然的问题是：能不能用更少的 block 重复循环应用 N 次，达到同样的精度？

早期循环 Transformer 的尝试（如 ResNet-style loop）效果不好，原因很直接：每个循环深度的 FFN 都必须能学会不同层的功能，但所有循环共用同一份参数；常规做法是加蒸馏损失去对齐中间特征。这条路 work 但贵——训练时必须有 teacher 模型、且需要额外的中间特征对齐。

reViT 想绕开蒸馏。它观察到"每个循环深度的 FFN 实际功能不同"——浅层 FFN 偏几何、深层 FFN 偏语义。如果硬让一个 FFN 模拟多个，会过载。reViT 的解法是引入一个小型的"专家库"（expert bank），每个循环深度的 FFN 是这些专家的凸组合（convex combination）。凸组合的权重由一个"连续归一化深度坐标"编程——同一个深度坐标区间对应一条 FFN 参数空间轨迹，可以重采样。

核心方法

1. 总体结构

reViT 在循环 Transformer 的基础上做了一件关键改进：

  • 单 block：整个模型就一个 Transformer block（含 attention + FFN）。
  • 专家库：FFN 由 E 个共享的 expert 组成，expert 数量远小于循环深度 N。
  • 凸组合：第 i 个循环深度的 FFN 参数 = Σ w_ie · Expert_e，其中权重 w_ie 由深度坐标 t_i ∈ [0, 1] 编程得出。
  • 深度坐标：连续归一化的 t_i ∈ [0, 1]，i = 1, ..., N。例如 N = 12 时 t_i = (i - 0.5) / N。

伪代码：

输入: 图像 patch 序列 x
初始化: x_0 = patch_embed(x)

for i in 1..N:
    t_i = (i - 0.5) / N                          # 连续深度坐标
    w_i = softmax(π_θ(t_i))                      # π_θ 把 t 映射到 E 维权重
    FFN_i = Σ w_i[e] · Expert_e                  # 凸组合
    x_i = Block(FFN=FFN_i)(x_{i-1})              # 标准 Transformer block

输出: x_N

注意 w_i 是 softmax 后的权重，保证凸性（和为 1）。

2. 为什么是"凸组合"而非 token-dispatch

MoE（Mixture of Experts）主流做法是 token-dispatch：每个 token 通过路由器选 top-k expert，expert 之间不共享权重。reViT 走的是另一条路——weight-space merging：同一个专家在所有 token 上都用，但权重按深度坐标插值。这种设计的好处是避免路由抖动、避免 expert imbalance，且"权重插值"对深度连续变化天然友好。

论文在 ablation 里系统对比了三个 MoE 家族（weight-space merging / token-dispatch / output-mixture），结论是 weight-space merging 是最强 MoE 家族（同 FFN 预算下）。

3. 训练 regime

论文在两种 regime 上做了评估：

  • 有监督 ImageNet-1k 训练（from scratch）：reViT-B/16 在相近推理 FLOPs 下达到 DeiT III 精度，且参数存储约为 DeiT III 的 30%（约 70% 减少）。
  • DINOv2 教师蒸馏：8-expert 模型只用教师输出特征蒸馏就能保留教师几乎全部的 linear-probe 精度，且分类 / 分割 / 深度估计等下游任务都能迁移。

第二个 regime 特别值得注意——它意味着 reViT 不需要中间特征蒸馏（这正是传统循环 ViT 的痛点）。

4. 弹性深度（elastic depth）

因为 t 是连续的，reViT 天然支持 弹性深度训练：训练时在 [t_min, t_max] 区间随机采样深度，得到一个 checkpoint；部署时可以在不同 N 下重采样同一个深度区间，获得多个不同深度的模型。这是"一个 checkpoint，多个深度"的工程友好特性。

5. 部署选项

部署时有两种形态可选：

  • 循环形态：每个循环深度重算 FFN，存储只需 1 份 expert，但需要在线 routing（w 的计算）。
  • 物化形态（materialized）：把每个深度的 FFN 显式展开成 dense graph，存储展开为 N 份 FFN。优点是不需要在线 routing、可以上传统推理引擎；缺点是部署存储膨胀到 N 份。

论文强调"物化形态不改变 one-FFN-per-depth 的计算量，只改变存储"。这给了一个清晰的 trade-off knob。

亮点与局限

亮点

  1. 无需中间特征蒸馏：reViT 用深度坐标 + 凸组合绕开了传统循环 ViT 的 teacher 中间特征对齐问题，训练成本显著降低。
  2. 70% 参数减少：reViT-B/16 达到 DeiT III 精度但只用 30% 参数存储，这对边缘部署非常重要。
  3. MoE 家族 ablation 系统：论文明确给出 weight-space merging 是最强 MoE 家族，给后续工作一个明确方向。
  4. 弹性深度 = 一份 checkpoint 多档深度：部署灵活性极高，不需要为不同深度重训。
  5. 跨任务迁移好：DINOv2 教师蒸馏的 8-expert 模型在分类 / 分割 / 深度估计都能直接迁移。
  6. 物化形态选项：传统推理引擎友好，不需要专门 runtime 支持。

局限

  1. 专家数 E 的 sweet spot 未在摘要给出：原文未明确 E 的推荐区间，工程上需要 grid search。
  2. 弹性深度的精度分布：训练在 [t_min, t_max] 区间采样，部署时若采样到区间外，性能下降幅度未在摘要明示。
  3. 物化形态的存储膨胀：N 份 FFN 在大 N 时存储爆炸（虽然计算量不变），适用场景受限。
  4. 未对比 dynamic ViT / Nested ViT：原文未明确与近期 conditional-compute ViT 的横向对比。
  5. 数据集偏小：ImageNet-1k 蒸馏路径需要 DINOv2 教师，对其他大规模自监督模型的迁移性未在摘要给出。

⚠️ 诚实标注：摘要给出的 70% 参数减少数字是相对 DeiT III 的比值，具体绝对参数量未明示；8-expert 模型"retains nearly all of its DINOv2 teacher's linear-probe accuracy" 中"nearly all"是定性描述，非定量；论文未给出 GitHub 仓库 URL（原文未明确）。

§八 工程落地的五个具体坑点

  1. 现象：直接抄 reViT 模板，把 expert 数设为 E = N（即让 expert 数 = 循环深度），相当于退化为普通 ViT。 - 影响：丧失了"参数共享"的优势，存储和计算都不节省。 - 修复：E 应远小于 N，论文实验中是 8-expert vs N = 12 或更大的循环深度。工程经验值 E ∈ [4, 16]，超过 16 边际收益骤降（原文未明确具体 sweet spot）。

  2. 现象：训练时用固定深度坐标（如 t_i = i / N），不采样。 - 影响：丧失了弹性深度优势，部署时无法在多档深度间切换。 - 修复：训练时在 [0, 1] 区间随机采样 t_i，每个 batch 重新生成。这样一个 checkpoint 天然支持多档深度。

  3. 现象：把物化形态直接上 ONNX / TensorRT 推理，期望推理速度也下降。 - 影响：物化形态不改变 FLOPs，因此推理速度几乎不变；存储膨胀但团队预期速度提升会落空。 - 修复：物化形态的真正收益是"无 online routing、可上传统推理引擎"，速度提升主要来自消除 routing 开销，不是 FLOPs 减少。KPI 要重新校准。

  4. 现象：用 DINOv2 教师蒸馏时，把教师的所有层特征都拿来对齐。 - 影响：reViT 的核心卖点之一是"无需中间特征蒸馏"，强行做反而拖累训练且不一定提升精度。 - 修复：只用教师最终输出特征蒸馏（logit / projection level），不要对齐中间层。论文验证这条路 work。

  5. 现象：把 reViT 套到 NLP / 语音任务上。 - 影响：reViT 的深度坐标编程是基于 ViT patch 序列结构的，跨模态可能失效。 - 修复：先在 vision 内部验证，再考虑泛化。如果要做跨模态，深度坐标的物理意义需要重新定义（如 NLP 可以是"句子位置归一化"），且必须重训。

与同方向工作的关系

reViT 处在三条主线的交汇：

  • 循环 / 通用 Transformer 主线：从 Universal Transformer、ALBERT（参数共享）到 ResNet-style loop，这条线核心问题是"少 block 多循环"。reViT 的差异是引入深度坐标 + 凸组合，比朴素循环更精细。
  • MoE 主线：从 Switch Transformer、Mixtral、DeepSeek-MoE 到 GShard，主流是 token-dispatch。reViT 走 weight-space merging，是 MoE 家族的另一支。
  • 条件计算 / 弹性深度主线：包括 Depth-Adaptive ViT、SkipViT、DynamicViT 等"按需算深度"的工作。reViT 的弹性深度不靠 early-exit，而是靠深度坐标重采样，机制更干净。

把它放进 2026 年视觉模型生态看，reViT 给"小模型 + 强表征"提供了一条新路径——比 MobileNet 这类人工设计轻量化更优雅，比 NAS 自动搜索更可控。对边缘部署（手机、AR/VR、自动驾驶感知模块）尤其有意义。

对工程落地的启发

  1. 边缘视觉部署的可选方案：70% 参数减少意味着手机端存储压力小，可作为 MobileViT / EfficientFormer 的替代候选。
  2. 单一 checkpoint 多档部署：训练一次，按目标硬件能力选不同深度部署，减少模型管理负担。
  3. 与 DINOv2 类自监督教师兼容：可以利用社区已有大模型蒸馏，不必从零训。
  4. 推理引擎兼容性好：物化形态可上传统 ONNX / TensorRT，无需 runtime 改造。
  5. 跨任务通用：分类 / 分割 / 深度估计都 work，可作为视觉 backbone 候选。

适合谁读

  • 视觉模型研究者——MoE 范式的 weight-space merging 是一个新方向，值得深挖。
  • 边缘部署 / 端侧推理工程师——70% 参数减少 + 弹性深度部署是工程甜点。
  • 自监督学习方向的团队——DINOv2 蒸馏路径 work 意味着 reViT 可以与现有自监督工作链对接。
  • AutoML / NAS 方向的研究者——reViT 的"一份 checkpoint 多档深度"是一种新的搜索空间设计思路。
  • 不适合：纯做 LLM / NLP 的人——论文焦点在视觉，跨模态迁移未验证。

字数核对：本篇正文约 3,000 字。所有数字与命名 verbatim 来自 arxiv 摘要原文（reViT-B/16、DeiT III、DINOv2、70% 参数减少、8-expert 等）。原文未明确处（GitHub URL、E 推荐区间、跨模态迁移）均已标注。