清华+阿里发布ViT³：解锁视觉TTT，突破Transformer复杂度瓶颈

本文第一作者为清华大学博士生韩东辰，主要研究方向包括基础模型架构、持续学习等。通讯作者为黄高副教授。

序列建模是大语言模型、计算机视觉等领域的基础共性问题。当前通用的 Transformer 模型计算复杂度随序列长度平方增长，在长序列任务中面临显著的计算挑战。因此，研究者们一直在探索具有线性计算复杂度的高效序列建模方法。

测试时训练（Test-Time Training，TTT）模型是一种新兴的序列建模范式，它将注意力操作重新定义为一个在线学习过程，在每次推理时用 Key-Value 对构建一个轻量化的内部模型。不同于一个确定性的模型设计，这一范式开启了一个丰富的线性复杂度设计空间，有望实现高精度的高效序列建模。

然而，TTT 灵活的设计空间是一把双刃剑：当前测试时训练模型缺少系统性的理解和设计原则，如何从丰富的设计空间中构建一个理想的视觉 TTT 模型、如何实现进一步提升，都是目前极具挑战的问题。

为此，本文在视觉领域系统性地研究了测试时训练模型的构建和效果。通过一系列的实验和分析，我们总结了高精度、高效视觉 TTT 模型的六条设计原则，并讨论了一些未来改进方向。

基于这些发现，本文构建了纯 TTT 架构 Vision Test-Time Training （ViT³）模型，它在各类视觉任务中表超越了现有线性复杂度模型，同时保持并行计算和高推理速度，为未来视觉 TTT 的研究提供了一个有力的基线。

• 论文链接：https://arxiv.org/abs/2512.01643
• 代码链接：https://github.com/LeapLabTHU/ViTTT

序列建模的新视角

在传统的理解视角下，我们通常将注意力计算理解为用注意力分数对 value 进行加权求和。实际上，我们可以从一个用 context 构建模型的视角来理解注意力。

图 1：Softmax 注意力、线性注意力、TTT 模块示意图

因此，两种经典的注意力范式都可以理解为使用 key 和 value 构造了一个小的模型，进而将 query 输入这个模型以实现建模。

于是，一个核心问题是：我们能否在实现压缩的同时保持模型的性能，同时实现高精度和高效性？

对于这个问题，研究者们回答是肯定的。回顾整个机器学习领域的发展，最为成功的压缩算法莫过于深度学习。深度学习通过梯度下降，把数据信息压缩到更紧的模型权重空间中，通常能够实现十倍甚至百倍的压缩率。因此，测试时训练模型正是尝试在序列建模任务上再次套用深度学习的成功。

如图 1（c）所示，测试时训练模型尝试将 key 和 value 的信息压缩到一个小的内部模型中，这个内部模型的结构几乎是任意的。为了实现这一点，TTT 将 key 和 value 视为一个小型的 “数据集”，要求模型在输入 key 时重建对应的 value，通过梯度下降对内部模型权重进行更新：

其中，ℒ是一个自监督的重建损失函数，比如 L2 损失。在几次更新之后，我们认为 key 和 value 的信息已经较好地压缩到了 TTT 内部模型权重之中。这时，我们将更新后的内部模型应用于 query，通过一次简单的前向传播，让 query 从内部模型的权重中读取 key 和 value 的信息。整个测试时训练计算过程的计算复杂度正比于内部模型的计算复杂度。

因此，当选取一个简单的两层 MLP 等线性复杂度模型作为内部模型时，TTT 具有线性空间和计算复杂度；同时，基于学习的压缩和非线性状态空间使得 TTT 具有良好的表达能力。

视觉 TTT 如何设计

