ViT³：把 Attention 改写成测试时训练的视觉骨干

Date 2026-05-23 · Category tech · Status finished · Confidence likely
论文解读, 计算机视觉

论文信息

论文题目：ViT³: Unlocking Test-Time Training in Vision

论文方向：视觉序列建模 / 高效视觉骨干 / 线性复杂度注意力替代

代码地址：https://github.com/LeapLabTHU/ViTTT

这篇论文关注的不是一个小的注意力改写技巧，而是一个更大的问题：

如果 Softmax attention 的 O(N^2) 成本越来越难承受，那么视觉模型除了 linear attention 和 Mamba 之外，还有没有另一条线性复杂度路线？

作者给出的答案是 TTT，也就是 Test-Time Training。但这里的 TTT 不是传统意义上的测试时域适应，而是把一次 attention 运算重写成：

先用当前输入里的 K、V 临时训练一个小 inner model
再用这个更新后的 inner model 去处理 Q

ViT³ 方法总览，图源：论文首页 Figure 1 裁剪

截至 2026-05-23，官方仓库已经公开了最小 TTT block 实现，以及 ViT³ / H-ViT³ 的 ImageNet 分类代码。README 中也明确标注该工作为 CVPR 2026 oral。

核心动机：Softmax attention 太贵，线性替代又常常不够强

这篇论文的出发点非常直接。

视觉 Transformer 之所以强，很大程度上来自 Softmax attention 的全局建模能力。但 Softmax attention 对 token 数量的复杂度是：

O(N^2)

图像分辨率一高，token 数就会迅速增长。分类还勉强能承受，到了检测、分割、生成这类高分辨率任务，二次复杂度会很快变成吞吐和显存瓶颈。

社区当然已经提出过很多线性复杂度替代方案，例如：

  • linear attention
  • Mamba / selective scan
  • 其他状态空间模型

但这类方法常见的问题是：状态表达力不够强。简单说，它们往往把很长的上下文压缩进一个相对固定、相对有限的状态里，虽然成本降下来了，但也容易丢掉关键信息。

作者对三种路线的理解可以压缩成下面这张图：

  • Softmax attention：很强，但 O(N^2)
  • linear attention：很省，但压缩太朴素
  • TTT：用一小段在线学习过程代替显式注意力矩阵

所以这篇论文真正想回答的问题不是“能不能再造一个注意力变体”，而是：

能不能把注意力改写成一种在线学习过程，从而同时保住线性复杂度和更强的表达能力？

核心想法：把 attention 改写成 online learning

论文最关键的概念是：

K、V 不再直接参与 QK^T 的显式匹配
而是被当成当前输入的一个临时训练集

也就是说，模型不再直接做：

O = Softmax(QK^T)V

而是改成：

1. 用 (K, V) 临时训练一个小 inner model F
2. 得到更新后的参数 W*
3. 用 F_{W*}(Q) 输出结果

这个视角非常重要，因为它把“上下文压缩”从一个固定矩阵，改成了一个可学习的小函数。作者认为，这正是 TTT 相比 linear attention 更有潜力的地方。

TTT Block 长什么样

从整体骨架看，TTT block 并没有推翻 Transformer。它依然是：

Norm -> TTT Block -> Residual -> Norm -> FFN -> Residual

TTT block 与 Transformer block 的对应关系，图源：论文 Figure 2 裁剪

真正变化的是中间这个 TTT Block。

它内部大致分成三步：

  1. 先把输入投影成 Q, K, V
  2. 用 K, V 临时训练 inner model
  3. 再用更新后的 inner model 去处理 Q

如果只看 TTT block 的内部机制，论文首页右侧那张图其实更关键。它把三种序列建模方式放在了一起：

  • Softmax attention：直接做 QK^T，再和 V 结合
  • Linear attention：把 K, V 压成一个线性状态
  • TTT layer：把 K, V 当成当前输入的临时训练集，更新 inner model 的参数

TTT layer 内部机制，图源：论文 Figure 1 裁剪

这张图里最需要记住的是最下面那一行：Q 并不是直接和 K 做匹配，而是先让 K, V 通过一次 inner training 改写权重 W*，然后再用更新后的 W* 去处理 Q。这就是整篇论文最核心的概念变化。

从官方代码看，作者最终采用的是两路并行结构：

  • 一路是 GLU 风格的主分支
  • 一路是 3x3 depthwise conv 分支

最后把这两路结果拼起来，再投影回原始维度。

这也是为什么这篇论文虽然是“替代 attention”，但最后实现出来的模块其实非常有视觉味道：它不仅保留了全局上下文学习，还显式混入了卷积式局部归纳偏置。

H-ViT³：为什么它是分层骨干

论文最后给了两套模型：

  • ViT³：不分层，类似经典 ViT / DeiT
  • H-ViT³：分层，类似 Swin 一类的现代视觉 backbone

