💬 NLP

LoopMTP: A looped transformer guided by latent multi-token prediction

该论文提出了 LoopMTP，这是一个参数高效的框架，通过利用潜在的多标记预测将中间隐藏状态与未来标记的嵌入进行对齐，从而增强循环 Transformer，以此缓解过度思考并显著提高推理准确性。

原作者： Behzad Shomali, Markus Frey, David Berghaus, Joachim Koehler, Mehdi Ali

发布于 2026-08-05

📖 1 分钟阅读☕ 轻松阅读

原作者： Behzad Shomali, Markus Frey, David Berghaus, Joachim Koehler, Mehdi Ali

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

想象一下你正在试图解决一个非常棘手的谜题，比如一个复杂的数学问题，或者是在写一个带有反转的故事情节。你可以雇佣一支庞大的专家团队，每位专家都有自己的办公桌来处理它。这就是当今大多数强大 AI 模型的工作方式：它们规模巨大，拥有成千上万层堆叠在一起的“专家”。但如果你的办公桌很小，而你只有一个才华横溢的专家呢？如果让你让他反复思考这个问题，通过不断完善答案来解决谜题，他是否依然能成功？这就是人工智能领域的一个重大课题，特别是在一个被称为“循环 Transformer”（looped transformers）的领域中。

长期以来，科学家们一直认为，要让 AI 变得更聪明，需要让它变得更大。但这里有一个陷阱：更大的模型既昂贵又难以在普通计算机上运行。它们有时还会陷入自身思维的死循环，对着一个简单的答案过度思考，直到把事情搞砸。另一种被称为“多 Token 预测”（Multi-Token Prediction）的想法则建议，AI 不应该只是猜测句子中的下一个词，而应该尝试同时预测接下来的“几个”词。这迫使 AI 进行预判，就像一名棋手不仅看一步，而是看三步之后的情况。目前的挑战在于，如何将“反复思考”的方法与“预判未来”的方法结合起来，而不至于让 AI 感到困惑或浪费精力。

于是，由研究员 Behzad Shomali 及其团队提出的新方法 LOOPMTP 登场了。你可以把 LOOPMTP 想象成那个单一位才华横溢专家的聪明策略。与其让专家盯着谜题并不断重写笔记（这往往会导致抹掉好的想法），LOOPMTP 为他提供了一个特殊的指南。每当专家对谜题进行新一轮观察（一个“循环”）时，他都会被轻轻地引导，目光始终聚焦在解题过程中某个特定的未来环节上。

以下是它的实际运作方式：想象 AI 正在写一个故事。在标准的循环中，它可能会写出一个句子，然后立即忘记刚才写了什么，接着尝试重写它，这往往会导致错误。LOOPMTP 改变了规则。当 AI 进行第二轮故事创作时，它会被秘密告知：“嘿，记住你还需要为两个步骤之后的那个词做好准备。”在第三轮时，它会被要求看向三个步骤之后。这就是“多 Token 预测”作为引导的部分。它并不是强迫 AI 直接喊出未来的词，而是让 AI 的内部思维与那些未来词汇的“概念”保持一致。这就像一位徒步旅行者，口袋里揣着地图，不是为了向前跳跃，而是为了确保每走一步都朝着正确的方向前进。

研究人员还加入了一个“门控”（gatekeeper）机制。过去，当 AI 进行循环时，它往往会丢弃前几次循环的工作成果，就像厨师每次品尝酱汁时都把它倒掉一样。LOOPMTP 则保留了所有版本的酱汁，但使用一个智能门控来决定将每个版本中的多少比例混合进最终的成品中。这确保了即使 AI 对问题进行了多次思考，也不会丢失任何有用的信息。

这种方法的测试结果非常令人振奋。当团队在包括数学、编程和常识问答在内的各种任务上测试 LOOPMTP 时，它的表现明显优于其他循环模型。事实上，与不使用这种循环技巧的标准模型相比，它的准确率平均提高了高达 8.1%。更令人印象深刻的是，即使要求 AI 连续“思考” 15 次，它依然能保持稳定并持续进步。而其他类似的模型往往在仅仅几次循环后就会崩溃或陷入混乱。

论文指出，这种方法是让较小的 AI 模型变得更聪明，而无需构建庞大且昂贵的模型的有效途径。它表明，通过引导 AI 的内部思维向未来靠拢，我们可以防止它过度思考，并帮助它更高效地利用有限的计算能力。虽然研究人员也注意到，这种方法在处理数学和推理任务时效果尤为显著，但也指出存在一个极限：如果你让模型思考次数过多，它最终还是可能会变得困惑。但就目前而言，LOOPMTP 提供了一种新鲜且富有启发性的方式，帮助 AI 在不需要更大“大脑”的情况下，“思考得更深入”。

技术摘要：LoopMTP

问题陈述
大语言模型（LLMs）展示了强大的推理能力，这种能力通常与架构深度密切相关。然而，增加深度需要显著增加参数量，这在资源受限的环境中（例如，本地部署的医疗保健环境）是不切实际的。循环 Transformer 提供了一种参数高效的替代方案，通过在多次迭代（ $T$ ）中重复使用单个层堆栈来模拟有效深度。尽管具有潜力，但现有的循环模型存在两个关键局限性：

潜在过度思考（Latent Overthinking）： 中间表示在每次迭代中都会被覆盖。如果预测在早期已经正确，随后的循环可能会将其修正为错误。
无差异计算（Undifferentiated Computation）： 随着迭代次数增加，隐藏状态变得越来越相似，导致冗余工作和计算浪费。
此外，现有方法缺乏跨循环的引导，且多 Token 预测（MTP）——旨在鼓励前瞻性推理的技术——在历史上仅在大型规模或使用昂贵的完整词表投影时才显示出优势。