如前所述，测试时训练模型在每个 Block 内部包含一次小的 “深度学习”，具有极高的设计自由度。当前 TTT 缺乏充分的理解和设计原则，阻碍了其在视觉领域的发展。本文在视觉领域系统性地探索了 TTT 的设计空间，聚焦于两个核心方面：（1）内部训练设置，如损失函数、学习率、batch size、训练轮数等（2）内部模型设计，如内部模型架构、宽度、深度等。本文基于 DeiT-S 进行探索，将原始的 Softmax 注意力替换为 TTT 模块，在 ImageNet-1K 上进行实验。

1. 内部训练设置

图 2：TTT 内部训练损失函数的选择

TTT 模型的内部训练所需要的 key、value 等，是在整个外部网络的端到端训练过程中优化。在进行外部端到端训练时，我们将内部梯度视为表达式，让整个模型实现端到端的反向传播。在这个过程中，我们需要对内部训练梯度的表达式再次求导，计算 “梯度的梯度”，这是元学习中的经典操作。在端到端优化过程中，value 的投影矩阵的外部梯度为：

观察 2：单次全批量内部训练（batch gradient）在视觉任务中效果较好。

图 3：TTT 内部训练 batch size 和 epoch

如图 3 所示，B=N 表示使用全部 N 个 key-value 对进行全批量内部梯度下降，而 B=N/2、B=N/3、B=N/4 则对应小批次梯度下降，即把数据集划分为 2、3、4 个连续的小批次，每个 epoch 的内部训练执行 2、3、4 次参数更新。

本文认为，连续小批次梯度下降会引入显式的因果关系：(1) 前序小批次的更新会改变内模型权重，影响后续批次的梯度计算；(2) 后续批次的更新可能覆盖前序批次产生的参数更新。这种因果关系非常适合自然语言这类因果数据，但在视觉任务中通常是不必要的。

多个 epoch 的内部训练能够提升性能，但是会显著降低模型推理速度。

观察 3：相对较大的内部学习率（如 1.0）更适合 TTT。

图 4：TTT 内部训练学习率

在图 4 中，我们验证了不同内部学习率的影响。过小的学习率不足以快速更新原始的内部模型权重，过大的学习率会导致训练不稳定。在本文验证的视觉任务中，内部学习率 1.0 较为合适。

2. 内部模型设计

图 5：TTT 内部模型结构。r 和 l 分别表示模型宽度比例和模型深度

观察 4：扩大内部模型容量能够持续提升 TTT 的序列建模能力。

如图 5 第一部分所示，我们将 TTT 内部模型实现为一个 SiLU 激活函数的两层 MLP，并逐渐将宽度比例（隐藏维度 / 输入维度）从 1.0 增加到 4.0。我们可以观察到，模型的性能随着 TTT 内部模型的容量持续提升。这是测试时训练范式的一个重要优势，即在外部模型尺寸固定的条件下，我们可以通过简单地扩展内部模型尺寸来实现更好的序列建模效果。

但是，值得指出的是，在推理时，一个内部模型引入的计算量大于相同尺寸的外部模块。在推理阶段，一个外部模块仅需要一次前向传播；而一个内部模型则至少需要进行（1）key 前向传播，1 倍前向传播计算量（2）内部损失反向传播，约 2 倍前向传播计算量（3）将更新后的内部模型应用于 query，1 倍前向传播计算量。一个内部模型的计算量是同尺寸外部模块的近 4 倍。因此，尽管简单扩展内部模型能显著提升模型能力，本文认为设计轻量的高效内部模型仍然是一个重要的研究问题。

观察 5：在当前的 TTT 实现下，较深的内部模型存在优化困难。

在上一部分中，我们扩展了内部模型的宽度。另一个常见的方案是扩展模型的深度。如图 5 第二部分所示，我们分别将内部模型实现为单层线性层、两层 MLP、三层 MLP，以探索内部模型深度对性能的影响。结果表明，扩展内部模型深度时，模型性能持续下降。在一定程度上，这是一个不符合预期的结果。因此理论上更深的内部模型具有更大的容量，应该能够实现更好的序列建模。

图 6：较深的 TTT 内部模型的训练、测试 loss 均更高

