🤖 machine learning

Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization

本文介绍了分布匹配策略优化（DMPO），这是一种新颖的强化学习框架，通过将策略分布与最优奖励倾斜的目标对齐，显著增强了扩散大语言模型的推理能力，在无需监督微调的情况下，实现了相对于现有基准模型的实质性准确率提升。

原作者： Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, Yongxin Chen

发布于 2026-07-14

📖 1 分钟阅读☕ 轻松阅读

原作者： Yuchen Zhu, Wei Guo, Jaemoo Choi, Petr Molodyk, Bo Yuan, Molei Tao, Yongxin Chen

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

想象一下，你有一个超级聪明的机器人，它能写故事、解数学题，还能玩游戏。如今大多数机器人就像是一个正在读书的人，一次只能读一个词，从左到右。它们无法跳着看或者回头看；它们只能根据之前的词来猜测下一个词。这很慢，就像试图通过只看眼前的一块拼图碎片来完成一个巨大的拼图。

扩散大语言模型 (Diffusion Large Language Models, dLLMs) 正是为此而生的。它们是这个领域的新成员。这些模型不再是一个词一个词地写作，而是从一团混乱、被搅碎的“神秘标记”（就像一个每个碎片都被迷雾覆盖的拼图）开始，然后慢慢清理，从而一次性或以任何顺序揭示出答案。这使得它们在思考方面潜力巨大，速度更快。

但问题在于：虽然这些扩散机器人很快，但在处理像高级数学或逻辑谜题这类棘手的推理任务时，它们并不总是最聪明的。为了让它们变得更聪明，科学家通常会使用一种叫做强化学习 (Reinforcement Learning, RL) 的技术。你可以把它想象成一个电子游戏，机器人给出好答案会得分，给出坏答案会扣分。

问题所在：“寻模”陷阱 (The "Mode-Seeking" Trap)

旧的训练这些机器人的方法（使用像 GRPO 这样的方法）有点像一个只钻研自己认为正确的那个答案的学生。如果机器人发现了一种能获得分数但感觉很“稳妥”的方法，它就会停止探索。它会陷入一种僵局，忽略掉其他同样优秀的巧妙解决方案。在论文中，作者称之为**“寻模” (mode-seeking)**。这就像一个机器人只学会了用一种特定的方式来解决数学问题，一旦这种方式被堵死了，它就会陷入恐慌。它也倾向于忽略那些看起来与第一个找到的答案不同、但依然正确的“杂乱”答案。

论文指出，仅仅追求最高分这种旧方法对于扩散模型来说是有缺陷的。它认为，通过只关注单一的最佳路径，我们会错失扩散模型同时探索多条不同路径的独特能力。

解决方案：DMPO（“地图匹配型”机器人）

作者提出了一种名为分布匹配策略优化 (Distribution Matching Policy Optimization, DMPO) 的新方法。

DMPO 不再告诉机器人：“只需找到最高分。”它说的是：“这是你可以给出的所有好答案的完整地图。你的任务是学会匹配这张完整的地图。”

想象你在教一只狗去捡球：

旧的方法： 你扔出一个球，狗就跑向球落下的那一个点。如果球落在灌木丛里，狗就会学会只去灌木丛里捡球。
DMPO 的方法： 你向狗展示一张包含球可能落下的所有位置的地图（灌木丛、草地、沙地、水里），并对它说：“学会从任何这些地方去捡球，并根据该位置的好坏进行加权。”这样，狗学会了在整个院子里灵活探索，而不是只盯着一个角落。

在技术层面，DMPO 使用了一种特殊的数学技巧，称为加权去噪交叉熵 (Weighted Denoising Cross-Entropy, WDCE)。这使得机器人可以从过去的尝试中学习（即使那些尝试并不完美），而无需每次都从头重新生成所有内容。这就像拥有一个“经验回放池”，机器人可以反复研究它玩过的旧游戏，从好的和坏的动作中学习，而不会感到困惑。

核心秘诀：“权重基准” (The "Weight Baseline")

作者发现，在使用小规模样本组（小批次大小）进行训练时，存在一个棘手的问题。有时，机器人会因为看到的例子太少，而错误地奖励一些糟糕的答案。

为了解决这个问题，他们发明了一个聪明的技巧，叫做权重基准减法 (Weight Baseline Subtraction)。

这就像一位老师在批改试卷。如果学生答对了问题，他会得到一颗金星。但如果老师只看到了这一个问题，他可能会认为学生做的所有事情都很棒。而“基准”就像是一个“平均水平”。老师会从学生的得分中减去这个“平均值”。

如果学生的表现优于平均水平，他会得到一颗大大的金星。
如果他的表现差于平均水平，即使他确实拿到了一些分数，也会受到“惩罚”（一个负权重）。

这确保了机器人不会对平庸的答案过度兴奋，并始终朝着真正卓越的方向努力。