这里“分层”的意思很简单：

  • 前面保持高分辨率
  • 后面逐 stage 下采样
  • token 数越来越少
  • 通道数越来越高

H-ViT³ 正是 4-stage 金字塔结构。以论文附录里的结构表为例：

H-ViT³ 架构表，图源：论文附录 Table 12 裁剪

例如 H-ViT³-T 的 4 个 stage 是：

Stage 1: B(64, 2) × 1
Stage 2: B(128, 4) × 3
Stage 3: B(320, 10) × 9
Stage 4: B(512, 16) × 4

其中 B(C, H) 表示：

  • 一个 TTTBlock
  • 通道维度是 C
  • head 数是 H

所以 H-ViT³ 你可以直接理解成：

一个现代分层视觉骨干，只不过把每个 stage 里的 attention block 换成了 TTT block。

设计空间研究：这篇论文真正的主贡献

如果只看最终模型，很容易低估这篇论文的贡献。它真正有价值的地方，其实是前半部分的大量 controlled study。

作者系统地研究了两类设计问题：

1. inner training 怎么配

  • 什么损失函数合适
  • 学习率该多大
  • full-batch 还是 mini-batch
  • 一次 inner update 做几轮

结论是：

  • MAE / Smooth L1 这类会让关键二阶导数消失的损失不适合
  • 视觉里最合适的是 single-epoch + full-batch
  • inner learning rate 开到 1.0 很有效

2. inner model 怎么配

  • 增大宽度通常有利
  • 但加深 inner model 反而很难优化
  • depthwise conv 作为视觉 inner module 特别合适

这两部分研究最后凝结成了 ViT³ 的最终配方。

所以更准确地说，ViT³ 不是作者拍脑袋设计出来的，而是他们先用消融把“视觉 TTT 应该怎么做”摸清楚，再把结果落到最终架构里。

分类实验：TTT 路线已经能进入第一梯队

论文先在 ImageNet-1K 上做分类验证，而且把分层和不分层分开比较。

ImageNet 分类结果，图源：论文 Table 5/6/7 裁剪

这张图有两类信息。

分层 backbone 对比

上半部分是 H-ViT³ 和 ConvNet、Transformer、Mamba、Linear attention 的比较。

比较亮眼的结果是：

  • H-ViT³-T‡ = 84.0
  • H-ViT³-S‡ = 84.9
  • H-ViT³-B‡ = 85.5

尤其 H-ViT³-B‡ = 85.5，已经进入最强分层骨干的竞争区间。

不分层 backbone 对比

右下角是 ViT³-T / S / B 与 DeiT / Vim / Agent-DeiT 的比较。

例如：

  • ViT³-S = 81.6
  • DeiT-S = 79.8
  • Vim-S = 80.3
  • Agent-DeiT-S = 80.5

这说明即便不引入分层金字塔结构，只把 attention 替换成 TTT block，模型本身也确实带来了收益。

一句话总结分类实验：

TTT 不只是概念验证，它已经是能够和强视觉骨干正面竞争的一条路线。

检测与分割：长序列场景下优势更明显

论文在 COCO 检测和 ADE20K 分割上都做了验证。

COCO 检测结果，图源：论文 Table 8 裁剪

这张检测表最值得记住的不是某个单一数字，而是整体趋势：

  • H-ViT³ 稳定强于多数 Mamba 和 linear attention 模型
  • 在高分辨率输入下，它比很多线性复杂度替代路线更能撑住性能

作者的解释也很重要：

检测和分割的输入分辨率高
token 序列很长
常常出现 N >> d

在这种情况下，很多线性状态模型会受到状态容量限制，而 TTT 因为通过 inner learning 形成了更强的非线性状态，往往更不容易吃亏。

语义分割与生成结果，图源：论文 Table 9/10 裁剪

看左半边的 Table 9：

  • H-ViT³-T = 48.0 mIoU
  • H-ViT³-S = 50.2
  • H-ViT³-B = 51.7

它们明显强于：

  • VMamba
  • SOFT++
  • VVT

但作者也很诚实：在最强的 dense prediction Transformer 面前，例如 TransNeXt-B = 53.0，H-ViT³-B = 51.7 仍然有差距。

所以分割实验传达出的正确信息不是“TTT 全面超越 Transformer”，而是：

它已经是非常强的线性复杂度分割骨干，但距离最强视觉 Transformer 还有继续提升的空间。

生成实验：替换进 DiT 后，FID 全面改善

还是看上面那张图的右半边 Table 10。

作者把 DiT 中的 Softmax attention 替换成 ViT³ block，得到 DiT3，然后在 ImageNet-1K 的 class-conditional generation 上比较。

最重要的观察是：

