💬 NLP

LoopFormer: Elastic-Depth Looped Transformers for Latent Reasoning via Shortcut Modulation

本文提出了 LoopFormer，一种通过捷径一致性训练方案使循环 Transformer 能够根据计算预算动态调整推理深度的模型，从而在资源受限和充裕场景下均能实现稳健的自适应语言建模与推理。

原作者： Ahmadreza Jeddi, Marco Ciccone, Babak Taati

发布于 2026-02-13

📖 1 分钟阅读☕ 轻松阅读

原作者： Ahmadreza Jeddi, Marco Ciccone, Babak Taati

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇论文介绍了一种名为 LoopFormer 的新型人工智能模型。为了让你轻松理解，我们可以把传统的 AI 模型想象成**“流水线工厂”，而 LoopFormer 则像是一个“拥有无限耐心的超级工匠”**。

1. 传统 AI 的困境：要么太慢，要么太浅

想象一下，你让一个传统的 AI 模型（比如现在的聊天机器人）去解一道复杂的数学题或写一个故事。

传统做法：它像一条固定的流水线。无论问题简单还是困难，它都必须经过完全相同的 24 个步骤（层）才能给出答案。
- 问题：如果问题很简单（比如“苹果是什么颜色的？”），它还要走满 24 步，这就浪费算力（就像为了买一瓶水，非要跑完整个马拉松）。
- 问题：如果问题很难，它可能还没走到第 24 步就“晕”了，或者因为步骤固定，无法在中间停下来思考。

以前的尝试（比如“早退机制”）试图让模型在简单问题时提前结束，但这就像让一个工匠在还没想清楚时就匆忙交卷，结果往往是一团糟（论文中称为“表示坍塌”）。

2. LoopFormer 的创意：像“揉面团”一样的弹性深度

LoopFormer 的核心思想是：同一个模型，可以根据你给的时间（算力预算），灵活地决定思考多久。

核心比喻：揉面团
想象你在揉面团。
- 传统模型：规定必须揉 24 下，不管面团熟没熟。
- LoopFormer：它只有一个“揉面动作”（共享的模块），但它可以重复这个动作。
  - 如果你只有1 分钟（低算力预算），它就揉 3 下，然后告诉你：“虽然没完全熟，但这是目前最好的状态。”
  - 如果你有10 分钟（高算力预算），它就揉 24 下，把面团揉得极其细腻。
- 关键点：无论揉几下，它都能保证“揉面”的质量是连贯的，不会因为揉的次数少就变成一坨烂泥。

3. 它是怎么做到的？（两个魔法）

为了让这个“弹性深度”不翻车，LoopFormer 用了两个聪明的技巧：

魔法一：给每一步都贴上“时间标签”和“步长标签”

以前的循环模型只是机械地重复。LoopFormer 给每一次重复都加上了两个“调料”：

当前时间 ( $t$ )：告诉模型“现在是在刚开始，还是快结束了？”
步长 ( $\Delta t$ )：告诉模型“这次是迈大步（粗略思考）还是迈小步（精细思考）？”

比喻：就像你在学习一门语言。

如果是大步（步长大），你快速浏览文章大意（适合时间紧）。
如果是小步（步长小），你逐词逐句分析语法（适合时间充裕）。
LoopFormer 知道自己在哪个阶段，所以它不会在刚开始时就试图分析语法，也不会在最后还在浏览大意。

魔法二： “短跑”与“长跑”的同步训练（捷径一致性）

这是论文最厉害的地方。在训练时，LoopFormer 会同时做两件事：

让它跑一次全程（比如 24 步），得到完美答案。
让它跑一次短途（比如只跑 4 步），然后强行要求它：“虽然你只跑了 4 步，但你的答案必须和跑完 24 步的‘完美版本’在方向上是一致的！”

比喻：这就像教练训练运动员。
教练不仅让运动员跑全程，还让他跑短距离。但教练会拿着全程的录像（完美答案）告诉短跑运动员：“你虽然只跑了 400 米，但你跑出的姿势和节奏，必须和跑完 4000 米时的状态保持‘神似’。”
这样，无论用户给多少时间（预算），模型都能给出一个**“虽短但精”**的答案，而不是胡言乱语。

4. 实际效果如何？

