用自适应循环Transformer改进测试时扩展

联合训练主干与迭代决策器，把额外迭代精准分配给受益token，让循环模型测试时扩展超越非循环基线

前置知识

测试时扩展（Test-time Scaling）

在推理阶段投入更多算力以换取更高准确率的范式，包括生成更长的思维链、多次采样投票、树搜索，以及在潜空间重复应用共享层。本文用“精度–算力斜率”来度量：每把解码FLOPs翻一倍，准确率能提升多少点，通过对不同输出token截断下的准确率关于 $\log_2$(解码FLOPs) 做线性拟合得到。

本文的核心问题就是“循环结构能否改善测试时扩展曲线”，全文所有结论都建立在这个斜率指标上，不理解它就无法读懂实验。

循环Transformer（Looped Transformer）

让同一个（或一组）Transformer层被重复执行多次的架构：隐藏状态 $h_t^{(m)} = F_\theta(h_t^{(m-1)})$ 迭代更新后再输出。参数共享使得有效深度可以超过物理层数，用“潜空间计算”代替生成更多token。代表工作有Universal Transformer、Huginn（循环中间块）和Ouro（循环全部层）。

这是本文的研究对象：作者从预训练模型出发通过后训练引入循环，并比较不同循环范围（全层 vs 中间块）的扩展行为。

自适应计算与早退（Early Exit）

不同输入所需的计算量不同，自适应计算让每个token动态决定执行多少层/多少次迭代。常见做法包括退出门控（Ouro的exit gate）、路由器（MoR）和带置信度惩罚的pondering（PonderNet类）。决策通常由一个小网络输出继续概率 $g$，低于阈值即停止。

TaH2的本质就是一个token级的深度决策器，理解早退机制才能看清它与Ouro门控、TaH离线标签的差异。

KV缓存与duo-causal注意力

自回归解码时每个token的Key/Value会被缓存以避免重复计算。循环模型中每次迭代可维护独立的KV流；duo-causal规则规定位置 $t$、深度 $m$ 的查询只能关注位置 $s\le t$、深度 $j\le m$ 的状态。把各迭代KV拼接后，该规则变成块状掩码，训练时可并行。

TaH2的“扩展duo-causal注意力”是支撑训练时前瞻监督的关键工程基础，前瞻迭代就是靠掩码规则与梯度截断并行实现的。

后训练/SFT与蒸馏

在预训练模型上用监督数据继续训练。本文用教师模型（Qwen3-8B或235B）重新生成数学/代码/问答/工具使用响应，再对学生模型做SFT。数据质量直接影响学生上限——论文发现8B教师蒸馏出的学生反而比235B教师高1.8点。

所有对比模型（Standard、Ouro、Huginn、TaH2）都用同一检查点和同一SFT配方训练，这是“控制变量比较架构”这一实验设计的前提。

交叉熵损失与BCE

下一token预测损失 $\ell_t^{(m)}=-\log q_t^{(m)}[y_t^*]$ 度量预测分布与真实token的差异；二元交叉熵 $\mathrm{BCE}(g,c)$ 监督决策器的继续概率 $g$ 逼近0/1标签 $c$。本文用实测损失下降量 $\delta$ 生成标签，并用 $|\delta-\delta_{\mathrm{cut}}|$ 加权BCE。

深度标签完全由逐token交叉熵的变化推导而来，看不懂 $\ell$ 和 $\delta$ 的定义就看不懂TaH2的训练目标。

研究动机

循环Transformer的既有评测口径与测试时扩展脱节：先前工作只在参数量对齐或单token FLOPs对齐的条件下比较循环与非循环模型，无法回答“随着输出token预算增大，循环是否带来更优的精度–算力曲线”。作者用同一Qwen3-1.7B-Base检查点、同一后训练数据对比Standard（非循环）、Ouro（全层循环，M=2，固定深度或自适应退出门）与Huginn（中间14层循环，M=3），扫描4K–16K输出截断并拟合精度对 $\log_2$ 解码FLOPs的直线，发现固定Ouro、自适应Ouro和Huginn的斜率分别为2.12、2.27、2.26点/倍增，均陡于Standard的1.79，但在重叠算力区间内精度全部更低。逐token分析（图3）揭示了根因：从第1次到末次迭代，Ouro有52.3%的token损失变化不超过 $10^{-3}$、21.7%的token损失反而恶化；Huginn对应为33.5%和15.9%。固定深度循环把额外算力花在大量零收益甚至负收益的token上。

