🤖 machine learning

Looped Diffusion Language Models

本文介绍了 LoopMDM，这是一种针对掩码扩散模型的新方法，它选择性地循环使用早期至中层的 Transformer 层，从而在训练效率和推理性能方面优于标准掩码扩散模型，同时还能在推理时实现灵活的算力扩展。

原作者： Sanghyun Lee, Chunsan Hong, Seungryong Kim, Jonghyun Lee, Jongho Park, Dongmin Park

发布于 2026-05-26

📖 1 分钟阅读☕ 轻松阅读

原作者： Sanghyun Lee, Chunsan Hong, Seungryong Kim, Jonghyun Lee, Jongho Park, Dongmin Park

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

想象一下，你正在尝试解决一个复杂的谜题，比如数独或一道棘手的数学题。你有一支专家团队（一个计算机模型）试图找出答案。

在 AI 语言模型的世界里，这些团队主要有两种工作方式：

“单次通过”团队（自回归式）： 他们从左到右阅读谜题，逐个猜测下一个词。如果他们在早期犯错，后期就很难修正。
“草稿纸”团队（掩码扩散式）： 他们从一个所有答案都被隐藏（掩码）的谜题开始。他们一次性猜测所有隐藏位置，检查自己的工作，然后修正猜测。他们重复这种“猜测 - 检查”循环，直到谜题被解决。

这篇论文为“草稿纸”团队引入了一种新技巧，称为LoopMDM。

问题：工作量太大，算力不足

通常，为了让“草稿纸”团队更聪明，你必须雇佣更多专家（在神经网络中添加更多层）。但雇佣更多人既昂贵又缓慢。研究人员问道：我们能否在不雇佣更多人的情况下，让现有团队变得更聪明？

解决方案：“循环”策略

作者们意识到，在解谜过程中，团队并不需要搬进新房间（网络的新一层）。相反，他们可以留在中间房间，来回穿梭，反复完善他们的想法。

这就像一群侦探在破案：

标准模型： 侦探 A 查看线索，将文件传给侦探 B，B 再传给 C。一旦文件离开他们的手，他们就无法改变主意。
LoopMDM： 侦探 A 查看线索，然后循环回去再次查看，再循环回去第三次以双重检查逻辑，然后才将文件传给侦探 B。

他们称之为**“选择性循环”。他们不会让每个侦探都循环回去；他们只让处于过程中间**的侦探这样做。这是团队需要深入思考但尚未确定最终答案的甜蜜点。

他们的发现（结果）

论文声称，这个简单的技巧威力巨大：

更便宜的训练： 因为团队重复使用同一个“中间房间”而不是建造新房间，他们可以用少 3.3 倍的计算量（FLOPs）来训练模型，达到与标准模型相同的技能水平。这就像通过调校现有引擎获得法拉利引擎，而不是购买一辆新车。
更聪明的数学： 在数学问题（如 GSM8K 基准测试）上，与同规模的标准模型相比，该方法将准确率提高了8.5 个百分点。它甚至击败了那些物理上更大（更深）但未使用此循环技巧的模型。
灵活的速度： 你可以通过改变循环次数来控制模型的“聪明”程度。
- 需要快速答案？循环一次。
- 需要解决难题的完美答案？循环 12 次。
- 你让模型在那个中间循环中“思考”得越多，模型就越出色。

为什么有效？（“原因”）

研究人员使用了一个“数独”实验来证明其有效性。

当他们强制模型按严格顺序（像读书一样）解决数独时，标准模型惨败。
但LoopMDM可以完美解决。为什么？因为“循环”允许模型将隐藏位置（空方格）视为共享工作空间。
随着模型循环，隐藏位置会“互相交流”。如果一个猜测是错误的，循环允许整个团队在确定最终答案之前共同意识到并修正它。这就像拥有一块白板，团队可以一起擦除和重写想法，而不是仅仅沿着一条线喊出答案。

“自适应”优势

论文还提出了一种使用此方法的智能方式：自适应循环。
与其强制模型对每个词都正好循环 12 次，模型可以检查自己的置信度。

如果答案显而易见，它循环一次然后继续。
如果答案棘手，它循环 12 次。
这在保持高质量的同时节省了能量。

总结

LoopMDM是一种让 AI 模型更聪明、更快速的方法，方法是让它们在处理过程的中间“循环思考”。与其建造一个更大、更昂贵的大脑，他们只是让现有大脑多花几秒钟来双重检查自己的工作。结果是一个学习更快、解决难题更好、且耗电更少的模型。

技术摘要：循环扩散语言模型（LoopMDM）

问题陈述

掩码扩散模型（MDMs）已成为自回归模型（ARMs）在语言建模方面的有力竞争者，通过训练目标、噪声调度和架构的进步缩小了性能差距。然而，专门针对 MDM 设计有效的 Transformer 架构仍未得到充分探索。虽然 ARMs 已成功利用循环 Transformer（重复应用共享块，将深度转化为循环）来实现参数效率和测试时计算扩展，但这种方法在 MDM 中的效用尚不明确。MDM 本身已在去噪步骤中执行迭代计算；目前尚不确定增加内部循环是否能带来额外收益，还是仅仅重复了现有机制。此外，以往关于循环的研究主要集中在参数效率或推理时优化上，而在 MDM 框架内如何提升训练计算效率以及在固定参数数量下的性能方面，仍存在理解上的空白。