论文通过大量实验证明：

省钱：在算力紧张时（比如手机端运行），它只跑几步，依然能给出不错的回答，不会像以前的模型那样直接“变傻”。
更强：如果给足算力，它能通过反复思考（增加循环次数），把答案打磨得比传统模型更精准，特别是在逻辑推理任务上。
不崩溃：以前的模型一旦减少步骤，内部逻辑就会崩塌（就像积木塔少了几块就倒了），而 LoopFormer 无论几步，内部逻辑都是稳固且不断进化的。

总结

LoopFormer 就像是给 AI 装上了一个**“智能油门”**。

以前，AI 要么全速冲刺（费电），要么急刹车（变笨）。
现在，你可以根据路况（算力预算）随意控制油门。路宽就深思考，路窄就浅思考，但无论开多快，车都不会散架，方向始终正确。

这项技术让未来的 AI 既能在大服务器上处理复杂的科学难题，也能在你的手机上快速回答日常问题，而且不需要重新训练，真正实现了**“按需计算”**。

1. 研究背景与问题 (Problem)

背景：
循环 Transformer（Looped Transformers，即参数共享的 Transformer）因其高效的参数利用和在算法推理任务中的出色表现而受到关注。研究表明，这类模型具有“隐式推理”（latent reasoning）的归纳偏置，即通过内部状态的迭代更新来模拟链式思考（Chain-of-Thought），且推理能力通常随计算深度（循环次数）的增加而提升。

核心问题：
现有的循环 Transformer 方法在训练和推理时通常固定循环次数。这导致了一个关键缺陷：

缺乏灵活性：模型无法根据可用的计算预算（Compute Budget）动态调整深度。
分布外失效（Out-of-Distribution Failure）：如果在训练时固定为 $L$ 次循环，当在推理时减少循环次数（例如只运行 $M < L$ 次）时，模型的表示往往会崩溃或性能急剧下降，因为短路径的表示分布与长路径不同。
计算浪费：为了适应不同难度的任务，用户无法灵活地分配计算资源，导致要么计算不足，要么过度计算。

目标：
构建一个**弹性深度（Elastic Depth）**的循环 Transformer，使其能够在不重新训练的情况下，根据用户指定的计算预算（循环次数 $M$ ）灵活调整，且在短路径下保持信息丰富，在长路径下持续优化。

2. 方法论 (Methodology)

作者提出了 LoopFormer，一种基于**捷径调制（Shortcut Modulation）**的循环语言模型。其核心思想是将迭代表示的 refinement 视为表示空间中的一条轨迹，并借鉴扩散模型（Diffusion Models）和神经 ODE 的思想。

2.1 核心架构：条件化循环

LoopFormer 在每次循环迭代 $i$ 时，不仅依赖前一步的隐藏状态，还显式地条件化于两个变量：

归一化时间 $t_{i-1}$ ：当前迭代在总轨迹中的累积位置（ $0 \le t \le 1$ ）。
步长 $\Delta_i$ ：当前迭代的步长大小（ $\Delta_i = t_i - t_{i-1}$ ）。

调制机制：
这两个标量通过正弦 - 余弦频率编码和 MLP 投影，生成条件向量 $c_i$ 。该向量用于调制 LoopFormer Block 中的组件：

RMSNorm 缩放：调整归一化层的尺度。
残差门控（Gating）：控制 MHSA（多头自注意力）和 FFN（前馈网络）残差连接的权重。
这使得模型能够感知当前处于轨迹的哪个阶段以及步长的粗细，从而在不同粒度的调度下保持一致的行为。

2.2 训练策略：捷径一致性（Shortcut Consistency）

为了训练出能处理任意长度轨迹的模型，作者提出了一种捷径一致性训练方案：

多轨迹采样：在训练每个 Batch 时，除了计算完整的 $L$ 步轨迹外，还随机采样一个较短的 $S$ 步轨迹（ $1 \le S < L$ ）。
一致性损失（Consistency Loss）：
- 计算完整轨迹 $L$ 的最终输出 logits。
- 计算短轨迹 $S$ 的最终输出 logits。
- 引入一个停止梯度（Stop-Gradient）操作，强制短轨迹的表示向长轨迹的最终表示对齐。
- 损失函数公式： $\mathcal{L} = \mathcal{L}_{LM} + \lambda_1 \mathcal{L}_{S} + \lambda_2 \mathcal{L}_{cons}$ ，其中 $\mathcal{L}_{cons}$ 是短轨迹与长轨迹（停止梯度后）的均方误差。

