清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

序列建模是大语言模型与计算机视觉等领域的核心基础问题。目前广泛采用的Transformer模型,其计算复杂度会随序列长度呈平方级增长,这在处理长序列任务时带来了显著的计算瓶颈。因此,学术界一直在积极探索具备线性计算复杂度的高效序列建模新方法。

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

然而,TTT灵活的设计空间也是一把双刃剑。当前,这类模型缺乏系统性的理解和设计原则,如何从广阔的设计空间中构建理想的视觉TTT模型,以及如何进一步提升其性能,都是极具挑战性的问题。

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

基于这些发现,我们构建了纯TTT架构的Vision Test-Time Training(ViT³)模型。该模型在多种视觉任务中超越了现有线性复杂度模型,同时保持了并行计算能力和高推理速度,为未来视觉TTT研究提供了一个强有力的基线。

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

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

序列建模的新视角

在传统理解中,注意力计算通常被视为用注意力分数对value进行加权求和。实际上,我们也可以从“用上下文构建模型”的视角来重新审视注意力机制。

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

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

如图1所示,对于经典的Softmax注意力,我们可以将其理解为:利用当前上下文的key和value构建一个两层MLP模型。其中,MLP第一个线性层的权重是key矩阵的转置,激活函数为Softmax,第二个线性层的权重则是value矩阵。整个注意力计算过程,相当于将query输入到这样一个由key和value构造的两层MLP模型中。类似地,线性注意力可以理解为:用key和value构造了一个线性层权重,再将query输入这个线性模型,从而实现线性注意力计算。

因此,两种经典的注意力范式本质上都是利用key和value构造一个小型模型,然后将query输入该模型以完成建模。

在这个视角下,两者的核心区别在于是否进行压缩。Softmax注意力使用完整的key和value构造两层MLP,不进行压缩,效果出色但计算成本高昂。相反,线性注意力对key和value进行了大幅压缩,通过矩阵乘法将完整的key和value直接压缩到d×d空间。因此,其空间与计算复杂度显著低于未压缩的Softmax注意力。然而,由于d×d空间完全是线性的,且压缩方式过于粗糙,线性注意力会导致明显的性能下降。

那么,一个核心问题随之而来:我们能否在实现压缩的同时,保持模型性能,兼顾高精度与高效性?

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

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

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

其中,ℒ是一个自监督的重建损失函数,例如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. 内部训练设置

观察1: 二阶混合偏导数消失(近似为0或等于0)的损失函数L不适合TTT。

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

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

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

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

若混合导数(近似)为0,则的外部梯度消失,无法优化。图2验证了这一点:MAE(L1)损失的一阶导数是符号函数,混合二阶导数几乎处处为0,效果最差。Smooth L1的混合二阶导数在特定区域为0,同样效果不佳。

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

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

图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。

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

图4:TTT内部训练学习率

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

需要特别指出的是,在部分特殊场景下,内部学习率能够被整合到 key 与 value 的缩放操作中。然而,这并不代表内部学习率的设定无关紧要。一个可供类比的情况是 Softmax 注意力机制中的缩放因子,尽管它可以被 query 和 key 吸收,但其重要性依然不可忽视。

2. 内部模型设计

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

图 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,以此考察内部模型深度对性能的影响。结果表明,随着内部模型深度的增加,模型性能反而持续下降。这在一定程度上与我们的预期相悖,因为理论上更深的内部模型拥有更大的容量,理应能够实现更优的序列建模效果。

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

图 6:较深的 TTT 内部模型在训练与测试阶段的 loss 均更高。

本文认为,优化问题正是导致更深内部模型性能不佳的根本原因。为了验证这一猜想,我们在图 6 中展示了上述三种模型的训练与测试 loss。可以看出,采用更深内部模型时,其训练 loss 反而更高,这表明性能下降并非源于过拟合,而恰恰是欠拟合。更深层的内部模型本应具备更强的拟合能力,但在实际训练中却比浅层模型更欠拟合,这充分说明深层内部模型存在严重的优化困难,远未达到其理论上的表征能力。图 5 第二部分的实验结果进一步印证了这一点:当强制将两层 MLP 或 SwiGLU 的输出层限制为 identity 映射时,模型性能得到了显著提升,这直接表明较深的内部模型存在明显的欠优化问题。

这种优化困难与 TTT 的元学习特性密切相关。一方面,内部模型的初始权重是外部训练过程中的可学习参数,更深的内部模型可能会使这些初始参数难以优化。另一方面,即使获得了较好的初始参数,深层内部模型在内部训练过程中也可能遭遇梯度消失或梯度爆炸等问题,从而影响最终效果。

如何解决较深内部模型的优化问题,是 TTT 未来发展中一个至关重要的方向。理论研究表明,神经网络的拟合能力会随着深度呈指数级增长,这正是当前神经网络取得成功的关键原因。因此,较深的内部模型在实现高精度测试时训练序列建模方面,具有巨大的潜力。

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

在 Transformer 问世之前,卷积操作一直是计算机视觉领域的基石。在 TTT 框架中,我们完全可以将内部模型实现为一个轻量级的卷积网络,而无需将其局限于 MLP、SwiGLU 这类单 token 处理组件。在图 5 的第三部分,我们展示了两种最简单的实现方式:将内部模型分别构建为 3×3 卷积或 Depthwise 卷积。实验结果显示,这两种实现都带来了显著的性能提升。

