循环Transformer中的递归推理步长调度 Scheduling Recursive Reasoning in Looped Transformers
追踪进展-波动平衡，在线自适应调度循环推理步长，免重训提精度并最高加速1.56倍
前置知识
循环推理 / Looped Transformer
一类通过反复施加同一变换来精炼潜在表示的模型：编码器生成初始状态 $X_0=E_\theta(u)$，之后每步执行 $X_{k+1}=\Phi_\theta(X_k;u)=X_k+\Delta_\theta(X_k;u)$。参数在循环间共享，使有效推理深度与参数量解耦，代表有 TRM、HRM、Ouro、Huginn 等。
本文的全部讨论对象就是这类模型的标准推理循环，TAPS 只是把其中的隐含单位步长 $\eta=1$ 换成可调的 $\eta_k$。
测试时计算扩展
在推理阶段投入更多计算换取更高性能的思路，典型手段包括增加循环次数、并行采样、自适应提前退出等。其收益取决于每次额外计算是否真正改善输出质量。
论文指出该领域已有工作集中在"重复什么计算"与"重复多少次"两个维度，而本文把"每步走多远"立为第三维度，直接决定测试时扩展的边际收益。
Adam 的一阶/二阶矩估计
Adam 优化器用指数滑动平均（EMA）维护梯度的首矩 $M_k$ 与未中心化二阶矩 $v_k$：$M_k=\beta M_{k-1}+(1-\beta)g_k$。首矩反映更新的一致方向（持久进展），$v_k-\|M_k\|^2$ 反映更新的方差（波动）。
TAPS (Adam) 变体直接借用这套统计量来估计轨迹的进展能量 $\hat P_k$ 与波动能量 $\hat S_k$，是方法的核心构件。
雅可比矩阵与伴随反向传播
状态雅可比 $D_j=D_X F(X_j,\eta_j)=I+\eta_j D_X\Delta_\theta(X_j;u)$ 描述第 $j$ 循环对状态扰动的线性传播。把终端损失梯度 $G_K$ 用其伴随 $D_j^*$ 反向传播得到 $G_j$，即可求出损失对任意 $\eta_k$ 的灵敏度。
论文的 Proposition 1（$\partial L/\partial\eta_k=\langle G_{k+1},\Delta_k\rangle_F$）和全部理论结果都建立在这一灵敏度分析之上，理解它才能读懂分解定理。
Barzilai–Borwein (BB) 步长
无导数步长方法：利用相邻两步的位移与梯度差之比估计局部曲率尺度，从而自动选择合适步长，是优化中利用轨迹信息做步长自适应的经典方案。
TAPS 的设计灵感明确来自 BB、Malitsky–Mishchenko 等轨迹式步长自适应；TAPS (BB) 变体在多个实验中表现最佳，是理解实验表的关键。
研究动机
循环推理模型（TRM、Ouro、Huginn 等）通过共享参数反复精炼潜在状态来扩展测试时计算，但标准推理对所有循环一律采用固定单位尺度，即 $X_{k+1}=X_k+\Delta_\theta(X_k;u)$，等价于隐含地令松弛因子 $\eta=1$。这种一刀切在更新持续指向同一方向（进展型轨迹）时过于保守，在更新来回震荡（波动型轨迹）时又过于激进，导致额外循环的收益递减。论文附录 B.2 给出反例说明常数更新并不令人满意；实验也强烈佐证：对 Ouro-1.4B 施加固定缩放 $\eta=1.2$ 会使 MMLU/ARC-C/HellaSwag/GSM8K 四基准平均精度从 61.10 崩至 32.85，在 Sudoku 上准确率从 89.67% 跌到 85.43% 且达不到基线速度。理想的步长应由其对终端任务损失的影响决定，但该信号需要标签 $y$ 和下游梯度，推理时不可得。
本文的目标是本文的目标是把"每次更新走多远"确立为与"重复什么计算"（架构 $\Delta_\theta$）、"重复多少次"（深度/循环数 $K$）并列的第三个控制维度，并构造一个因果自适应调度策略 $\pi$：只依赖输入 $u$ 和截至第 $k$ 循环的轨迹信息，不接触标签 $y$ 与未来状态，在线选择松弛因子 $\eta_k\in[\eta_{min},\eta_{max}]$（$\eta_{min}\le1\le\eta_{max}$），使得在固定预算 $K$ 下期望终端损失满足 $J_K(\pi)0$，达到目标质量所需循环数更少，即 $K_\epsilon(\pi)<K_\epsilon(\pi_0)$。同时希望该机制无需重训即可即插即用于各类递归架构与推理策略。
与已有工作不同的是，已有工作集中于前两个轴：架构方法设计学到的更新 $\Delta_\theta$（Universal Transformer、HRM/TRM、MoE 递归、全模型递归等），自适应深度方法决定迭代多少次（ACT、自适应退出、定点停止）。而更新尺度 $\eta_k$ 基本未被作为独立研究对象，仅以稳定性手段出现——固定缩小步长或定点法的残差阻尼，目标是收敛而非有限预算下的任务性能。本文的独特切入是对终端损失关于调度的灵敏度做时间平均的精确分解：$\bar A_w^k=P_w^k-S_w^k$，把不可观测的任务相关量转化为可由轨迹估计的持久进展能量与中心化波动能量，再借鉴 Barzilai–Borwein、Malitsky–Mishchenko 等优化中的轨迹式步长自适应思想，首次架起"轨迹统计 → 任务改进方向 → 循环数减少"的理论桥梁。
核心方法
直觉：递归更新持续进展时应放大步长，波动占主导时应阻尼。分三步。其一，灵敏度分析：Proposition 1 证明灵敏度 $\partial L/\partial\eta_k=\langle G_{k+1},\Delta_k\rangle_F$（$G_{k+1}$ 为伴随反传梯度）；时间加权平均可精确分解为 $\bar A_w^k=P_w^k-S_w^k$（Proposition 2），即持久进展与中心化波动。其二，在线估计：类 Adam 的 EMA 维护 $M_k,v_k$，得 $\hat P_k=\frac1{nd}\|M_k\|_F^2$、$\hat S_k=v_k-\hat P_k$，平衡比 $\hat B_k=\frac{\hat P_k-\gamma\hat S_k}{\hat P_k+\gamma\hat S_k}$，热身后取 $\eta_k=\mathrm{clip}(1+\rho\hat B_k,\eta_{min},\eta_{max})$。其三，共设计：加入 $\mu\sum_k[\hat S_k-\kappa\hat P_k]_+$ 惩罚。全链路只用轨迹观测。
核心创新是把终端任务损失对每个松弛因子的灵敏度的时间平均精确分解为持久进展与中心化波动两项，并证明在温和条件下（可观测能量与 oracle 贡献条件可比、加上漂移界）这一分解足以保证因果控制器的方向正确性。与已有方法的本质区别有三。其一，固定缩放法（常数 $\eta\neq1$）无视轨迹状态，实验显示 $\eta=1.2$ 在 Ouro-1.4B 上平均掉 28 分；其二，定点推理只朝收敛方向单调阻尼，其目标是潜在动力学稳定而非有限预算下的任务精度，而 TAPS 允许步长双向调整；其三，也是最重要的，TAPS 保持学习更新方向 $\Delta_k$ 完全不变、只调制幅度，且调制准则直接对齐终端任务损失——Theorem 4 给出单因子期望损失下降下界 $g_k^{lb}$，Theorem 5 进一步给出 $J_K(\pi)\le J_K(\pi_0)-G_K^{lb}$ 的累计保证，把每一步的局部步长选择与"用更少循环达到目标质量"联系起来。
方法步骤详情
推理流程：(1) 编码器生成 $X_0=E_\theta(u)$，初始化 $M_{-1}=v_{-1}=0$；(2) 每循环计算 $\Delta_k=\Delta_\theta(X_k;u)$；(3) 偏差修正 EMA：$M_k$ 累加 $\Delta_k$，$v_k$ 累加 $\|\Delta_k\|_F^2$；(4) 由此得进展能量 $\hat P_k$、波动能量 $\hat S_k$ 并合成 $\hat B_k\in[-1,1]$；(5) 热身 $K_{warm}$ 循环内 $\eta_k=1$，之后 $\eta_k=\mathrm{clip}(1+\rho\hat B_k,\eta_{min},\eta_{max})$；(6) 执行 $X_{k+1}=X_k+\eta_k\Delta_k$ 至预算 $K$ 或自适应退出。共设计训练加 $\mathcal L_{co}=\mathcal L_{task}+\mathcal L_{ACT}+\mu\sum_k[\hat S_k-\kappa\hat P_k]_+$，仅当波动超过进展 $\kappa$ 倍。附录 D 另给五种变体。
技术新颖性
概念上，首次把更新尺度 $\eta_k$ 立为递归推理的第三控制轴：架构决定重复什么、深度决定重复多久、尺度决定每步多强。理论上，Proposition 1–2 给出灵敏度恒等式与精确分解；Theorem 3 证明当 $b_-/a_+\le\gamma\le b_+/a_-$ 时 $\hat B_k$ 的符号可靠指示增减方向；Theorem 4 给出单步增益下界 $g_k^{lb}$；Theorem 5 证明 $J_K(\pi)\le J_K(\pi_0)-G_K^{lb}$，累计增益够大即 $K_\epsilon(\pi)<K_\epsilon(\pi_0)$。工程上，控制器只需 EMA 与范数统计，零重训即插即用，与自适应退出、层次递归、定点与并行递归均正交。与固定缩放的本质区别：固定缩放无视轨迹状态，$\eta=1.2$ 使 Ouro 平均从 61.10 崩至 32.85；与定点法的区别：TAPS 步长可双向调整，目标对齐有限预算下的任务精度而非渐近收敛，且只调制幅度不改方向。
实验结果
推理-only：Sudoku (H=32) 六个变体超基线 89.67%，最好 TAPS (BB) 90.19% @ 1.039×；Maze (H=16) 基线 78.80%，TAPS (Momentum) 78.90% @ 1.509×；固定 $\eta=1.2$ 均不可用。共设计：Sudoku 91.39% @ 1.300×（Momentum），Maze 79.90% @ 1.561×（P/S-Sign），即摘要的 1.56× 加速。自适应退出：Sudoku-Extreme 上 91.2% Exact 仅用 316.5 有效更新，比 FPRM（91.1%）省 4.7× 计算，定点法 1000 次的精度约 260 次匹配。层次递归：内环优于外环，共设计后 $\Delta_{I-O}$ 最大 +0.67。语言模型：Ouro-1.4B 平均 61.10→62.86（TAPS (BB)），GSM8K 51.78→58.98（+7.20）；Huginn-0125 平均 38.87→39.15。中间层递归：Qwen3-4B（$K=3$）平均 59.22→59.37（TAPS (BB)）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Sudoku (H=32)，推理-only | Accuracy (%) / 达基线速度加速比 | TAPS (BB) 90.19%，1.039×；六变体全部 ≥89.98% | 标准 η=1.0：89.67%，1.000× | +0.52pp，同时省 3.9% 推理时间 |
| Sudoku (H=32)，训练-推理共设计 | Accuracy (%) / 加速比 | TAPS (Momentum) 91.39%，1.300× | 标准共设计：90.60%，1.194× | +0.79pp，加速比再提升 9% |
| Maze (H=16)，训练-推理共设计 | Accuracy (%) / 加速比 | TAPS (P/S-Sign) 79.90%，1.561× | 标准 η=1.0：78.80%，1.000× | +1.10pp，墙钟加速 1.561× |
| Sudoku-Extreme，自适应退出 | Exact 准确率 / 有效更新数 | 91.2% Exact，316.5 有效更新 | FPRM 定点退出：91.1% | 相当精度下计算节省 4.7× |
| Ouro-1.4B 语言模型递归（MMLU/ARC-C/HellaSwag/GSM8K 平均） | 四基准平均准确率 (%) | TAPS (BB) 62.86（GSM8K 58.98） | 标准 η=1.0：61.10（GSM8K 51.78） | 平均 +1.76，GSM8K 单项 +7.20 |
| Qwen3-4B-Instruct 中间层循环 K=3（CommonsenseQA/GPQA-Main/MMLU-Pro 平均） | 三基准平均准确率 (%) | TAPS (BB) 59.37，平均步长 $\bar\eta$=1.065 | Loop η=1.0：59.22 | +0.15，全部控制器自动选 $\bar\eta>1$ |
局限与改进
作者承认三点局限：(i) 理论依赖条件 (4)(5)——只要求可观测能量与 oracle 贡献条件可比，但常数 $a_\pm,b_\pm$ 无法直接测量，漂移界 $\delta_k^{drift}$ 在真实网络上也难以验证，理论保证更多是设计指引而非可检验承诺；(ii) TAPS 仅使用局部轨迹统计（EMA），未纳入更丰富的信号；(iii) 训练-推理共设计只是塑造动力学的初步尝试。我的补充观察：语言模型上的绝对增益偏小（Huginn-0125 平均仅 +0.28，Ouro-1.4B 的 +1.76 主要由 GSM8K 单项拉动，个别单项如 TAPS (Adam) 的 MMLU 68.12 还低于基线 68.23），稳健性存疑；方法需要调节 $\beta,\gamma,\rho,K_{warm},\eta_{min},\eta_{max}$ 六个以上超参；在 $n\times d$ 维隐状态上每循环维护两个 EMA 有额外显存与计算开销；控制器只调制幅度不改方向，若更新方向本身有误则放大只会更快走偏；共设计需要重新训练，与推理-only 的零成本卖点存在一定张力。
独立分析的弱点
第一，增益规模：核心提升集中在合成任务（+0.4~1.1pp）与小型语言模型（平均 +0.1~1.8），对大模型实践者吸引力有限，且 Table 3 个别单项低于基线（TAPS (Adam) 的 MMLU 68.12 vs 68.23），应报告多种子置信区间并扩充基准。第二，超参与热身：需调 $\beta,\gamma,\rho,K_{warm},\eta_{min},\eta_{max}$ 等多个超参；$K_{warm}$ 期间无增益，过长浪费预算；$\gamma$ 的容许区间依赖不可测常数，可引入自适应 $\gamma$ 或元学习控制器。第三，开销：每循环在 $n\times d$ 隐状态上维护两个 EMA，对 4B 级模型不可忽略，可改为子空间或按头统计。第四，幅度调制有天花板：不改方向，方向错误时放大只会更快走偏，可与 Anderson 加速等方向修正法结合。第五，共设计代价：Table 1 的主要精度增益来自共设计（91.39 vs 90.10），削弱免训练卖点，可用 LoRA 式轻量微调缓解。
未来方向
作者提出三个方向：(i) 把理论保证从当前条件 (4)(5) 扩展到更广的递归动力学；(ii) 引入比局部 EMA 统计更丰富的观测信号来驱动调度；(iii) 更深度的训练-推理集成并覆盖更广架构。基于其成果还可延伸：把统一的 $\eta_k$ 推广为 token 级或序列位置的异质步长，处理长序列中不同 token 收敛速度不同的问题；将进展-波动原则迁移到扩散模型与流模型的迭代采样步长调度；与自适应退出结合成"深度×步长"的联合控制器（论文的 TAPS 退出规则已迈出第一步）；在 RL 与蒸馏的 rollout 生成中用 TAPS 提升样本效率；探索 $\eta_k$ 调度与谱归一化、权重约束等训练稳定技术的相互作用；以及把进展-波动平衡作为诊断工具，用于刻画和比较不同递归架构的动力学健康度。
复现评估
可复现性总体良好。所有评测基准公开：Sudoku-Extreme、Maze、MMLU、ARC-C、HellaSwag、GSM8K、CommonsenseQA、GPQA-Main、MMLU-Pro；所有基线模型为开源权重：TRM、Huginn-0125、Ouro-1.4B/2.6B、Qwen3-4B-Instruct。附录 D 给出六种控制器（GD/Sign/Momentum/RMSProp/Adam/BB）的估计量与映射公式，附录 C 给出完整评测协议，附录 B 给出理论假设讨论。核心控制器只是 EMA 加 clip 映射，几十行代码即可实现。算力方面：推理-only 实验只需单卡推理即可复现，门槛低；共设计需重训 TRM 量级模型（数百万参数），成本可控，但复现 Ouro-1.4B/2.6B 级别的重训成本较高。风险点：正文未明确声明开源代码仓库，若实现细节（尤其 $\beta,\gamma,\rho$ 的具体取值与 $K_{warm}$ 设置）不在附录中给出，共设计部分复现难度中等偏难；推理-only 部分则容易。
论文图表