本文的目标是本文的目标是：不从头预训练，仅通过后训练就让循环LLM在相同测试时算力下超越非循环基线。具体拆解为三点：(1) 让模型学会判断哪些token能从额外迭代中受益，把深度算力按token动态分配，而非对每个token无差别执行M次迭代；(2) 深度决策必须与主干联合训练，监督信号来自当前模型实际测得的迭代增益，避免离线标签与当前策略失配；(3) 同时改进两条扩展曲线——精度随解码FLOPs的斜率（test-time scaling）和精度随最大迭代深度M的增长（depth scaling），并在4B/8B规模以及数学之外的代码、问答、工具使用域上验证泛化性。

与已有工作不同的是，本文的独特切入是“在线前瞻深度监督 + 主干与决策器联合训练”。既有自适应循环方法各有结构性缺陷：Ouro分两阶段，先在全深度下训练骨干、再冻结骨干训练退出门，导致骨干从未见过早退状态的训练–推理失配；TaH用分阶段训练加离线top-1失配标签（当前预测错就标“继续”），但标签只反映“现在对不对”，不反映“再迭代一次是否真能改进”；MoR和Ponder类方法靠容量约束或置信度惩罚间接学深度。TaH2则在每个迭代点直接度量预测损失变化 $\delta_t^{(m)}=\ell_t^{(m)}-\ell_t^{(m+1)}$，在线生成“继续迭代是否有收益”的标签，并用排序截断过滤边际增益噪声——深度策略与骨干在同一损失下共同演化，训练时token就走当前决策器选出的路径（on-policy），把“该不该多想一步”从启发式问题变成有直接因果度量的监督学习问题。

核心方法

直觉起点：一句话中各token所需的计算深度并不相同——论文的token级可视化（图11）显示数学表达式和最终代码用的迭代次数明显少于其前面的自然语言推理。整体路线是在预训练Qwen3-Base上插入三个轻量组件构成TaH2：共享主干 $F_\theta$、学习型输入注入更新器 $U_\psi$、token级迭代决策器 $D_\phi$，新增参数不足3%（1.7B为46.16M，占2.61%）。每次迭代对所有 $L$ 层执行扩展duo-causal注意力（token $t$ 在深度 $m$ 关注 $s\le t$、$j\le m$ 的KV状态）；首次迭代输入词嵌入 $e_t$，其后按 $h_t^{(m+1)} = F_\theta\big(U_\psi(e_t, h_t^{(m)})\big)$ 更新；决策器拼接 $[e_t; h_t^{(m)}; \mathrm{TopK}(q_t^{(m)})]$ 输出继续概率 $g_t^{(m)}\in(0,1)$，$g<\tau_{\mathrm{exit}}=0.5$ 即停止；各深度预测按停止概率加权混合 $q_t=\sum_m \omega_t^{(m)} q_t^{(m)}$。训练采用前瞻深度监督：用无梯度前瞻迭代实测增益 $\delta_t^{(m)}$ 在线生成继续标签，与下一token预测损失联合优化。

核心创新是“用实测迭代增益在线生成深度标签，并与骨干联合训练”，这改变了深度决策的监督来源。与TaH的本质区别：TaH的标签是离线top-1失配（预测错→继续），无法区分“错了但再迭代也没用”与“再迭代一次就能修正”，且分阶段训练使标签与当前策略脱节；TaH2的标签来自当前骨干在当前决策路径上的真实损失下降 $\delta$，直接回答“再想一步值不值”，且随骨干更新在线再生。与Ouro的本质区别：Ouro骨干只在全深度下训练、门控后训练且冻结骨干；TaH2训练时token按决策器的确定性阈值决策选择深度（on-policy），骨干从第一步就适应早退后的状态分布。另有两个配套设计：代价敏感加权（离截止 $\delta_{\mathrm{cut}}$ 越远的样本BCE权重越大）和覆盖率 $\rho=0.99$ 的增益截断（过滤边际增益噪声），消融显示均匀加权会损失3.6点、$\rho=1$ 损失1.9点，证明二者缺一不可。

方法步骤详情

