Scheduling Recursive Reasoning in Looped Transformers
Boyuan Wang, Chengyao Yu, Jiaxi Ren, Hongxin Wei, Bingyi Jing, Yuxin Tao
cs.LG
2026-09-29
南科大提出TAPS,按循环轨迹的进度与波动在线调步长;不重训即可提点,训练协同后迷宫匹配精度最多加速1.56倍。
Looped Transformer 把同一组参数反复套在隐状态上,用循环次数换有效推理深度。已有工作主要改两条轴:循环什么(层、MoE、整模),以及循环多久(固定预算、自适应退出、定点收敛)。每一步更新乘多大的系数,默认一直是 1。
单位步长会在两种阶段同时犯错。更新方向持续一致时,步子偏保守,循环被浪费;方向来回抖时,步子又偏猛,把已有进度冲掉。按终点任务损失选 η 最理想,可推理时没有标签,也没有损失梯度。南科大、港中深和深圳环区研究所把轨迹当成唯一可用信号,问进度和波动的平衡够不够用来在线改步长。
标准循环是 X{k+1}=Xk+Δθ(Xk;u)。TAPS 插入松弛因子 ηk,变成 X{k+1}=Xk+ηk Δθ。η=1 退回原循环;大于 1 沿学到的更新多走一截,叫过松弛;小于 1 把更新缩小,叫阻尼。调度器只能看当前输入和到第 k 步为止的轨迹,不能看答案。
终点损失对 ηk 的偏导等于回传梯度与当前更新的内积。其时间平均可精确拆成持续进度 P 减去中心化波动 S。P 大就该加大步,S 大就该收步。推理时任务版看不见,就用更新 Δ 的指数滑动平均(EMA)来估:
前几步 warmup 强制 η=1,只积累统计量。这套叫 TAPS (Adam),一二阶追踪和 Adam 同构。换成别的估计器,得到 GD、符号、Momentum、RMSProp、Barzilai-Borwein(BB) 五个变体。
理论给充分条件:可观测能量与任务贡献对齐且窗口漂移小时,按 B̂ 调步长能降低条件期望终点损失,并用更少循环打到目标精度。训练可加协同项,只在波动超过进度时惩罚,任务损失和 ACT 保持原样,让循环更顺、调度器才敢把 η 拉过 1。
主设置是 TRM 潜空间循环,数独预算 32 步、迷宫 16 步。
| 设置 | 方法 | 准确率 | 加速 |
| 数独 推理-only | Standard η=1 | 89.67% | 1.000× |
| 数独 推理-only | TAPS (BB) | 90.19% | 1.039× |
| 数独 协同训练 | TAPS (Momentum) | 91.39% | 1.300× |
| 迷宫 推理-only | Standard η=1 | 78.80% | 1.000× |
| 迷宫 协同训练 | TAPS (P/S-Sign) | 79.90% | 1.561× |
不改权重,六种 TAPS 在两个任务上都高于单位步,并更快摸到单位步终点精度。固定 η=0.8 或 1.2 做不到:数独上 η=1.2 掉到 85.43% 且预算内追不上;迷宫上两种常数步都是 N/R。
协同训练后数独最好 91.39%(Momentum,1.300×),迷宫最好 79.90%(P/S-Sign,1.561×),对照单位步 89.67% 与 78.80%。协同后的 Standard 也能到 90.60%/1.194× 与 78.80%/1.249×,TAPS 再往上拉一截。
叠到其他策略:Sudoku-Extreme 自适应退出到 91.2% Exact、316.5 次有效更新,相对 FPRM 在 91.1% 时少用约 4.7 倍算力;定点法 1000 次才到的精度,TAPS 大约 260 次对上;分层循环打在内层比外层更准,协同后内-外差 0.58–0.67 点。并行 PTRM 的具体数字主文未列。
语言模型循环上,Ouro-1.4B 四项平均 61.10→62.86(TAPS BB),增量几乎全在 GSM8K(51.78→58.98),MMLU 从 68.23 微降到 67.95。Huginn-0125 平均 38.87→39.15。Ouro 上固定 η=1.2 平均只剩 32.85。冻结 Qwen3-4B-Instruct 中间层循环 3 次,原模型平均 56.77,单位步 Loop 59.22,TAPS (BB) 59.37,相对 Loop 只多 0.15 点,且平均 η̄ 都大于 1。
已经在跑 looped 推理的人,多了一个几乎不改权重的旋钮:维护一组 EMA,按轨迹把学到的更新拉长或收短。结构推理上,不重训大约 0.3–0.5 点、约 1.04× 加速;愿意重训则迷宫匹配单位步精度能到 1.56 倍墙钟。相对 FPRM 的 4.7 倍差说明同一信号也能驱动早停。
还在训 recurrent LM、或给现成 LLM 套中间层循环的人,预期要压低。Ouro 平均大约 1.8 点且集中在 GSM8K,MMLU 没有变好;Qwen3 上 TAPS 和单位步 Loop 几乎持平。步长作为第三条轴在概念上成立,通用语言任务上目前很细。
固定步长对照把问题坐实了:找一个更好的常数 η 不够,沿轨迹变才有用。
作者列了三条:理论依赖轨迹统计与任务贡献对齐,推广到更一般循环动力学仍开放;TAPS 只用局部轨迹当代理;训练-推理协同只是初步。
主表上推理-only 提升很薄,数独 90.19% 对 89.67%,迷宫 79.10% 对 78.80%。1.56× 绑定协同训练、迷宫、匹配单位步精度,不能读成即插即用的普遍加速。语言模型没有方差或多种子,Ouro 上 TAPS (BB) 的 MMLU 低于单位步。Qwen3 只循环 3 次,0.15 点很难当成泛化证据。4.7× 比的是 FPRM 有效更新次数,没有和固定预算单位步比墙钟。并行 PTRM 与 Ouro-2.6B 的幅度主文未展开。