这种机制本质上是在循环内部进行自蒸馏，确保较短的“捷径”路径也能学习到高质量、信息丰富的表示，而不仅仅是长路径的截断。

2.3 推理阶段：弹性深度

在推理时，用户指定计算预算 $M$ （循环次数）和步长调度 $\Delta_M$ 。模型直接运行 $M$ 次循环，无需微调。由于训练时的条件化机制，模型能够根据 $t$ 和 $\Delta$ 自动调整内部状态演化，实现平滑的性能缩放。

3. 主要贡献 (Key Contributions)

提出 LoopFormer 架构：定义了一类基于捷径调制的循环 Transformer，通过条件化归一化时间和步长，实现了轨迹一致性的表示学习。
捷径一致性训练协议：引入了一种新的训练目标，通过在不同长度的轨迹间强制表示对齐，实现了无需重新训练的弹性深度推理。
表示动力学分析：利用各向异性（Anisotropy）、曲率（Curvature）、熵（Entropy）和 CKA 相似度等指标，证明了传统早退（Early Exit）机制会导致表示停滞（Collapse），而 LoopFormer 保持了随深度演化的非退化状态。
实证性能突破：在语言建模（Perplexity）和零样本推理（Zero-shot Reasoning）基准测试中，LoopFormer 在低预算下表现优异，且随着计算预算增加，性能平滑提升，显著优于其他循环变体。

4. 实验结果 (Results)

实验在 The Pile 数据集上进行，对比了固定深度的非循环 Transformer、标准循环模型、TMLT（时间调制循环 Transformer）以及 LoopFormer。

困惑度（Perplexity）表现：
- 在相同的 FLOPs 预算下，LoopFormer 显著优于其他循环模型（如 Base-Loop, TMLT）。
- 即使在极低的预算下（例如仅运行 2 次或 4 次循环），LoopFormer 的困惑度也接近同参数量的非循环基线，且远优于其他弹性深度模型。
- 随着循环次数增加，性能平滑提升，没有出现性能崩塌。
推理能力（Reasoning）：
- 在 10 个零样本推理基准（如 COPA, HellaSwag, ARC 等）上，LoopFormer 展现了最强的推理能力。
- 特别是在中等预算下，LoopFormer 的准确率显著高于其他循环模型，证明了其隐式推理能力的有效性。
表示分析：
- 避免崩溃：图 4 和图 5 显示，传统的早退模型（Naive-Loop-EE）在增加循环步数时，表示变化极小（高 CKA 相似度，平坦的曲率），表明陷入停滞。
- 持续演化：LoopFormer 的表示在中间步骤表现出显著的曲率变化和熵增加，说明模型正在有效地利用额外计算进行信息提炼，直到接近终点才趋于稳定。
轨迹选择：
- 研究发现，在固定预算下，“先粗后细”（早期步长大，后期步长小）的调度策略通常能获得最佳性能，这与模型在早期进行粗略探索、后期进行精细修正的机制相符。

5. 意义与展望 (Significance)

可控的大语言模型：LoopFormer 为构建**预算感知（Budget-aware）**的大语言模型开辟了新路径。用户可以根据延迟或成本限制，动态调整模型的推理深度，而无需为不同深度训练多个模型。
隐式推理的机制验证：该工作证实了循环 Transformer 具有内在的隐式推理能力，且通过正确的条件化训练，这种能力可以灵活地适应不同的计算资源。
效率与性能的平衡：它解决了循环模型长期以来“固定深度”的局限性，使得参数共享模型在灵活性和性能上都能与深度非共享模型竞争。
未来方向：论文指出了实例级（Instance-level）自适应预算调度、更深层的理论分析以及因果性表示分析是未来的研究方向。

总结：LoopFormer 通过引入时间/步长条件化和捷径一致性训练，成功解决了循环 Transformer 在变长推理路径下的表示退化问题，实现了真正的弹性深度推理，为高效、可控的下一代语言模型设计提供了重要参考。