d2: Improving Reasoning in Diffusion Language Models via Trajectory Likelihood Estimation

该论文介绍了 d2，一种用于掩码扩散语言模型的创新强化学习框架，该框架采用了专门的轨迹似然估计器（d2-AnyOrder 和 d2-StepMerge），显著增强了在逻辑和数学基准测试上的推理能力，并取得了新的最先进性能。

要点

想象一下，你有一位才华横溢的艺术家，他可以通过从一张空白画布开始，逐渐添加细节，直到完成一幅杰作。这就是**扩散语言模型（Diffusion Language Models, DLMs）**的工作方式。与传统的 AI 作家（像堆砖块一样一个词一个词地构建句子）不同，这些模型从一团混乱的“噪声”（就像旧电视上的静电雪花）开始，通过逐步“去噪”，最终揭示出一个连贯的故事或解决方案。

这篇论文介绍了一种名为 d2 的新训练方法，旨在教导这些艺术家如何更好地解决复杂的难题（如数学问题或逻辑游戏）。以下是其工作原理的详细拆za，使用了简单的类比。

问题所在：“盲目”的评论家

为了教会 AI 进行推理，研究人员使用了强化学习（Reinforcement Learning, RL）。你可以把它想象成一场游戏，AI 试图解开谜题，而一个“评论家”（奖励系统）会给它评分。如果 AI 得到了高分，它就会学习再次这样做。

然而，这里有一个陷阱，专门针对扩散模型。为了进行有效学习，AI 需要知道生成该特定单词序列的“可能性”究竟有多大。

• 对于传统 AI（自回归模型）： 这就像写信。你知道你先写了哪个词，然后是第二个、第三个。计算这种“可能性”既简单又快速。
• 对于扩散 AI： 这就像观察一座雕塑如何从一块原石中显现出来。AI 会进行成千上上的次微调。要准确知道最终雕塑的概率，理论上你必须倒带视频，并重放雕塑过程中的每一帧。为每一个训练步骤都这样做，在计算上是不可能的（太慢且太昂贵）。

因为无法准确计算这个“分数”，以往的方法只能靠猜测，导致推理能力较差。

解决方案：d2（聪明的计分员）

作者创建了 d2，这是一个框架，充当了一个超级高效的计分员。它能在无需每次都重放整个电影的情况下，计算出 AI 旅程的“可能性”。他们为两种不同的艺术家打造了两种不同的工具：

工具 1：d2-AnyOrder（“魔镜”）

有些扩散模型非常特殊；它们可以以任何顺序揭示最终图像。你可以先画眼睛，再画鼻子，最后画背景，或者反过来，这都没问题。

• 类比： 想象你有一面魔镜。你不需要逐帧观看绘画过程，只需看一眼完成后的画作，就能瞬间看到艺术家在每一步中所表达的内容，且是一次性看清所有步骤。
• 结果： 这个被称为 d2-AnyOrder 的工具，只需通过一次计算机运行（single computer pass）即可计算出完美的分数。它极其精确，能让 AI 学习得非常快。

工具 2：d2-StepMerge（“分块策略”）

大多数流行的扩散模型（例如著名的 LLaDA）不能以任何顺序绘画；它们必须遵循特定的序列。它们没有那面“魔镜”。

• 类比： 既然我们无法一次性看完整部电影，我们就把电影分成若干块（chunks）。我们不再观看每一帧，而是以 10 秒为一个片段来观看。我们估算每个片段的分数，然后将它们相加。
• 结果： 这个被称为 d2-StepMerge 的工具是一种近似值。它并不完美，但作者在数学上证明了其误差是微小且可控的。它用极小的精度损失换取了巨大的速度提升，使其对标准模型具有实用价值。

结果：从新手到大师

研究人员在一些对 AI 最难的谜题上测试了这些工具：

1. 逻辑游戏： 数独（Sudoku）和 Countdown（一个数学游戏）。
2. 数学问题： GSM8K 和 MATH500（标准的数学推理基准测试）。