方法论：LoopMDM

作者提出了LoopMDM（循环掩码扩散模型），这是一种将选择性循环集成到 MDM 去噪网络中的架构。该设计遵循两个原则：

选择性应用：循环次数选择性地应用于特定层，这些层在单位训练 FLOPs 下通过迭代优化能获得最高收益，而非均匀应用于整个网络。
可扩展性：该模型支持随循环次数（ $S$ ）的有效扩展，随着推理时 $S$ 的增加，性能呈现单调提升。

架构

去噪 Transformer 被划分为三个顺序组：

头部（Head）： $n_h$ 层，每次前向传播应用一次（输入嵌入）。
中间块（Mid-block）： $n_m$ 层，构成循环堆栈。这些层共享权重，并连续应用 $S$ 次。
尾部（Tail）： $n_o$ 层，应用一次（输出投影和语言模型头）。

前向传播展开后的有效深度为 $n_h + S \cdot n_m + n_o$ 层，而参数量与具有 $n_h + n_m + n_o$ 层的非循环网络相同。

训练策略

随机循环次数：在训练期间，循环次数 $S$ 在每一步从 $\{1, \dots, S_{max}\}$ 中均匀采样。这使得共享的中间块暴露于一系列有效深度之下，确保跨迭代次数的表示稳定性，并支持以任意 $S$ 值（包括超过 $S_{max}$ 的值）进行部署。
计算匹配：为确保公平比较，LoopMDM 的训练步数按比例少于基线模型，以匹配总训练 FLOPs。这将架构差异与原始计算量的增加隔离开来。

自适应推理

本文提出了一种自适应策略，其中循环次数不是固定的，而是在采样过程中动态确定。当连续循环迭代之间的相对隐藏状态变化低于阈值 $\epsilon$ ，或达到最大预算时，过程停止。这将更多计算分配给更新幅度较大的时间步（通常是中间噪声水平），而分配给稳定时间步的计算则较少。

主要贡献与结果

1. 训练计算效率

在三个预训练语料库（FineWeb-Edu、OpenWebText、LM1B）上，LoopMDM 实现了与同规模非循环 MDM 相当或更优的测试负对数似然（NLL），使用的训练 FLOPs 最多减少 3.3 倍。

在匹配总 FLOPs 进行训练时，LoopMDM 在零样本困惑度、生成困惑度和下游任务准确率方面始终优于基线模型。
消融研究结果：收益高度依赖于在哪里以及循环多少。选择性地循环少量早期中间层（具体为零基索引中的第 1–2 层）即可捕获大部分收益。循环整个网络或将循环置于输入/输出层附近会降低性能或毫无优势。

2. 超越深度扩展的性能

在 GSM8K 数学推理基准测试上，LoopMDM 比同规模的非循环基线提高了 +8.5 分。

关键在于，当匹配每步训练 FLOPs 时，LoopMDM 的表现优于更深的非循环 MDM（例如 21 层模型）。这表明收益源于共享计算的重复应用（迭代优化），而不仅仅是增加模型深度。
尽管由于匹配 FLOPs 所需的训练步数较少而导致看到的总训练 token 较少，LoopMDM 仍保持或提升了下游知识迁移能力，这表明在循环上花费的每步计算补偿了 token 暴露的减少。

3. 机制洞察

掩码位置作为工作空间：通过具有受限（从左到右）生成顺序的数独实验，作者表明 LoopMDM 可以解决浅层基线无法解决的全局约束任务。模型利用掩码位置作为并行工作空间，来修正早期局部合理但全局不一致的预测。
注意力动态：分析显示，掩码到掩码的注意力随着循环次数的增加而增强，特别是在中间时间步。这表明循环促进了掩码位置之间的交互，使模型能够在承诺 token 预测之前优化隐藏状态。
自适应效率：自适应循环策略将平均循环次数从 12 次减少到约 5 次，同时保持性能，证明了并非所有时间步都需要相同的计算预算。

意义与主张

本文主张选择性循环是一种简单而强大的架构选择，可用于改进掩码扩散模型。其意义在于证明了：

效率：MDM 可以在不增加参数数量的情况下实现显著的训练效率提升（FLOPs 减少高达 3.3 倍）。
质量：循环带来的收益不仅仅是深度增加的代理；它们源于迭代计算与扩散模型固有的掩码工作空间之间的特定交互。
可扩展性：该方法支持推理时灵活的计算扩展，自适应策略提供了进一步的效率。

作者将这项工作定位为迈向 MDM 更好架构设计的一步，表明选择性循环可以作为在固定参数和计算预算下提高训练效率和泛化能力的标准方法。他们承认了局限性，指出当前设计依赖于手动指定循环位置，未来的工作可以探索动态学习这些策略或将该方法扩展到十亿级规模的模型。