复用 MLP 做 fast weights,字节让 4B 模型免从头重训学会测试时训练

In-Place Test-Time Training

Guhao Feng, Shengjie Luo, Kai Hua, Ge Zhang, Di He, Wenhao Huang, Tianle Cai

ICLR 2026 Oral Presentation

cs.LG, cs.AI, cs.CL, stat.ML

2026-04-08

把 gated MLP 输出投影当 fast weights 在推理时按 chunk 在线更新,目标对齐下一 token 预测,Qwen3-4B 的 RULER 128k 从 74.8 升到 77.0,256k 外推同步占优。

这篇在解决什么

LLM 的常规路径是先训练、后部署,权重到了推理期就冻死。新来的信息只能靠 attention 塞进上下文窗口,窗口一长,二次复杂度和位置外推都是麻烦。Test-Time Training(测试时训练)指了另一条路:挑一小块参数当 fast weights,推理时随输入在线更新,把上下文压进权重。但现有 TTT 方案有三个硬伤:通常做成替换 attention 的专用层,想用就得从头预训练;逐 token 的串行更新喂不饱 GPU;更新目标沿用「重建当前 token」的自监督,和下一 token 预测这个真正的目标没有对齐。这篇出自字节跳动 Seed 与北大(ICLR 2026 Oral),主张不换任何零件,把 TTT 塞回现成的 Transformer。

方法

三个设计对应三个硬伤。

fast weights 放在哪:不加新层,直接把 gated MLP 的输出投影矩阵 Wdown 当 fast weights,推理时原位更新;Wup、Wgate 和全部 attention 冻结。MLP 本来就被看作存通用知识的 key-value 记忆,现在让同一块矩阵顺带在线吸收上下文。架构没动,预训练权重完整保留,可以从 checkpoint 热启动。
按 chunk 更新:attention 原样保留,MLP 的更新就不必逐 token 做。序列切成 chunk,每个 chunk 先用当前 Wdown 算输出(apply),再用该 chunk 的激活 Z 和目标 V 做一步梯度下降(update)。更新量具有结合律,可以用 parallel scan 加前缀和在 context parallelism 下并行计算,数学上严格等价于串行的因果更新。消融显示 chunk 取 512 到 1024 最好,摆脱了独立 TTT 层那种小 chunk 依赖。
换目标:不再要求 fast weights 重建当前 token 的表示,目标 V 改为携带未来 token 信息,由词嵌入过一维卷积再乘可学习投影得到。理论部分在 induction head 设定下证明:用「下一个 token」当目标,一次更新能让正确 token 的 logit 期望提升 λ·cnorm²·calign,其他 token 几乎不动;重建目标推不动正确 logit。配合内积损失,更新有闭式解 W ← W + ηVZᵀ,一步更新没有反向传播开销。

结果

Drop-in 实验:Qwen3-4B-Base 走同一条继续预训练课程(约 20B token@32k 加 15B token@128k,配 YaRN),唯一变量是是否启用 In-Place TTT。RULER 平均准确率:

设置	Baseline	In-Place TTT
RULER 64k	74.3	78.7
RULER 128k	74.8	77.0
RULER 256k(超出训练长度)	41.7	43.9

4k 上基线还略高(96.6 对 96.1),8k 起反超并随长度扩大差距。同一张表里 Llama3-8B 在 128k 只有 69.5,Phi3-medium-14B 只有 46.1;不过不开 TTT 的基线也有 74.8,跨模型的优势主要来自底子和继续训练,TTT 在 128k 的净贡献是 +2.2。LLaMA-3.1-8B 和 Qwen3-14B 上结论一致(64k 分别 +2.1、+2.7),且与 YaRN 叠加不冲突。

从头训练(500M/1.5B,32k 上下文,对照 SWA、GLA、DeltaNet、LaCT):滑窗困惑度全程低于所有基线,并随上下文到 32k 持续下降。4B 规模(120B token、8k 训练长度)收益最大:Full Attention 模型 RULER-16k 从 6.58 到 19.99,SWA 模型 RULER-8k 从 9.91 到 26.80。常识推理基本不动(HellaSwag 55.67 到 55.85)。prefill 吞吐与峰值显存的开销可忽略(图 4)。

消融(1.7B 模型):启用越多层的 MLP(状态越大)效果越好;chunk 512/1024 最优;去掉一维卷积或投影都掉点,卷积管长上下文,投影管短上下文。

为什么重要

TTT 一直是想法好、落地要推倒重来。这篇证明可以用继续预训练的代价给现有 checkpoint 加上在线适应,不动 attention、不改架构。对做长上下文的团队,这条路线与 RoPE 外推、稀疏注意力正交,可以叠加。MLP 占 Transformer 参数大头,把它当可扩容的记忆状态,等于给「记忆大小随模型规模增长」提供了一个天然 scaling 轴。

局限与存疑

Drop-in 不等于免费:Qwen3-4B 的实验花了约 35B token 的继续预训练,比从头训练便宜,离即插即用还远。RULER 涨幅是几个点量级,256k 外推的 43.9 本身远未可用。理论建立在 induction head 玩具设定加两条假设(词嵌入近似正交、key-query 对齐)上,解释的是方向而非最终性能。常识任务几乎没动,收益集中在长上下文。作者自己留的口子:损失函数与优化器的选择未探索,与线性注意力、SSM 骨干结合是未来工作。文中也没有与「把 TTT 层嫁接到预训练模型」等其他集成方式的直接对照,drop-in 的优势是论证出来的,不是比出来的。

术语

TTT(Test-Time Training):推理阶段用自监督目标在线更新一小块参数,让模型适应当前输入。
fast weights:推理时快速更新的参数子集,充当随上下文演化的记忆;其余冻结参数叫 slow weights。
RULER:长上下文综合基准,在不同长度上测检索与推理任务的平均准确率。
YaRN:一种 RoPE 位置编码缩放方法,把有效上下文拉长到预训练长度之外。
Context Parallelism:把序列沿长度切分到多卡并行的训练与推理策略。
滑窗困惑度:固定最后一段 token,看前文变长时这段的困惑度是否下降,衡量模型是否真用上了上下文。

原文与代码

社区讨论

LLM 的下一个 scaling 轴是「睡觉」?解读睡眠期计算与记忆巩固研究 bycloud