最终成果：

• 无需任何额外的“小抄”（监督微调/Supervised Fine-Tuning）： d2 框架教导模型的推理能力比以往任何方法都更强。
• 击败竞争对手： 在数独和 Countdown 测试中，经过 d2 训练的模型分别达到了接近 92% 和 56% 的得分，彻底碾压了此前处于 22% 和 42% 左右的现有最佳方法。
• 数学精通： 在复杂的数学基准测试中，d2 创下了新的“最先进水平”（State-of-the-Art）记录，证明了扩散模型可以像传统的文本生成模型一样具备优秀的推理能力，同时还拥有更快、更具可控性的优势。

总结

论文的核心观点是：“我们找到了一种教导扩散 AI 模型清晰思考的方法。”

• 他们意识到，旧的评估 AI 作品的方式要么太慢，要么不准确。
• 他们发明了 d2，一种全新的评分系统，它要么是完美精确的（针对特殊模型），要么是聪明高效的（针对标准模型）。
• 利用这个新系统，AI 模型在不需要预先学习答案的情况下，显著提升了解决逻辑和数学难题的能力。

简而言之，d2 让扩散模型获得了能够准确“回顾”自身思考过程的能力，从而使它们能够从错误中学习，成为真正的推理专家。

技术摘要

技术摘要：d2：通过轨迹似然估计提升扩散语言模型的推理能力

问题陈述

扩散语言模型（DLMs），特别是掩码扩散语言模型（Masked DLMs），已成为自回归（AR）模型之外具有竞争力的替代方案，具备可控性和并行生成能力。然而，通过强化学习（RL）来增强其推理能力仍然具有挑战性。虽然 RL 已成为诱导自回归模型进行推理的标准方法，但将其应用于 DLM 则并非易事。核心难点在于策略梯度公式化：与自回归模型中似然度在 Token 位置上可以清晰分解不同，扩散轨迹的精确似然计算在计算上是难以实现的。天真地计算这些似然度需要 $T$ 次前向传播（其中 $T$ 为扩散步数），这使得像组相对策略优化（GRPO）这样的标准策略梯度方法在计算上变得极其昂贵。现有的近似方法通常会引入显著偏差，导致次优的策略更新。

方法论

作者提出了 d2，一个专门为掩码 DLM 设计的推理框架。该框架的核心是一个源自 GRPO 目标的全新策略梯度算法，该算法明确要求对采样轨迹似然进行准确估计。为了解决计算瓶颈，论文针对不同的模型类别提出了一系列估计器：

1. 理论基础

作者通过对时间潜变量进行边缘化并采用重要性采样，推导出了适用于掩码 DLM 的 GRPO 式目标函数。由此产生的梯度分解（定理 3.1）取决于当前策略 $\pi_\theta$ 与旧策略 $\pi_{old}$ 之间的轨迹似然比。关键挑战在于如何高效地评估 $\pi(x_{0:T})$。

2. 估计器

d2 框架根据模型架构提供了两种特定的估计器：

• d2-AnyOrder（精确估计器）：

  • 目标： 支持**任意顺序解码（AO-dLLMs）**的 DLM。这些模型允许以任何顺序解码 Token，只要注意力掩码遵循特定的因果约束（独立掩码和顺序因果性）。
  • 机制： 它构建一个长度为 $2L$ 的序列（拼接干净的 Token 和掩码 Token），并使用特定的位置编码和注意力掩码。这使得模型能够通过单次前向传播计算出所有 Token 的精确条件似然。
  • 局限性： 标准的掩码 DLM（如 LLaDA）原生并不支持任意顺序解码；它们的训练目标无法保证该估计器所需的注意力属性，从而保证其无偏性。

