arXiv:2607.20519cs.LG2026-07被引 2

发现循环Transformer的自适应深度问题本质是轨迹与读出的协同,而非门控机制本身。

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

论文配图:Adaptive Depth in Looped Transformers: Diagnosing Learned Halting Gates and Trajectory Readouts
图 1 · 摘自论文原文
  • 通过分解轨迹与读出,发现固定先验监督可生成更合理的中间状态路径。
  • 简单置信度读出在多数任务中表现优于学习型门控,且显著降低平均退出深度。
  • 适用于希望提升推理效率、理解自适应计算机制的研究者。

循环Transformer通过重复应用共享的递归模块增加测试时计算量。现有学习式停止机制通常使用单一出口分布同时作为推理阶段的停止规则和训练阶段各深度损失的加权依据,导致出口选择与轨迹形成相互纠缠:门控不仅决定使用哪个递归状态,还影响每个中间状态的监督强度。因此,自适应计算性能不佳可能源于读出方式、诱导的轨迹或二者交互。本文通过轨迹-读出视角系统研究循环Transformer的自适应深度,在受控的合成任务(模运算与二进制奇偶性)以及大规模Ouro-1.4B和2.6B检查点上展开分析。结果表明,固定先验深度监督能生成具备难度感知能力的轨迹,其中间状态蕴含有效停止信号;而简单的后处理置信度读出常可达到甚至超过学习型线性与MLP门控的表现。在冻结轨迹上拟合门控发现失败主要源于联合训练诱导的轨迹质量,而非门控表达能力不足。该模式在Ouro评估中同样存在:预训练的权衡门控具有竞争力但非全局最优,实测延迟也证实平均退出深度下降带来实际推理时间节省。本研究将自适应深度重新定义为轨迹形成与出口读出的联合问题,揭示了以往学习停止工作常隐含的区分。

原文摘要 · Abstract (English)

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. 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. 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. 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推理优化轨迹分析

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。