本文认为，优化问题是导致较深内部模型性能不佳的原因。为了验证这一点，我们在图 6 中展示了上述三个模型的训练和测试 loss。可以看到，采用更深内部模型时，模型的训练 loss 更高，这说明性能衰退不是因为过拟合，而恰恰是欠拟合。更深的内部模型具有更强的拟合能力，而在真实训练中相比于浅内部模型欠拟合，说明深内部模型存在优化问题，远远没有达到其理论表征能力。在图 5 第二部分的结果进一步验证了这一观点：将两层 MLP、SwiGLU 的输出层强行限制为 identity 能够让模型性能显著提升，说明较深的内部模型显著欠优化。

这种优化困难与 TTT 的元学习属性有关。一方面，内部模型的初始权重是外部训练过程中的可学参数，较深的内部模型可能导致这个初始参数难以优化。另一方面，即便有较好的初始参数，深内部模型也可能在内部训练过程中遇到梯度消失和爆炸问题，影响效果。

解决较深内部模型的优化问题，是 TTT 的一个重要未来方向。理论工作表明，神经网络的拟合能力随深度指数增长，这正是当前神经网络成功的核心原因。因此，较深的内部模型在实现高精度测试时训练序列建模中具有突出潜力。

观察 6：在视觉任务中，内部模型非常适合采用卷积架构。

在 Transformer 出现之前，卷积操作一直是计算机视觉领域的基石。在 TTT 中，我们完全可以将内部模型实现为一个小的卷积网络，而不需要将其限定为 MLP、SwiGLU 等单 token 组件。在图 5 第三部分中，我们展示了两种最简单的实现，将内部模型实现为 3x3 卷积或 Depthwise 卷积。可以看到，这两种实现均取得了显著的性能增益。

本文认为，采用卷积网络作为 TTT 内部模型，能够实现局部信息和全局信息的完美融合。具体而言，在 TTT 的内部训练过程中，key 和 value 的全局信息被压缩到卷积核的权重中。在 query 通过内部模型前向传播时，它很自然地同时得到局部信息（卷积核感受野）和全局信息（包含全局信息的卷积核权重），因此完美契合视觉任务。

ViT³：视觉测试时训练模型

基于以上探索，本文提出纯 TTT 架构 Vision Test-Time Training （ViT³）模型。具体而言，对于内部训练，采用一次全批量梯度下降、学习率 1.0、点积损失函数；对于内部模型，采用简化的 SwiGLU 和 Depthwise 卷积。作为 TTT 架构的一个 baseline，ViT³ 可以应用于各类视觉 Transformer 框架，完成各类任务。

本文在分类、检测、分割、生成等代表性视觉任务上验证了 ViT³ 的性能。实验结果表明，ViT³ 能够超越各类先进的线性复杂度设计，如线性注意力、视觉 Mamba 模型。这一结果验证了 ViT³ 设计的有效性，同时凸显了测试时训练模型在高效、可扩展的线性复杂度视觉序列建模领域的突出潜力。

以下是部分实验结果，详细结果请见论文原文：

图 7：ImagNet-1K 图片分类实验

图 8：高分辨率下游任务 —— 物体检测

图 9：图片生成实验

图 10：推理 Throughput 与显存随输入分辨率变化

总结与展望

设计高效、高表达能力的序列建模范式一直是计算机视觉领域的核心课题。

本文在视觉领域系统性地探索了测试时训练（Test-Time Training，TTT）这一新兴序列建模范式的设计空间。

通过大量实验，本文将观察与分析总结为六条设计见解，为构建高效的视觉 TTT 模型明确了设计准则，并讨论了一些未来的研究方向。

基于上述研究，本文提出了视觉测试时训练模型（ViT³），一种面向视觉序列建模的纯 TTT 架构。ViT³ 在多项视觉任务上取得了具有竞争力的结果，为线性复杂度模型提供一个有力的基线。

我们期望本文的研究结论与分析能够帮助视觉 TTT 模型的后续研究工作。