方法论：LoopMTP
作者提出了 LoopMTP，这是一种将多 Token 预测（MTP）作为潜在引导信号，以解决过度思考和无差异计算问题的循环 Transformer 架构。其核心组件包括：

MTP 引导的循环块（MTP-Guided Looped Block）： 模型对整个 Transformer 堆栈进行 $T$ 次循环。为了防止有用信息的覆盖并引导轨迹，该架构采用了：
- 迭代索引嵌入（Iteration-Index Embedding）： 将归一化的迭代索引附加到 Token 向量中，以指示当前步骤。
- 逐迭代归一化（Per-Iteration Normalization）： 特别针对每次迭代的状态应用 LayerNorm，以处理特定于迭代的统计数据。
- 输入融合（Input Fusion）： 在进入共享骨干网络之前，将前一次迭代的输出和重新归一化的 Token 嵌入进行拼接并进行投影。
- Loop-LNS（循环层归一化缩放）： 对残差流应用修改后的缩放因子 $1/T$ ，以维持在不同展开次数下的稳定性，避免在渐进深度缩放中观察到的梯度中和现象。
软 MTP 对齐（潜在引导）（Soft MTP Alignment (Latent Guidance)）： 为了避免计算昂贵的完整 Token 分布预测，LoopMTP 使用了一种软对齐目标。对于第 $t$ 次迭代（其中 $t \ge 2$ ），隐藏状态通过与未来第 $t$ 个步长的 Token 嵌入（ $u_{i+t}$ ）的余弦相似度进行对齐。这作为一个表示正则化项，引导每一次迭代去预判特定的未来 Token，而无需完整的词表投影。第一轮迭代（ $t=1$ ）保持不受约束，以允许模型学习丰富的初始表示。
学习型状态聚合（Learned State Aggregation）： LoopMTP 不会丢弃中间状态，而是使用 Token 权重的加权和对所有 $T$ 次迭代的输出进行聚合。一个轻量级的共享线性门控计算逐迭代权重，使模型能够保留来自早期循环的有用信息，同时整合来自后期循环的新计算。一个“沉思正则项”（Ponder Regularizer，基于 KL 散度）鼓励在迭代之间进行均匀分布的注意力分配，以防止过早收敛。
训练目标（Training Objective）： 总损失结合了标准的下一 Token 预测（NTP）交叉熵、软 MTP 对齐损失以及沉思正则项。

核心贡献

架构创新： 引入了 LoopMTP，它在结构上将循环次数与对未来 Token 的预判联系起来，为解决过度思考和无差异计算提供了自然的补救措施。
高效的 MTP 集成： 证明了通过轻量级的余弦对齐信号，MTP 可以有效地应用于参数高效的循环模型中，避免了标准 MTP 的计算开销。
稳定训练： 所提出的架构能够在广泛的循环次数范围内（高达 $T=15$ ）实现稳定训练，而无需针对每个循环次数进行大量的超参数网格搜索，这是以往循环模型（如 LoopFormer）中常见的问题。
小型专家模型： 展示了 LoopMTP 可以训练出在相同参数预算下优于更大非循环基准模型的领域特定小型专家模型（例如数学领域）。

结果
在 260M 参数、基于 6.8B Token 的模型上进行的评估结果如下：

性能提升： 在通用任务上，LoopMTP 在相对平均准确率方面比参数匹配的非循环基准模型高出达 8.1%。在 QA、数学和代码任务上，它也比非循环基准模型降低了高达 7.5% 的每字节比特数（BPB），并显著优于最先进的 LoopFormer（在 28 项匹配对比中赢得了 27 项）。
可扩展性： 与在较高循环次数下会出现训练不稳定和性能下降的 LoopFormer 不同，LoopMTP 随着 $T$ 的增加，在困惑度（Perplexity）和通用任务准确率方面表现出单调改进。
机制分析：
- 表示（Representations）： 结合 MTP 后，各次迭代间的隐藏状态变得更加独特，证实了模型学习到了独特的变换而非冗余变换。
- 真值检索（Ground-Truth Retrieval）： 带有 MTP 引导时，真值 Token 排名显著提高（高达 35.6 倍），表明对齐信号有效地传递到了语言模型头。
- 门控（Gating）： 学习到的门控机制会随深度调整，后期迭代显示出更高且更具变动性的门控值，表明其具备一种依赖输入的、用于信息保留的策略。
领域专家： 在专门的数学训练设置中，一个 $T=15$ 的 LoopMTP 模型在 GSM8K 上达到了 19.03% 的准确率，相对于相同参数量的非循环基准模型（7.05%）实现了约 170% 的相对提升。

意义与主张
论文声称 LoopMTP 成功地弥合了参数效率与强大推理能力之间的差距。通过利用 MTP 作为潜在引导信号，它解决了以往循环 Transformer 的根本不稳定性及冗余问题。作者强调，这种方法允许在有限硬件（例如本地部署）上部署强大的推理模型，这对于数据无法离开现场的敏感领域（如医疗保健）至关重要。这项工作表明，当通过前瞻性目标进行引导时，通过循环获得的“有效深度”可以超越针对特定推理任务的原始参数缩放。作者保持了谦逊的态度，指出在专门领域（除数学之外）的收益仍需进一步验证，且有效循环深度的上限仍是一个开放性问题。

技术摘要：LoopMTP

类似论文