完整流程分四步。第一步，架构搭建：从Qwen3-{1.7B,4B,8B}-Base出发，插入更新器（RMSNorm+拼接投影 $W_u$+SwiGLU MLP，融合 $e_t$ 与 $h_t^{(m)}$）和决策器（同构MLP+sigmoid打分，接收Top-2048预测概率），并将注意力改为扩展duo-causal：每次迭代维护独立KV流，查询可见 $s\le t$、$j\le m$ 的状态；训练时对停止token追加无梯度前瞻迭代，前瞻查询不可见其他token的前瞻状态。第二步，数据与配置：1.7B用Qwen3-8B重新生成的273K条数学/代码/问答响应加Nemotron工具使用样本（共3.4B token，3 epoch，16K上下文，batch 128，学习率 $4\times10^{-5}$，bf16）；4B/8B用Qwen3-235B-A22B原始响应，固定TPP=2 token/参数。第三步，训练循环：token按当前决策器阈值 $g<0.5$ 走深度路径；对执行了第 $m$ 次迭代的token计算 $\ell_t^{(m)}=-\log q_t^{(m)}[y_t^*]$，停止token由无梯度前瞻给出 $\ell_t^{(m+1)}$；增益 $\delta_t^{(m)}=\ell_t^{(m)}-\ell_t^{(m+1)}$，正增益按降序累积至总量的 $\rho=0.99$ 得截止 $\delta_{\mathrm{cut}}^{(m)}$，$\delta\ge\delta_{\mathrm{cut}}$ 标“继续”$c=1$；总损失 $\mathcal{L}_{\mathrm{SFT}}=\frac{1}{N}\sum_t[-\log q_t[y_t^*]]+\alpha_D\frac{1}{N}\sum_{t,m}w_t^{(m)}\mathrm{BCE}(g_t^{(m)},c_t^{(m)})$，其中 $\alpha_D=0.05$，$w$ 随 $|\delta-\delta_{\mathrm{cut}}|$ 增大。第四步，推理与服务：token在首次 $g<\tau_{\mathrm{exit}}=0.5$ 时停止、否则跑满 $M$；用扩展的Mini-SGLang做连续深度批处理，不同深度请求合入同一前向。

技术新颖性

技术新颖性有四点。(1) 评测视角创新：首次在“后训练引入循环”的设定下，用精度–算力斜率系统刻画循环模型的测试时扩展，揭示了“斜率更陡但绝对精度更低”这一此前被匹配参数/FLOPs口径掩盖的悖论，并给出逐token损失增益的机制解释。(2) 监督信号创新：深度标签建立在实测的逐token损失增益上，随骨干更新在线再生，克服了TaH离线失配标签与Ouro冻结骨干门控的缺陷；前瞻通过扩展duo-causal掩码在训练中并行完成，前瞻项仅占训练FLOPs的20–24%。(3) 架构组合的贡献分解：学习型更新器（对比TaH的top-100嵌入加权好2.1点）与停止概率加权混合输出（对比只用末次迭代预测好1.0点）各有独立、可消融验证的贡献。(4) 工程验证完整：扩展Mini-SGLang实现连续深度批处理，给出真实的端到端延迟–精度曲线（图5）而非仅理论FLOPs，并以表9–10公开了逐组件的FLOPs记账方法。

**Figure 2**: TaH2's architecture and training scheme. Left: the decider decides whether to continue after each iteration. Right: the updater reinjects token embeddings at each iteration, and stopping probabilities weight per-iteration predictions.

**Figure 7**: Extended duo-causal attention in TaH2. Each cell denotes a (token id, iteration depth) pair; green arrows and cells show the keys visible to token 3.

实验结果

实验覆盖三个规模、十个基准，主要发现如下。第一，测试时扩展（图1a/图4b）：1.7B上TaH2(M=2)的斜率为2.74点/倍增解码FLOPs，比Standard的1.79高53%，也高于其他循环模型的2.12–2.42；评测延长到32K后，Standard在191.7 TFLOPs/响应处饱和于12.0%，TaH2在相同算力下达15.4%（+3.4点）。第二，深度扩展（图1b/图4a）：提高上限 $M$ 时Ouro/Huginn基本平台化，而TaH2对Standard的AIME领先从 $M=2$ 的+2.8点增至 $M=8$ 的+3.9点；验证损失上自适应TaH2在 $M=8$ 时比Standard低0.0106。第三，十基准平均（表1，1.7B）：TaH2在 $M=2/4/8$ 分别得40.6/40.9/42.5，对Standard的37.7提升+2.9/+3.2/+4.8点，而单token解码FLOPs仅为1.21×/1.47×/2.37×；分项上HumanEval 72.6 vs 63.6（+9.0）、OlympiadBench 47.4 vs 40.4（+7.0）。第四，规模泛化（表3）：4B平均49.6 vs 46.4（+3.2，AIME26提升6.9点），8B平均58.5 vs 56.1（+2.4，AIME24提升4.4点）。第五，并行扩展：AIME24–26的cons@32达27.3–29.3%（M=2–8），显著高于Standard的21.9%。第六，真实服务开销（表2/图5）：TaH2每token解码14.09 GFLOPs（2.00×）、延迟2.35×、吞吐0.44×，但在相同延迟预算下精度仍高于Standard。第七，机制验证（图6）：决策器继续概率与实测损失增益单调对齐，$p\approx0$ 的token增益为负，阈值0.5处增益接近零。