• d2-StepMerge（近似估计器）：

  • 目标： 不支持任意顺序解码的标准掩码扩散模型（MDMs）。
  • 机制： 受块组合似然（Block Composite Likelihood）的启发，该估计器通过将 $T$ 步轨迹划分为 $N$ 个连续的时间段来近似轨迹似然。它不是评估每一步，而是通过每个段内的一次前向传播来评估在该段内未被掩码的 Token 的似然度。
  • 权衡： 它将前向传播次数从 $T$ 减少到 $N$。论文为近似误差（KL 散度）提供了一个理论上界（定理 4.1），表明误差随 $N$ 的增加而单调递减，从而量化了计算量与偏差之间的权衡。

核心贡献

1. d2 框架： 一个针对掩码 DLM 的原则性 RL 后训练框架，它推导出了一个依赖于准确轨迹似然估计的 GRPO 目标。
2. d2-AnyOrder： 一种针对 AO-dLLMs 的精确、单次传递似然估计器，证明了无偏似然估计在 RL 设置中显著优于现有基线。
3. d2-StepMerge： 一种针对标准 MDMs 的确定性、多路径估计器，具有解析可界的近似误差，提供了比 diffu-GRPO 等方法更优的效率-精度权衡。
4. 实证验证： 大量实验表明，d2 在无需依赖监督思维链（CoT）微调的情况下，在推理基准测试上达到了最先进的性能。

实验结果

作者在逻辑推理（Countdown, Sudoku）、数学推理（GSM8K, MATH500）基准测试以及编程任务和毒性控制方面对 d2 进行了评估。

• 在 AO-dLLMs 上的表现： 当应用于支持任意顺序解码的模型（如 Eso-LM、Set Diffusion 以及自定义的 Any-Order Causal LLaDA）时，d2-AnyOrder 显著优于 DDPO 和 diffu-GRPO 等基线。例如，在基于 Qwen3-1.7B 的 GSM8K 测试中，d2-AnyOrder 实现了 67% 的准确率，而 diffu-GRPO 为 63%。
• 在标准 MDMs 上的表现： 应用于 LLaDA-8B-Instruct（一种标准掩码 DLM）时，d2-StepMerge（取 $N=16$）创下了新的 SOTA 结果：
  • Sudoku： 91.9%（对比 d1 使用 SFT 的 22.1%）。
  • Countdown： 56.6%（对比 d1 的 42.2%）。
  • GSM8K： 85.0%（对比 d1 的 82.1%）。
  • MATH500： 41.6%（对比 wd1 的 39.0%）。
    值得注意的是，这些结果是在没有额外 CoT 数据进行监督微调的情况下实现的，而某些基线则依赖于此类数据。
• 泛化能力： 该方法成功泛化到了其他架构（Dream 7B）和任务（HumanEval, MBPP），显示出比 diffu-GRPO 更强的长度泛化能力。
• 消融实验： 实验证实，精确似然（d2-AnyOrder）比近似似然能提供更高保真度的信号，且 d2-StepMerge 在最优 $N$ 值（如 16）下比极端值能更好地平衡计算成本与准确度。

意义与主张

论文声称，准确的轨迹似然估计是提升 DLM 推理能力的实现成功的核心。作者认为，以往的方法之所以无法达到最优性能，是因为它们依赖于有偏或计算昂贵的似然估计。

• 理论贡献： 本研究首次为掩码 DLM 提供了严谨的 GRPO 目标推导，并建立了关于步长合并策略近似误差的理论界限。
• 实际影响： d2 证明了通过适当的 RL 训练，扩散模型可以获得足以媲美甚至超越自回归模型的推理能力，且无需大量的监督 CoT 数据。
• 未来展望： 作者指出，虽然 d2-AnyOrder 表现更优，但其适用性目前受限于专门为任意顺序解码训练的模型。他们认为，开发更强大的通用型任意顺序 dLLM 是充分利用该精确估计器的自然下一步。

论文总结道，d2 在逻辑和数学推理任务上为 DLM 树立了新的基准，验证了其似然估计策略在释放基于扩散的语言模型推理潜力方面的有效性。