实验结果：提升了多少？

作者在一些非常困难的推理基准测试上测试了这种新方法。他们将 DMPO 应用于预训练模型（如 LLaDA-Instruct 和 Dream-Instruct），采用了两种不同的方式：

直接应用： 他们直接将 DMPO 应用于基础模型，而没有进行任何先前的“课后作业”（即针对推理数据集的有监督微调或 SFT）。这种类似于 “R1-Zero” 的方法是为了清晰地展示 DMPO 的原始潜力。
增强应用： 他们还将 DMPO 应用于已经过 SFT 处理的模型，以展示它作为现有模型强大升级方案的效果。

实验结果在两种设置下都非常令人印象深刻：

在一个名为 GSM8K 的数学谜题上，直接应用于预训练基础模型（未经过 SFT）的模型展现出了巨大的进步。具体来说，相比于之前的非 DMPO 强化学习基准，DMPO 实现了高达 39.63 个百分点的准确率提升；而相比于基础模型本身，提升幅度更是达到了惊人的 67.97 个百分点。
在一个名为 数独 (Sudoku) 的逻辑谜题上，他们看到了巨大的飞跃，其中一个版本的模型成功率从 16.41% 跳升到了 80.86%（差距超过 64 个百分点！）。
至关重要的是，即使应用于已经经过 SFT 的模型，DMPO 依然能带来显著的性能增益，这证明它是一种鲁棒的方法，无论你是从零开始还是在现有训练基础上构建，它都行之有效。

论文表明，DMPO 不仅仅是一个小小的改进，而是一个根本性的转变。它允许机器人实现离策 (off-policy) 学习，这意味着它可以高效地从旧数据中学习；同时也实现了仅前向 (forward-only) 过程，意味着它不需要进行会拖慢其他机器人的昂贵反向计算。

总结

作者确信，这种方法对于他们测试的特定模型（如 LLaDA 和 Dream）以及这些特定的推理任务非常有效。他们通过在标准数据集上的严格实验测量了这些改进，证明了无论是在刚完成预训练的模型上，还是在已经经过有监督微调的模型上，DMPO 都能提升性能。然而，他们也承认，目前尚未在所有可能的模型类型或任务上进行测试，因此虽然结果非常强劲，但其全部潜力仍有待进一步探索。

简而言之，DMPO 教会了扩散机器人不要死盯着单一的“完美”答案，而是学会欣赏所有优秀解决方案的整体景观，从而使它们成为更聪明、更快速、更有创造力的解题者。

技术摘要：通过分布匹配策略优化增强扩散大语言模型（Diffusion LLMs）的推理能力

问题陈述

扩散大语言模型（dLLMs）通过实现任意顺序生成和潜在更高的推理吞吐量，为自回归大语言模型（AR-LLMs）提供了一种极具前景的替代方案。然而，将 dLLMs 的规模扩展到足以匹配 AR-LLMs 的推理能力仍面临重大挑战。现有的强化学习（RL）微调方法，例如组相对策略优化（GRPO），由于以下两个主要因素而不适用于 dLLMs：

计算不匹配： dLLMs 依赖于双向处理，这使得估计序列对数概率（标准策略梯度方法所必需）在计算上变得昂贵且往往难以实现，这与 AR 模型相比具有显著差异。
算法局限性： 当前针对 dLLMs 的 RL 方法通常依赖于启发式近似（例如均值场分解），这些方法忽略了 Token 位置之间的相关性。此外，它们通常侧重于奖励最大化，这倾向于优化单一的奖励模态。这种“寻模行为”（mode-seeking behavior）会导致奖励作弊（reward hacking），并导致无法探索 dLLMs 理论上能够生成的各种高质量推理路径。

方法论：分布匹配策略优化 (DMPO)

本文提出了分布匹配策略优化（DMPO），这是一个基于随机最优控制（Stochastic Optimal Control, SOC）的有原则的 RL 微调框架。与传统的最大化期望奖励的方法不同，DMPO 旨在使模型的策略分布与最优的、经过奖励倾斜（reward-tilted）的目标分布相匹配。

核心框架

目标是学习一个策略 $\pi_\theta$ ，使得诱导的序列分布 $p_\theta(o|q)$ 与最优分布 $p^*(o|q)$ 相匹配，其中：
$p^*(o|q) = \frac{1}{Z(q)} p_{ref}(o|q) e^{r(q,o)/\alpha}$
这里 $p_{ref}$ 是预训练策略， $r$ 是奖励函数， $\alpha$ 是温度参数。

为了实现这一目标，作者将问题表述为最小化目标分布与当前策略之间的前向 Kullback-Leibler (KL) 散度（等价于交叉熵）：
$\min_\theta \mathbb{E}_{q \sim D} \mathbb{E}_{p^*(o|q)} \left[ \log \frac{p^*(o|q)}{p_\theta(o|q)} \right]$
这种“覆盖质量”（mass-covering）的目标鼓励模型探索多样的高奖励轨迹，而不是坍缩到单一模态。

