Improving Test-Time Scaling with Adaptive Looped Transformers
Yichen You, Tianyu Fu, Aosong Feng, Xingtai Lv, Xuefei Ning, Ning Ding, Yu Wang
cs.CL, cs.LG
2026-09-29
清华与Yale提出自适应循环后训练TaH2，只给真正受益的token加深。1.7B上AIME算力斜率从1.79提到2.74（+53%），同等解码FLOPs准确率比非循环基线高出约3.4个百分点。
循环 Transformer 把同一组层反复套在每个 token 上，用共享参数换更深的潜空间计算。Ouro 整栈循环，Huginn 只循环中间块。先前对照多卡在参数量对齐，或每个 token 的 FLOPs 对齐，没回答一个更实用的问题：输出变长、解码算力翻倍时，循环能不能把准确率抬得更快。
清华与 Yale 用同一份 Qwen3-Base 做后训练，把非循环基线 Standard 和后训练版 Ouro、Huginn 放到同一条准确率-算力曲线上。指标叫 accuracy-compute slope：解码 FLOPs 每翻一倍，准确率涨多少个百分点。固定深度 Ouro（M=2）、带退出门的 Ouro、Huginn（M=3）斜率分别是 2.12、2.27、2.26，Standard 只有 1.79。斜率更陡，重叠算力区间里循环模型却全程落后。验证集上，Ouro 从第 1 次迭代到第 2 次，52.3% 的 token 损失变化不超过千分之一，21.7% 变得更差；Huginn 对应比例是 33.5% 和 15.9%。固定深度在给每个 token 加一轮，很多位置并不需要。
TaH2 在后训练里同时学骨干和「这一轮还要不要再跑」的 iteration decider。updater 把原始 embedding 重新注入下一轮输入，和 decider 加起来不到总参数的 3%。每个 token 最多循环 M 次，decider 看 embedding、隐状态和 TopK 概率，吐一个继续概率，低于 0.5 就停。简单位置早停，难位置继续，按 token 加时。
深度标签在线算。当前骨干再跑一轮，停下来的 token 用无梯度 lookahead 补这一轮，看 next-token 损失降了多少。正增益按大小排序，覆盖总正增益的 99% 之后截断，剩下的边际改进当噪声丢掉。标签跟着当前骨干走，和前作 TaH 的分阶段离线 mismatch 标签不同，也和 Ouro 先冻骨干再训退出门不同。训练时 token 走的深度策略和推理一致，避免训满深度、推理早退的错位。
最终预测按各层停止概率加权混合。联合损失是 next-token 交叉熵，加上系数 0.05 的代价敏感二元交叉熵：离截断点越远的决策权越大。把这套权重改成均匀，AIME24–26 平均掉 3.6 点，是所有消融里伤最重的一项。学到的 updater 换成 Top-100 embedding 加权，平均掉 2.1 点；只用最后一轮输出，掉 1.0 点。
1.7B、同一份 273K 蒸馏数据、16K 上下文。十个基准平均如下。
| 方法 | 十基准平均 | 每 token FLOPs |
| Standard（M=1） | 37.7 | 1.00× |
| Huginn（M=3 / 7） | 38.2 / 38.4 | 1.90× / 3.72× |
| Ouro（M=2 / 4） | 37.3 / 36.4 | 1.91× / 3.75× |
| TaH2-fixed（M=2 / 4） | 39.4 / 40.3 | 1.99× / 3.97× |
| TaH2（M=2 / 4 / 8） | 40.6 / 40.9 / 42.5 | 1.21× / 1.47× / 2.37× |
AIME24–26 上，M=2 的 TaH2 斜率 2.74，对 Standard 的 1.79，高 53%。评测拉到 32K token 时，Standard 在每条回复 191.7 TFLOPs 处封顶 12.0%，同一算力下 TaH2 到 15.4%，高出 3.4 点。深度上限从 2 加到 8，相对 Standard 的 AIME 增益从 +2.8 长到 +3.9；Ouro 和 Huginn 基本不再涨。M=2 时平均实际深度只有 1.23，这就是每 token FLOPs 仅 1.21 倍的原因。
单科上，HumanEval 从 63.6 到 72.6（M=8），OlympiadBench 从 40.4 到 47.4。4B / 8B 的 TaH2（M=2）平均分别加 3.2 和 2.4 点，4B 的 AIME26 从 45.5 到 52.4。AIME 上 cons@32 为 27.3% 到 29.3%，Standard 是 21.9%。扩展 MiniSGLang 实测，M=2 每 token FLOPs 高 22%，端到端延迟高 30% 到 34%，吞吐约 0.71 倍，准确率-延迟曲线仍压过 Standard。
GPQA 几乎没拉开：Standard 31.5，TaH2 在 M=8 才 32.0，Ouro 在 M=2 给出 34.0。收益集中在数学和代码，问答上不稳定。
这是给现成 dense LLM 加循环深度的后训练方案，不用从零预训练原生循环模型。M=2 是更可落地的工作点：多数 token 只跑一轮出头，额外延迟三成左右，AIME 斜率已经拉开。再往 M=8 推还能涨点准确率，训练 FLOPs 会到 Standard 的 5.17 倍。
固定循环在后训练设定里经常既贵又不如不循环。浪费发生在 token 粒度，解法也得按 token 分配深度。它仍是 SFT 上的渐进改进，没有碰强化学习。4B 到 8B 平均增益从 3.2 收到 2.4，更大模型还没有数据。
训练比普通 SFT 贵。1.7B、M=2 是 Standard 的 1.87 倍，lookahead 标签占训练 FLOPs 的 20% 到 24%；M=8 到 5.17 倍。方法只在 SFT 下验证，在线蒸馏和强化学习留作后续。
对照里的 Ouro、Huginn 是按论文设定改到后训练的再实现，不是官方预训练 checkpoint。结论覆盖的是 dense 基座事后加循环，不能直接外推到原生循环预训练。主数字来自 30 题规模的 AIME avg@32，方差不会小。decider 阈值钉在 0.5，没有扫 serving 侧的延迟-准确率边界。自适应深度还要改推理引擎才能在同一次前向里混不同深度，普通 vLLM 路径没有报。