把 Test-Time Training 拆成可组合模块,系统消融发现浅层线性才是最强记忆

Modular TTT: Rethinking Test-Time Training as Composable Modules

Bohao Tang, Zhen Qin, Yuqi Pan, Zheng Li, Pengfei Liu, Ya Zhang

cs.LG, cs.CL

2026-08-07

把 TTT 内部学习器表示成有向无环图,各组件独立变化。消融发现小学习率、衰减、单层非线性有用,深层和归一化有害。1.45B 与 Gated DeltaNet 持平。

这篇在解决什么

TTT(test-time training)把序列建模看成在线学习:模型处理序列时,内部的一组「快权重」被一条学习规则持续更新。问题是 TTT 变体越来越多,每个都硬编码一套实现,想改一个组件(快权重网络、损失、学习率、权重衰减、归一化)常常牵一发动全身,既难设计新变体,也看不清每个组件到底贡献多少。

方法

Modular TTT 把内部学习器表示成有向无环图(DAG):节点是已注册的原语(线性层、非线性、归一化、损失),边是张量依赖。框架自动把「训练视角的前向和反向」与「因果查询视角」的规则组合成完整的图级 TTT 计算,包括快权重状态转移。换组件只换节点或子图,所以每个维度能单独变化。来自字节跳动 Seed 和上海交大。

结果

系统消融(在 160M 和 410M 规模)给出几条清晰结论:

损失函数:MSE 和内积损失表现接近(都好),其他明显差。
小学习率初始化有用。
权重衰减(标量或向量)有用。
单层非线性(线性层后接 SiLU 或 GELU)有用,它对写入做逐坐标门控,导数有界、不放大激活。
深层快权重网络(MLP)和归一化反而有害,会诱导过大的激活;残差连接和门控基本没收益。

按这些结论训 410M 和 1.45B(100B token):1.45B 训练损失 2.3150、多选题平均 56.4%,与 Gated DeltaNet(2.3042 / 56.3%)持平,和 LLaMA(2.3041)也接近。训练吞吐是官方 TTT 实现的 2.2 到 3.3 倍。

为什么重要

价值有两层。工程上,DAG 加自动组合让「造一个新 TTT 变体并和旧的公平对比」从手搓变成搭积木,配套的解析反向算子带来 2 到 3 倍吞吐。认知上,它给 TTT 的设计空间画了一张较清楚的地图:该往「浅、加一层非线性、带衰减」走,别往「深、带归一化」走。论文还给了深层快权重难优化的数学解释:同一个等效快权重的不同因子化,会诱导差异很大的更新方向。

局限与存疑

规模和上下文是短板。RULER 长上下文评测里 LLaMA 仍明显更强(尤其 8k),说明当前浅层 Modular TTT 在精确长程召回上还有差距;containment 类任务也偏弱。1.45B、100B token 按今天的标准不算大,能否上到更大规模仍保持持平,没有验证。「与 GDN 持平」是和另一个线性注意力或递归基线打平,不是超越。

术语

TTT(test-time training):把序列建模看成在线学习,模型在处理序列时内部更新一组快权重。
快权重(fast weights):随输入快速更新的参数,充当记忆。
DAG(有向无环图):节点为算子、边为数据依赖的计算图结构。
Gated DeltaNet(GDN):一类带门控的线性注意力或递归模型,本文的主要对比线。
消融(ablation):逐个去掉或改动组件看性能变化,用来定位各组件的贡献。

原文与代码

社区讨论

拆解测试时训练(TTT):消融实验揭示最优组件配置 burny_tech