**Table 1**: Accuracy (%) of Qwen3-1.7B models across ten benchmarks.

**Table 2**: Runtime efficiency on AIME26. Ratios are to Std. at the same batch size.

**Table 3**: Accuracy (%) at 4B and 8B. Subscripts give gains over Standard in points.

**Table 4**: Design choices at 1.7B. Each row varies one aspect from TaH2 (default in gray).

**Table 6**: Architectures used in the 1.7B comparison.

**Table 7**: Updater and decider parameters across backbone scales.

**Table 10**: Training FLOPs (10^18) of the 1.7B post-training runs, split into the terms of Equation 21.

**Table 11**: AIME accuracy (%) for Qwen3-1.7B-Base students trained on responses from different teachers.

**Table 12**: Effects of changing the inference depth ceiling.

**Figure 1**: TaH2 improves both test-time and iteration-depth scaling. Mean AIME24–26 accuracy over 32 samples for 1.7B models post-trained from the same checkpoint.

**Figure 4**: Depth and test-time scaling at 1.7B. (a) Final validation loss versus iteration depth. (b) Mean AIME24–26 accuracy versus decoding FLOPs.

**Figure 5**: Accuracy versus mean end-to-end latency.

**Figure 6**: Mean loss reduction from a second iteration versus the decider's continue probability for TaH2 (M = 2) on the validation set.

**Figure 9**: Per-benchmark test-time scaling on the six math benchmarks at 1.7B.

**Figure 10**: Parallel test-time scaling: mean AIME24–26 cons@n versus decoding FLOPs per problem.

**Figure 11**: Token-level iteration depth in sampled correct responses from OlympiadBench, HumanEval and GPQA.

查看结构化数据

任务	指标	本文	基线	提升
AIME24–26 测试时扩展（精度–算力斜率）	每倍增解码FLOPs的精度增益（点）	TaH2 (M=2)：2.74	Standard：1.79	+53%（0.95点/倍增）
AIME24–26 匹配算力精度（32K评测上限）	avg@32 (%)	TaH2：15.4% @ 191.7 TFLOPs/响应	Standard：12.0%（已饱和）	+3.4点
十个基准平均（1.7B，M=8）	平均准确率 (%)	TaH2：42.5	Standard：37.7	+4.8点（单token FLOPs 2.37×）
AIME24–26 深度扩展（M=2→8）	对Standard的领先（点）	M=8：+3.9	Standard（M=1，无迭代）	+2.8→+3.9（Ouro/Huginn则平台化）
十基准平均（4B，M=2）	平均准确率 (%)	TaH2：49.6	Standard：46.4	+3.2点（AIME26单项+6.9点）
十基准平均（8B，M=2）	平均准确率 (%)	TaH2：58.5	Standard：56.1	+2.4点（AIME24单项+4.4点）
AIME24–26 多数投票（并行扩展）	cons@32 (%)	27.3–29.3（M=2–8）	Standard：21.9	+5.4~7.4点
HumanEval（1.7B，M=8）	avg@8 (%)	72.6	Standard：63.6	+9.0点
OlympiadBench（1.7B，M=8）	avg@4 (%)	47.4	Standard：40.4	+7.0点

局限与改进

