在 Looped Transformers 中调度递归推理

TL;DR 针对循环 Transformer 固定步长问题，提出 TAPS 调度器，在线追踪更新轨迹的持续性进展与波动来调整步长，无需重新训练即可提升推理准确率，结合训练后最大加速 1.56 倍。

问题

问题背景

当前循环 Transformer（Looped Transformer）作为递归推理模型的核心架构，通过共享参数迭代细化潜在状态来扩展测试时计算，在结构化推理任务中表现突出。然而，这些模型在推理时通常以固定单位步长应用每次学习到的更新，未考虑轨迹中的动态特性。

现有方法局限

固定单位尺度（unit scale）引发两类问题：

• 当更新产生持续进展（persistent progress）时，单位步长趋于保守，限制了循环带来的收敛加速；
• 当更新出现波动（fluctuation）时，单位步长又趋于激进，可能引发振荡或偏离最优解。

论文通过分析终端损失对步长的敏感性，发现其时间平均可精确分解为持续进展项与中心化波动项。现有方法无法显式追踪这两者的平衡，因此无法根据轨迹状态在线调整步长，导致额外循环的收益递减。

为什么这个问题难/重要

• 技术挑战：步长调度需要同时考虑局部梯度信息和全局轨迹统计，且需在无标签、纯推理阶段有效；理论保证（如减少期望终端损失、更快达到目标质量）较难建立。
• 业界关注度：测试时计算扩展（test-time compute scaling）是当前高效推理的研究热点，循环推理模型作为其中的核心架构，其推理效率直接受步长影响。将步长作为继架构和深度之后的第三个控制轴，对实际部署有直接价值。

行业类比

类似于优化器中自适应学习率（如 Adam）利用梯度一阶/二阶矩动态调整步长，TAPS 将进展-波动平衡思想引入循环推理的测试时调度，可类比为推理阶段的“自适应步长控制器”。

核心洞察

• TAPS 将循环推理中的更新步长视为可在线调度的控制变量，而非固定为 1。这一视角区别于现有工作普遍依赖固定步长或仅通过增加循环深度、改进架构来扩展推理计算。TAPS 通过对终端损失关于步长的梯度进行精确分解，分离出持续进步项与中心化波动项，从而判断当前更新是应放大还是缩小，为推理阶段提供了动态的策略，使得步长成为与深度、架构并列的独立控制轴。
• 进度-波动分解为步长调度提供了可解释且无需额外训练的信号。作者证明终端损失对步长的敏感性可分解为持续进步和中心化波动两部分，前者反映更新方向的一致性收益，后者反映振荡带来的风险。据此设计的 TAPS 能在线跟踪二者平衡并调整步长，理论上保证减少期望终端损失并以更少循环达到目标质量。这区别于基于启发式或需要额外训练的自适应推理方法，且实验表明在不重训下即可提升准确率，与训练结合后最高获得 1.56 倍 wall-clock 加速，跨架构通用。

方法

输入与问题定义

对循环 Transformer 在推理时的迭代轨迹，每一步用共享参数对 latent state 施加更新，但其步长固定为 1，导致保守或激进。输入为某一任务下的初始 latent state 与循环模型参数。

关键模块：Progress–Fluctuation 分解

方法先分析终端损失对循环更新尺度的敏感性。对时间平均后的梯度进行推导，得到两项精确分解：

• Persistent Progress：更新方向在时间上保持一致、持续降低损失的分量。
• Centered Fluctuation：更新方向围绕零均值随机波动、造成噪声的分量。

TAPS 在每次循环中在线估计这两项的相对强度，通常使用指数滑动平均（EMA）跟踪 progress 与 fluctuation 的统计量，并计算其比例。

输出与调度策略

基于该比例动态调整下一步的步长（scale）：当 progress 占优时放大步长，加速收敛；当 fluctuation 占优时缩小步长，抑制震荡。调度器输出一个缩放因子，乘到原本的循环更新上。该方法无需重新训练，可直接嵌入现有推理流程；也可将 progress–fluctuation 目标作为训练正则项，进一步改善循环动态。

与同类方法的差异