关键技术组件

加权去噪交叉熵 (WDCE)：
由于 $p^*$ 是不可计算的（由于未知的配分函数 $Z(q)$ 以及无法直接从其进行采样），DMPO 采用了重要性采样。它将来自参考策略 $p_v$ （模型的一个停止梯度副本）的样本视为来自 $p^*$ 的加权样本。
由此产生的损失函数 WDCE 利用了 dLLMs 特有的前向（加噪）过程。它仅使用最终的干净样本及其相关的重要性权重来计算损失，丢弃了中间的展开轨迹。这使得 WDCE 成为一个离策（off-policy）且仅前向的目标。
权重基线减法 (Weight Baseline Subtraction)：
作者指出，在使用小训练批次时，使用 WDCE 会出现不稳定的关键实现挑战。在小批次中，对于“坏”响应（尽管在 $p^*$ 下仍有非零概率）的正权重可能会占据主导地位，从而导致错误的梯度更新。
为了解决这个问题，作者提出使用权重基线减法，即从重要性权重中减去一个基线权重 $w_{base}$ ：
$w_{real} = w - w_{base}$
作者提出了三种 $w_{base}$ 策略：

组权重基线 (Group Weight Baseline)： 一个常数（1），以鼓励收敛。
个体权重基线 (Individual Weight Baseline)： 基于负奖励倾斜。
模型权重基线 (Model Weight Baseline)： 基于当前模型自身的权重。
该技术有效地注入了针对低奖励样本的负梯度，防止了在低批次场景下模型促进坏响应的情况。

加权直接判别优化 (WDDO)：
作为交叉熵的替代方案，论文探索了一种受 GAN 启发的优化目标，它能自然地整合针对坏样本的负梯度，从而进一步增强分布匹配。

核心贡献

新型 RL 框架： 提出了 DMPO，这是一个专门为 dLLMs 设计的分布匹配框架，将范式从奖励最大化转向匹配整个奖励倾斜的策略分布。
理论基础： 利用重要性采样和随机最优控制推导出了 WDCE 损失，为无需昂贵的序列对数概率估计的离策训练提供了理论依据。
实际解决方案： 识别了前向 KL 目标在小批次下的失效模式，并引入了权重基线减法和 WDDO 来稳定训练并提高样本效率。
高效性： 证明了 WDCE 能够仅通过前向过程进行训练，从而允许复用展开过程（off-policy），并与快速推理技术（如 KV-cache）兼容，与基于轨迹的方法相比显著加速了训练。

实验结果

作者在 LLaDA-8B-Instruct 和 Dream-7B 模型上针对四个推理基准测试进行了评估：GSM8K、MATH500、Countdown 和 Sudoku。

性能提升： DMPO 一致地优于基础模型和现有的 RL 基线（包括 diffu-GRPO、cGRPO 和 d1）。
- 在 Sudoku 上，DMPO-LLaDA 比表现最好的非 DMPO RL 基线提升了 39.63 个百分点，比基础模型提升了 67.97 个百分点。
- 在 Countdown 上，观察到相比强基线最高达 16.40 个百分点的提升。
无需 SFT： 该方法在无需使用精选数据集进行监督微调（SFT）的情况下实现了这些结果，展示了其直接从预训练模型扩展推理能力的能力。
训练效率： 由于其离策性质和仅前向计算，DMPO 实现了 2–3 倍更高的样本效率，并且与 GRPO 类方法相比，实现了高达 8 倍的展开采样速度（结合快速 dLLM 推理时）。它仅用 31% 的训练计算预算就达到了 d1 的奖励水平。
稳定性： 即使使用高度过时的（stale）数据（高达 24 倍重复使用同一个展开批次），该方法依然保持稳定，这一特性归功于 DMPO 中使用的序列级重要性权重的稳定性。

意义与主张

论文声称 DMPO 代表了在推理任务上扩展 dLLMs 能力的重要一步。通过摆脱奖励最大化范式和基于轨迹的 RL 计算瓶颈，DMPO 充分利用了扩散模型的独特特性：

多样性： 前向 KL 目标防止了模态坍缩，允许 dLLMs 探索多样化的推理路径。
高效性： WDCE 的离策、仅前向性质与 dLLMs 的高吞吐量推理能力相契合，为实现可扩展的 RL 微调提供了一条计算上可行的路径，适用于大型模型。
泛化性： 该框架被呈现为一种有原则的方法，可以应用于各种 dLLM 骨干网络，这一点从其在 LLaDA 和 Dream 架构上的成功得到了证实。

作者总结道，虽然目前的工作侧重于特定的模型和基准测试，但分布匹配框架为改进生成模型和设计更有效的权重基线技术开辟了新的方向。