Small 模型

  • DiT-S/8: 153.60 -> DiT3-S/8: 143.49
  • DiT-S/4: 100.41 -> DiT3-S/4: 93.77
  • DiT-S/2: 68.40 -> DiT3-S/2: 62.65

Base 模型

  • DiT-B/8: 122.74 -> DiT3-B/8: 120.41
  • DiT-B/4: 68.38 -> DiT3-B/4: 65.25
  • DiT-B/2: 43.47 -> DiT3-B/2: 39.31

也就是说：

只把注意力模块替换成 TTT block，FID 在所有设置上都变好了，而且参数和 FLOPs 的额外开销并不大。

这说明 TTT 并不是只对分类有效，它对生成任务同样成立。这也大大增强了作者的主张：TTT 更像一种通用视觉序列建模范式，而不是某个专门为分类调出来的小技巧。

Figure 4：效率优势不是理论上的，而是真能跑出来

论文最后给了一张特别关键的效率图。

吞吐与显存对比，图源：论文 Figure 4 裁剪

它比较的是：

  • DeiT-T：标准 Softmax attention
  • ViT³-T：TTT 路线

左图是 FPS vs Resolution，右图是 Memory per image vs Resolution。

最重要的结论直接写在图上：

  • 在 1248 × 1248 分辨率下，ViT³-T 相比 DeiT-T 加速 4.6 倍
  • 单图显存占用 降低 90.3%

这张图的意义非常大，因为它把“线性复杂度”从公式变成了真实可见的工程收益：

  • 分辨率低时，优势还不夸张
  • 分辨率一旦上去，Softmax attention 的二次复杂度惩罚会迅速显现
  • 而 TTT 的吞吐和显存曲线都增长得平缓得多

因此，如果你的任务本来就要处理高分辨率输入，这篇论文的效率结论是非常有现实意义的。

从官方代码看，这个模块的落地实现有什么特点

我额外看了官方仓库里的实现，发现几件很值得记住的事情。

第一，仓库里有一个单独可复用的最小实现 ttt_block.py，说明作者确实把它当成通用 plug-in 模块，而不只是论文里的抽象概念。

第二，代码里的 TTT block 并不是“所有 head 完全对称”，更准确地说是：

  • 一个 GLU 风格主分支
  • 再加一个额外的 3x3 depthwise conv 分支
  • 最后融合回原维度

第三，作者为了效率，没有在 inner loop 上直接依赖 autograd.backward()，而是手工写了 closed-form gradient。这一点很说明问题：

TTT 不只是一个新架构想法，它的工程实现本身也在和效率做妥协与优化。

第四，截至 2026-05-23，官方仓库公开得最完整的是分类代码。检测、分割、生成的完整任务实现并没有一起放出。因此如果想完全复现论文全部实验，工程工作量仍然不小。

这篇论文真正值得你记住的 5 件事

1. 它的主要贡献不只是模型，而是设计原则

这篇论文真正有价值的地方，是第一次比较系统地回答了：

视觉 TTT 到底应该怎么设计？

2. 它把 attention 改写成了在线学习过程

不是直接算 QK^T，而是：

(K, V) -> 临时训练 inner model
Q      -> 用这个 inner model 读出结果

3. 它是 linear attention / Mamba 之外的第三条线性复杂度路线

相比简单线性状态压缩，TTT 更强调“把上下文写进一个非线性小模型”。

4. 当前瓶颈主要是优化，而不是思路

作者已经证明更深的 inner model 理论上更有潜力，但现阶段优化跟不上，所以这仍是未来最值得继续挖的方向。

5. 它已经是很强的 baseline，但还不是终局

ViT³ / H-ViT³ 已经能稳定强于多数线性复杂度替代方法，也能逼近甚至局部碰到强 Transformer；但在最顶级 dense prediction backbone 面前，它还没有完全封顶。

总结

ViT³ 最值得肯定的地方，不是单纯把 attention 换了个公式，而是把视觉序列建模问题重新表述了一遍：

注意力不一定要显式构造 token-token 相关矩阵，也可以理解成“当前输入驱动的一次在线学习”。

在这个视角下，作者不只提出了一个可运行的 TTT 架构，还通过系统实验把视觉 TTT 的设计空间初步摸清楚了。最终结果也很有说服力：

  • 分类上，它已经进入强视觉骨干的竞争区间；
  • 检测和分割上，它明显强于多数线性复杂度模型；
  • 生成上，把 ViT³ block 替换进 DiT 后，FID 全面改善；
  • 高分辨率下，它在吞吐和显存上都展现出真实的线性复杂度优势。

所以这篇论文最准确的定位不是“又一个高效 attention 变体”，而是：

它为视觉里的 TTT 路线建立了一套可复用的设计原则，并给出了一个足够强、足够完整的 baseline。

See also