本文认为,采用卷积网络作为 TTT 内部模型,能够实现局部信息与全局信息的完美融合。具体来说,在 TTT 的内部训练过程中,key 和 value 所包含的全局信息被压缩进了卷积核的权重中。当 query 经过内部模型进行前向传播时,它能够同时自然地获取到局部信息(来自卷积核的感受野)和全局信息(来自蕴含全局信息的卷积核权重),从而完美契合视觉任务的需求。

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

基于上述探索,本文提出了一个纯 TTT 架构——Vision Test-Time Training(ViT³)模型。具体而言,在内部训练方面,采用一次全批量梯度下降,学习率设为 1.0,并使用点积损失函数;在内部模型方面,则采用简化的 SwiGLU 和 Depthwise 卷积。作为 TTT 架构的一个基线模型,ViT³ 可以灵活应用于各类视觉 Transformer 框架,以完成多种不同的任务。

本文在分类、检测、分割、生成等代表性视觉任务上,全面验证了 ViT³ 的性能。实验结果表明,ViT³ 能够超越多种先进的线性复杂度设计,例如线性注意力机制和视觉 Mamba 模型。这一结果不仅验证了 ViT³ 设计的有效性,也凸显了测试时训练模型在高效、可扩展的线性复杂度视觉序列建模领域所展现出的巨大潜力。

以下是部分实验结果,更多详细结果请参考论文原文:

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

图 7:ImageNet-1K 图像分类实验

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

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

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

图 9:图像生成实验

清华团队提出ViT³:用深度学习革新序列建模,线性复杂度超越Transformer

图 10:推理吞吐量与显存占用随输入分辨率的变化

总结与展望

设计高效且具备高表达能力的序列建模范式,始终是计算机视觉领域的核心课题。

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

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

基于上述研究,本文提出了视觉测试时训练模型(ViT³),这是一种专为视觉序列建模设计的纯 TTT 架构。ViT³ 在多项视觉任务上均取得了具有竞争力的结果,为线性复杂度模型提供了一个强有力的基线。

我们期望本文的研究结论与分析能够为后续的视觉 TTT 模型研究工作提供有益的参考与启发。


关注“鲸栖”小程序,掌握最新AI资讯

本文来自网络搜集,不代表鲸林向海立场,如有侵权,联系删除。转载请注明出处:https://www.itsolotime.com/archives/35172

(0)
上一篇 14小时前
下一篇 14小时前

相关推荐

  • Agent原生架构:Claude Code 后时代该如何构建智能体应用

    最近,Claude Code 的流行不仅源于其作为“Vibe编程神器”的体验,更在于它正在重塑智能体的开发范式。过去那种依赖胶水代码或拖拽式构建的、面向过程的传统智能体,正面临被一种全新模式的挑战:这种模式只需开发者描述目标结果,然后交由智能体通过持续循环运行来达成目标。 Claude Code 配合其恰到好处的插件与技能机制证明,一个优秀的编程智能体,本身…

    2026年1月11日
    53800
  • 卡帕西力荐NanoClaw:仅4000行代码的AI执行中枢,开启本地化智能新纪元

    自从OpenClaw爆火后,各种Claw开始轮番登场。 Nano Claw 、Zero Claw 、Pico Claw 刷屏,连卡帕西都坐不住了,为了“抓虾”,他一个百米冲刺奔向苹果店抢Mac Mini,要好好拆解一番爆火的各种Claw们。 店员还奇了怪了,一脸懵地跟卡帕西嘟囔:不知道为啥,这玩意儿最近卖爆了…… 大神顺利购入Mac Mini过后,心满意足一…

    2026年2月22日
    96500
  • Vision Agents:开源框架革新实时视频AI,构建多模态智能体的终极解决方案

    如果你曾尝试构建一个能够“看见”、“听见”并即时“响应”的实时 AI 系统,就会知道其技术栈有多么复杂。 视频需要一个 SDK。 语音需要另一个。 目标检测需要另一个。 大语言模型(LLM)还需要一个。 之后,你仍需将所有组件集成起来,处理延迟问题,并设法让整个系统实时运行。 Vision Agents 改变了这一切。 这是一个开源框架,旨在帮助开发者构建能…

    2025年12月17日
    52300
  • 跨越模态边界:构建真正理解图像、表格与文本的多模态RAG系统

    构建多模态 RAG 系统的终极指南 三个月前,我们新开发的 AI 应用在诸多看似简单的问题上频频“翻车”。问题根源并非 AI 不够智能或数据不足,而是因为答案蕴含在一张图片里,而当时的系统仅能处理文本。 这一时刻迫使我直面一个在构建 RAG 系统时长期回避的核心问题:我们花费数年时间教 AI “阅读”文字,却忽略了人类同样通过图像、表格、公式和流程图来“表达…

    2025年12月16日
    56800
  • DualCamCtrl:几何感知扩散模型革新视频生成,相机运动误差降低40%

    本研究的共同第一作者是来自香港科技大学(广州)EnVision Research 的张鸿飞(研究助理)和陈康豪(博士研究生),两位研究者均师从陈颖聪教授。 你的生成模型真的「懂几何」吗?还是只是在假装对齐相机轨迹? 当前众多视频生成模型虽宣称具备「相机运动控制」能力,但其控制信号通常仅依赖于相机位姿。虽近期工作通过逐像素射线方向(Ray Condition)…

    2025年12月21日
    37800