不同于固定步长、手动退火或基于损失的启发式调度，TAPS 从损失对步长的梯度分解出发，显式分离进展与波动，且不依赖额外标注或复杂二阶信息。

实验

实验设计

作者在多个 结构化推理任务 上进行评估，覆盖无重训练推理时调度与训练-推理联合设计两种场景；同时对比了固定步长等不同推理策略，并在多种 循环架构 上进行验证通用性。

关键发现

• 推理时调度：TAPS 无需重训练即可提升终端准确率，表明在线调节更新步长能有效利用额外循环。
• 训练-推理联合：将 progress-fluctuation 原则 纳入训练后，在匹配基线准确率前提下实现最高 1.56 倍 wall-clock 加速。
• 理论支撑：TAPS 可在更少循环内达到目标质量，且降低期望终端损失。

基线对比解读

传统 固定单位步长 的循环更新在持续进步时显得保守，在波动时又过度激进；TAPS 通过追踪 持续进步项 与 中心化波动项 之间的平衡，在线调整步长，从而更精确地控制每条推理轨迹的收敛行为。与仅增加深度或架构改进不同，TAPS 将 更新尺度 作为独立控制轴，这一思路对实际部署中平衡延迟与准确率具有参考价值。

行业影响

落地场景：TAPS 适用于任何基于循环 Transformer 的推理服务，包括多步推理对话系统、代码生成与调试助手、数学/科学计算求解器、企业自动化决策引擎。这些场景往往需要模型迭代多次以提升答案质量，而固定步长导致算力浪费或收敛不足。

商业价值：推理成本是当前 AI 商业化的重要瓶颈。TAPS 能在不重训练的情况下减少平均循环次数，同时保持或提升准确率，直接降低 GPU 计算时间和成本。若将其原则融入训练，还能带来最高 1.56× 的墙钟速度提升，这意味着相同硬件可服务更多请求，提升毛利率；对终端用户则意味着更低延迟、更快响应。

接口集成：TAPS 作为推理时调度器，不改变模型权重，可无缝集成到现有推理引擎（如 vLLM、TensorRT-LLM）中，作为一个中间件模块，在每个递归更新后计算 progress-fluctuation score 并调整步长。也可以与自适应退出机制结合，进一步减少无效计算。

具体用例：

1. 教育科技数学辅导：使用递归推理模型解答竞赛级数学题，TAPS 根据解题进展动态调整步长，避免在简单步骤上浪费迭代，降低平均响应时间。
2. 企业级代码助手：在处理复杂代码生成或重构任务时，模型需要多次迭代才能输出高质量代码；TAPS 可减少冗余迭代，使 IDE 插件的补全延迟更低。

局限

• - 论文主要在结构化推理任务（如算术、逻辑推理等合成或特定数据集）上验证 TAPS，但未在开放式自然语言生成、长链推理或大规模语言模型（如 70B+ 参数）上系统评估。从摘要的“structured reasoning tasks”可推断，其效果在更复杂、更嘈杂的真实推理轨迹上可能衰减；且实验未明确报告 TAPS 与已有先进推理策略（如多数投票、自适应退出）结合时的鲁棒性。工程落地时，可能需要额外调参和领域适配。
• - TAPS 的核心分解基于 loss 对更新尺度的敏感性，理论保证需要若干假设（如更新有界、损失景观局部良态、progress/fluctuation 可在线估计）。但在实际循环 Transformer 中，这些假设可能不成立，导致调度器性能不稳定。此外，TAPS 需要跟踪每次更新的 progress-fluctuation 平衡，若涉及额外梯度计算或滑动统计，会增加推理时的计算与内存开销，可能抵消部分 wall-clock 加速收益（论文报告最高 1.56×，但未说明在最坏情况下的开销）。
• - 在优化领域，步长调度（如 Polyak 步长、Adam 类自适应学习率、test-time adaptation 动态调整）已有大量工作。TAPS 将这类思想迁移到循环推理的更新尺度上，但本质上仍是启发式在线调度，缺少与这些基线在相同条件下更深入的对比。论文仅在固定单元尺度或简单固定 schedule 上比较，未展示与更先进的 adaptive step-size 方法的差异，因此其相对增益可能源于基线较弱而非新原理。