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

# 诊断学习到的暂停门和轨迹读出 来源：https://arxiv.org/html/2607.20519 ## 循环变压器中的自适应深度：诊断学习到的暂停门和轨迹读出 Andrei Cristian Popescu 剑桥大学 Haitz Sáez de Ocáriz Borde 剑桥大学 Pietro Liò 剑桥大学 ###### 摘要 循环 Transformer 通过重复应用共享的循环块来增加测试时的计算量。循环 Transformer 中学习到的暂停目标通常使用单一的退出分布作为推理时的停止规则，并作为每个深度损失训练时的权重。这将退出选择与轨迹形成纠缠在一起：门不仅决定使用哪个循环状态，还决定了每个中间状态被监督的强度。因此，较差的自适应计算性能可能源于读出、诱导的轨迹，或它们的交互。我们通过轨迹-读出的视角研究循环 Transformer 中的自适应深度，涵盖受控合成任务（模运算和二进制奇偶性）以及大规模 Ouro-1.4B 和 2.6B 检查点。我们发现，固定先验深度监督（在不依赖输入停止策略的情况下塑造轨迹）会产生具备难度感知能力的轨迹，其中间状态暴露了有用的停止信号，而且简单的后验置信度读出通常匹配或超过学习到的线性门和 MLP 门。在冻结的轨迹上拟合门可以定位失败原因：似乎主要源于联合门训练所诱导的轨迹，而非门表达能力的限制。在 Ouro 评估中同样观察到这一模式，预训练的 ponder 门具有竞争力但并非统一的帕累托最优，测量的延迟证实了平均退出深度的降低转化为实际推理时间的节省。我们的系统性诊断评估将循环 Transformer 中的自适应深度重新定义为轨迹形成和退出读出的联合问题，而非仅仅是门的学习问题，强调了先前学习暂停工作通常隐含的区分。 参考图注 图 1：自适应深度是轨迹-读出问题，而不仅仅是门学习问题。*左：* 循环 Transformer 的循环轨迹 h_1, ..., h_T 由训练目标塑造，而单独的读出（学习到的门或简单的启发式方法）决定何时停止。*右：* 在 MANO 上，固定先验轨迹上的后验读出在低于固定深度基线的退出深度下达到接近上限的准确率，匹配或优于 MLP 门，后者的停止分布既选择退出又塑造其训练的轨迹。 ## 1 引言 本文对循环 Transformer 中的自适应深度进行了系统性诊断研究。学习到的暂停机制承诺通过将循环深度分配给输入并在无需进一步细化时停止来实现自适应计算。然而，当前循环 Transformer 中学习到的暂停目标耦合了两个应加以区分的角色。学习到的退出分布用作推理时的停止规则，但同时也对跨循环深度的训练损失进行加权。因此，学习到的门不仅仅决定何时停止；它还改变了从中选择退出的循环轨迹。这种耦合使得自适应深度难以诊断：性能-计算权衡的不足可能源于停止规则、诱导的轨迹，或它们的交互。因此，我们将自适应深度作为轨迹-读出问题进行研究，改变训练目标、学习到的门、后验读出、强制退出，并评估在 MANO（合成模算术任务）和奇偶性任务上训练的受控模型，以及预训练的 Ouro 检查点。我们的结果表明，自适应深度不仅仅是学习更好的暂停门的问题。固定先验深度监督（在不学习依赖输入的停止策略的情况下塑造轨迹）可以产生具备难度感知能力的循环轨迹，其中间状态支持强大的后验提前退出。在几种设置中，简单的置信度读出在计算-准确度权衡上匹配甚至超过学习到的门。我们不将其解释为原则上反对学习暂停。相反，我们的结果表明，当前学习到的门目标可能失败，因为它们纠缠了轨迹形成和退出选择，未来的自适应深度方法应将这两个组件都视为首要设计选择。我们的贡献包括：(i) 提出循环 Transformer 中自适应深度的轨迹-读出视角；(ii) 系统评估训练目标、门参数化、后验读出、强制退出轨迹、模型规模和延迟；(iii) 引入固定先验深度监督以解耦轨迹形成和退出选择；(iv) 在 MANO 上展示固定先验训练能产生具备难度感知能力的循环轨迹；(v) 表明简单的后验置信度读出可以匹配或超过学习到的线性门和 MLP 门；(vi) 在 Ouro-1.4B 和 Ouro-2.6B 上展示预训练的 ponder 门并非统一帕累托最优；(vii) 提供延迟测量，将平均退出深度的降低与实际的推理时间加速联系起来。 ## 2 相关工作 接下来，我们回顾与我们研究最相关的两个轴线的先前工作：循环深度 Transformer 架构和用于选择计算量的自适应计算方法。 ### 2.1 循环 Transformer 循环 Transformer 用共享层块的重复应用取代不同层的堆叠，将有效计算深度与参数数量解耦。循环 Transformer 领域的早期工作包括通用 Transformer (Dehghani et al., 2019 (https://arxiv.org/html/2607.20519#bib.bib13))、隐式深度模型如深度平衡模型 (Bai et al., 2019 (https://arxiv.org/html/2607.20519#bib.bib23))，以及 ALBERT 中的跨层参数共享 (Lan et al., 2020 (https://arxiv.org/html/2607.20519#bib.bib22))。最近的工作将循环重新框定为潜在计算，其中每个循环步骤在执行预测之前细化隐藏表示 (Saunshi et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib25); Geiping et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib18); Zhu et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib14); Jeddi et al., 2026 (https://arxiv.org/html/2607.20519#bib.bib26))。这种深度方向的循环不同于通过连续隐藏状态或潜在标记扩展序列的方法，因为计算仍保持在模型的深度轴上，而非上下文中 (Hao et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib27); Zelikman et al., 2024 (https://arxiv.org/html/2607.20519#bib.bib28); Huang et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib29))。循环模型已被证明在迭代和算法计算上具有有前景的性能，包括可编程计算、长度泛化以及合成任务上的多步细化 (Giannou et al., 2023 (https://arxiv.org/html/2607.20519#bib.bib24); Anil et al., 2022 (https://arxiv.org/html/2607.20519#bib.bib30); Yang et al., 2024 (https://arxiv.org/html/2607.20519#bib.bib31); Kohli et al., 2026 (https://arxiv.org/html/2607.20519#bib.bib32))。最近的更多语言建模工作研究这种循环深度缩放是否适用于更大规模的模型，在这些模型中额外的计算可以在潜在空间中花费，而不是通过更长的生成推理轨迹 (Geiping et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib18); Zhu et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib14); Jeddi et al., 2026 (https://arxiv.org/html/2607.20519#bib.bib26))。我们的工作建立在这种将循环 Transformer 视为增加测试时计算的手段的观点之上，但侧重于一个不同的问题：训练目标如何塑造循环轨迹，以及退出读出如何将它们转换为计算-质量权衡。 ### 2.2 自适应计算 自适应计算方法旨在通过为不同标记分配不同数量的计算来改善性能-计算权衡。早期方法如 ACT (Graves, 2017 (https://arxiv.org/html/2607.20519#bib.bib5)) 和 PonderNet (Banino et al., 2021 (https://arxiv.org/html/2607.20519#bib.bib12)) 学习概率性的暂停策略，而后来的工作将自适应深度应用于 Transformer 和语言模型 (Hou et al., 2020 (https://arxiv.org/html/2607.20519#bib.bib6); Elbayad et al., 2020 (https://arxiv.org/html/2607.20519#bib.bib7); Fedus et al., 2021 (https://arxiv.org/html/2607.20519#bib.bib8); Bae et al., 2023 (https://arxiv.org/html/2607.20519#bib.bib9); Elhoushi et al., 2024 (https://arxiv.org/html/2607.20519#bib.bib10); Raposo et al., 2024 (https://arxiv.org/html/2607.20519#bib.bib11))。这个想法自然适应于循环 Transformer。它们的循环结构产生中间状态，从而允许我们提前退出并承诺其中一个隐藏状态，或者通过额外迭代继续细化它们。因此，最近的循环语言模型使用学习到的门或路由器来分配循环深度 (Dehghani et al., 2019 (https://arxiv.org/html/2607.20519#bib.bib13); Zhu et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib14); Song et al., 2026 (https://arxiv.org/html/2607.20519#bib.bib15); Bae et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib17))。学习到的门并非唯一的停止信号。早期退出的前馈 Transformer 方法也使用置信度、熵或预测稳定性作为终止计算的廉价标准 (Schuster et al., 2022 (https://arxiv.org/html/2607.20519#bib.bib19); Xin et al., 2020 (https://arxiv.org/html/2607.20519#bib.bib20); Zhou et al., 2020 (https://arxiv.org/html/2607.20519#bib.bib21))，而链式思考轨迹上的熵动态也类似地与推理正确性相关 (Català et al., 2026 (https://arxiv.org/html/2607.20519#bib.bib39))。在循环 Transformer 中，类似的信号可以跨循环迭代计算，包括隐藏状态、logits 或预测分布的变化 (Geiping et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib18))。我们的工作建立在这一观察之上，但分离了通常被隐式耦合的两个角色：训练目标塑造循环轨迹，而读出则从诱导的轨迹中选择退出。先前的工作问的是如何通过学习的暂停策略、路由器或其他依赖输入的分配规则使循环深度成为动态 (Geiping et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib18); Zhu et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib14); Song et al., 2026 (https://arxiv.org/html/2607.20519#bib.bib15); Bae et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib17))。我们问的是一个互补的问题：当这样的方法成功或失败时，是系统的哪一部分负责？在循环 Transformer 中，暂停机制不仅是停止规则；在训练期间，它还决定了哪些循环深度受到监督，从而塑造了后来选择退出的轨迹。我们通过实验隔离这一区别，将轨迹形成与退出读出分开。 ## 3 将自适应深度形式化为轨迹和读出 在本节中，我们介绍全文使用的符号和机制。我们首先将循环 Transformer 轨迹形式化为通过重复应用共享 Transformer 块产生的隐藏状态序列。然后我们描述学习的暂停门，它在推理时定义退出分布，并在训练时对跨循环深度的损失进行加权。最后，我们引入固定先验深度监督和后验轨迹读出，这使我们能够将轨迹形成与退出选择分开。 ### 3.1 循环 Transformer 轨迹 设 x = (x_1, ..., x_M) 为长度为 M 的输入序列，并设 h_0 = E(x) ∈ R^(M×d) 为其标记嵌入，其中嵌入函数 E: R^|V| → R^d，词汇表大小 |V|，隐藏维度 d。循环 Transformer 重复应用一个由 L 个 Transformer 层组成的共享块 F_θ 来更新隐藏状态，其中 θ 表示模型参数： h_t = F_θ(h_{t-1}), for t = 1, ..., T 其中 T 是最大的可能循环次数。生成的隐藏状态可以使用层 lmhead: R^d → R^|V| 进行解嵌入。循环 Transformer 自然地在潜在空间中诱导一个轨迹： τ_θ(x) = {h_i}_{i=1}^T。 在这项工作中，我们研究这个诱导轨迹如何被自适应计算训练目标塑造，以及轨迹代理在预测计算何时应该停止方面是否与训练的门具有竞争力。循环 Transformer 的额外架构背景和符号在附录 A.1 (https://arxiv.org/html/2607.20519#A1.SS1) 中提供。 ### 3.2 自适应计算的门控机制 循环 Transformer 中自适应计算的一种常见方法是附加一个受 PonderNet 启发的中止模块，该模块预测循环深度上的退出分布 (Banino et al., 2021 (https://arxiv.org/html/2607.20519#bib.bib12); Zhu et al., 2025 (https://arxiv.org/html/2607.20519#bib.bib14); Song et al., 2026 (https://arxiv.org/html/2607.20519#bib.bib15))。设 e_t(x) = σ(g_φ(h_t)) ∈ (0,1) 表示在循环 t 退出（给定模型之前未退出）的条件概率。在这项工作中，我们将门实验为线性投影和 MLP。预测的条件停止概率在退出深度上诱导一个分布： q_t(x) = e_t(x) ∏_{j<t} (1 - e_j(x)) for t = 1, ..., T-1. q_T(x) = ∏_{j<T} (1 - e_j(x)). 注意，我们将所有剩余的概率质量分配给最终的循环深度 T，确保 q(·|x) 是关于退出的有效分布。在推理时，学习到的分布可以用作停止规则，当累积退出分布超过阈值 α 时停止计算。我们可以如下定义提前退出深度 d_α： d_α(x) = min { t : ∑_{s=1}^t q_s(x) ≥ α }。 在训练期间，退出分布用于对不同循环深度的损失进行加权。给定每个深度的交叉熵损失 ℓ_t，模型与门联合训练，使用以下目标： L(θ, φ) = E_(x,y) [ ∑_{t=1}^T q_t(x) ℓ_t(x,y) + β KL(q(x) || π) ], 其中 π 是深度的目标先验，β 控制正则化的强度。我们注意到，由于学习到的权重 q_t(x) 出现在损失内部，门不仅仅是推理时的停止规则。它还决定了每个循环深度在训练期间贡献多少预测损失。因此，学习到的轨迹 τ_(θ,φ)(x) = {h_t}_{t=1}^T