作者承认的局限有两条：(1) 训练FLOPs高于标准SFT——TaH2(M=2)为 $78.1\times10^{18}$，是Standard（41.8）的1.87倍，其中前瞻标签占20–24%（不过仍低于同上限的TaH2-fixed 2.01×和Ouro 2.00×）；(2) 方法只在SFT/蒸馏设定下研究，未扩展到on-policy蒸馏和强化学习。我补充的观察：(1) 推理延迟开销显著，端到端延迟2.19–2.35×、吞吐降至0.44×，对延迟敏感场景不友好；(2) 收益随规模递减（4B +3.2点 vs 8B +2.4点），在与更强基座的RL推理模型组合时优势可能进一步稀释；(3) 1.7B的绝对精度仍低（AIME均值约15–16%），增益是在弱基线上的相对改进；(4) 深度标签以单步next-token损失为代理，与长链推理的最终正确性未必一致；(5) 退出阈值 $\tau_{\mathrm{exit}}=0.5$ 固定，缺少算力–精度调节旋钮，且表12显示推理上限超过训练上限（M=8→12）无收益，深度泛化能力有限。

独立分析的弱点

独立分析的弱点及改进方向：(1) 串行迭代是延迟的本质瓶颈——即使自适应决策，$M=2$ 时平均深度也达1.23，延迟仍翻倍以上；可结合迭代间的投机解码或更激进的跨请求深度并行调度来压缩延迟。(2) 前瞻监督的训练成本——每个停止token需要一次无梯度的完整前向，占训练FLOPs的20–24%；可用轻量价值头近似 $\delta$、或只对部分token抽样前瞻来降低开销。(3) 标签代理偏差——单步损失下降不等于最终答案更正确，长数学推理中可能出现“单步损失降了但推理链走偏”；可用结果级奖励（RL的return差分）替代或校正token级 $\delta$。(4) 阈值僵硬——$\tau_{\mathrm{exit}}=0.5$ 无法按算力预算动态调节；可训练多阈值或将 $\tau$ 作为显式预算参数（类似budget forcing），实现单模型多档算力。(5) 与现代推理训练栈脱节——实验止步于SFT蒸馏，而当前强推理模型都经RL训练；若在GRPO等框架下让深度决策直接接受结果奖励，斜率优势可能被放大，这是最值得优先验证的方向。

未来方向

作者明确提出的新方向：把TaH2扩展到on-policy蒸馏与强化学习。基于本文成果可以自然延伸的：(1) RL+深度监督——在可验证奖励的RL中，把 $\delta$ 标签换成“多迭代一次对结果奖励的差分”，学习真正面向最终正确性的深度策略；(2) 与token轴扩展协同——潜空间迭代与长思维链、多数投票正交，cons@32已达27.3–29.3%，两轴乘性组合的上限值得系统测绘；(3) 动态退出阈值与算力可控推理——让 $\tau_{\mathrm{exit}}$ 成为用户可调的预算参数，一个模型覆盖多档算力–精度工作点；(4) 深度×宽度联合自适应——与MoE结合，让路由同时决定“用哪几个专家、迭代几轮”；(5) 机制与系统协同设计——图12/13显示不同注意力头分别偏好第一迭代、后迭代或混合的KV状态，可据此做迭代级KV压缩、跨迭代缓存共享；(6) 从后训练走向预训练——联合训练框架天然适配预训练，可能消除后训练引入循环的残余失配，验证循环模型的预训练scaling law。

复现评估

复现条件较好。所有基座（Qwen3-{1.7B,4B,8B}-Base）与数据集（AM-Qwen3-Distilled、Nemotron-Agentic-v1）公开；附录A/D给出完整损失定义、架构细节、FLOPs记账、超参数与评测协议：1.7B用273K条Qwen3-8B重生成响应、3.4B token、3 epoch、batch 128、16K上下文、学习率 $4\times10^{-5}$、bf16，训练约 $78\times10^{18}$ FLOPs（Standard的1.87倍），单节点多卡可承受；服务评测仅需单张A800-80GB。代码、配置与checkpoint承诺发表后放出，但截至论文版本尚未公开，需跟踪作者仓库。主要复现门槛有三：(1) 扩展duo-causal注意力（块状掩码+前瞻隔离）与连续深度批处理需要魔改推理引擎（Mini-SGLang）；(2) 前瞻标签的梯度截断与停止权重实现细节较多，容易出错；(3) 4B/8B规模训练需要较大算力（TPP=2，分别8B/16B token）。总体评估：1.7B级别的核心结论可在学术集群上以中低难度复现验证。

用自适应循环Transformer改进测试时扩展 Improving Test-Time Scaling with Adaptive Looped Transformers