Adaptive Depth in Looped Transformers: Diagnosing Learned Halting Gates and Trajectory Readouts

循环 Transformer 中的自适应深度：诊断学习型停止门与轨迹读取

Abstract: Looped Transformers increase test-time computation by repeatedly applying a shared recurrent block. Learned halting objectives in looped Transformers typically use a single exit distribution both as the inference-time stopping rule and as the training-time weighting of per-depth losses. This entangles exit selection with trajectory formation: the gate not only chooses which recurrent state to use, but also determines how strongly each intermediate state is supervised. Consequently, poor adaptive-compute performance can arise from the readout, the induced trajectory, or their interaction.

摘要： 循环 Transformer (Looped Transformers) 通过重复应用共享的循环块来增加推理时的计算量。在循环 Transformer 中，学习型停止目标通常使用单一的退出分布，既作为推理时的停止规则，也作为训练时各深度损失的加权依据。这使得退出选择与轨迹形成相互纠缠：门控不仅选择了使用哪个循环状态，还决定了每个中间状态受监督的强度。因此，自适应计算性能不佳可能源于读取机制、诱导轨迹或二者的相互作用。

We study adaptive depth in looped Transformers through this trajectory-readout lens, across controlled synthetic tasks (modular arithmetic and binary parity) and large-scale Ouro-1.4B and 2.6B checkpoints. We find that fixed-prior depth supervision, which shapes the trajectory without an input-dependent halting policy, produces difficulty-aware trajectories whose intermediate states expose useful stopping signals, and that simple post-hoc confidence readouts often match or outperform learned linear and MLP gates.

我们通过这种“轨迹-读取”视角，在受控合成任务（模运算和二进制奇偶校验）以及大规模 Ouro-1.4B 和 2.6B 检查点上研究了循环 Transformer 中的自适应深度。我们发现，固定先验深度监督（在没有输入依赖的停止策略的情况下塑造轨迹）能够产生具有难度感知能力的轨迹，其中间状态暴露了有用的停止信号；此外，简单的事后置信度读取往往能达到或超过学习型线性门和 MLP 门的效果。

Fitting gates on frozen trajectories localizes the failure: it appears to stem mainly from the trajectory induced by joint gate training rather than from limited gate expressivity. The same pattern is present in Ouro evaluations, where pretrained ponder gates are competitive but not uniformly Pareto-optimal, and measured latency confirms that the resulting reductions in average exit depth translate into practical inference-time savings.

在冻结轨迹上拟合门控可以定位故障点：结果表明，问题主要源于联合门控训练所诱导的轨迹，而非门控本身的表达能力受限。同样的模式也出现在 Ouro 的评估中，预训练的思考门（ponder gates）具有竞争力但并非在所有情况下都是帕累托最优的；延迟测量证实，由此带来的平均退出深度降低确实转化为实际推理时间的节省。

Our systematic diagnostic evaluation reframes adaptive depth in looped Transformers as a joint problem of trajectory formation and exit readout, rather than gate learning alone, highlighting a distinction that prior learned-halting work has often left implicit.

我们的系统性诊断评估将循环 Transformer 中的自适应深度重新定义为轨迹形成与退出读取的联合问题，而非单纯的门控学习问题，从而突显了以往学习型停止研究中常被隐含处理的一个关键区别。