💬 NLP

Mask-Aware Policy Gradients for Diffusion Language Models

本文介绍了掩码感知策略梯度（Mask-Aware Policy Gradients），这是一种将掩码扩散语言模型生成形式化为两阶段决策过程的强化学习框架，通过联合优化标记选择与位置掩码，从而克服了对数似然不可解的问题，并在数学推理和编程基准测试中达到了最先进的性能。

原作者： Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl

发布于 2026-07-17

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl

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

想象一下，你正试图教一个机器人如何写故事或解决数学题。长期以来，教机器人的最佳方式是让它像人类在键盘上打字一样，一次写一个词。这被称为“自回归”（autoregressive）生成。这种方法很可靠，但也很慢，就像蜗牛横穿高速公路一样。最近，科学家们发现了一种更快的办法，叫做“扩散”（Diffusion）。扩散模型不再是从头开始书写，而是从一页写满空白掩码（比如 [MASK]）的页面开始，逐渐填补这些空格，逐一揭示单词，直到句子完整。这就像一个拼图游戏，你从一张被遮盖的图片开始，然后慢慢揭开遮盖物，露出图像。

然而，通过奖励（一种被称为“强化学习”的技术）来教这些快速的“扩散”机器人变得越来越聪明，一直是一场噩梦。在旧有的“逐词”方法中，很容易精确计算出机器人做出某个特定选择的可能性。但在“扩散”方法中，机器人在每一步都会做出两个选择：它既要决定在空白处填入什么词，也要决定下一步要揭开哪个空白处。以往的方法试图通过只关注第一个选择（单词）来教机器人，而忽略了第二个选择（揭开的顺序）。这就像是在教一位厨师烹饪完美佳肴时，只通过评价他们挑选的食材来给予批评，却完全忽略了他们切菜和搅拌的顺序。

这篇发表在 COLM 2026 会议上的论文指出，我们一直遗漏了拼图中的一个巨大部分。作者 Haran Raajesh、Kulin Shah 以及他们在德克萨斯大学奥斯汀分校的团队认为，要真正掌握扩散艺术，你必须同时奖励机器人对于“选对词”和“选对揭开顺序”的表现。他们开发了一种名为“掩码感知策略梯度”（Mask-Aware Policy Gradients）的新方法。该系统不再仅仅是猜测下一个词，而是将“下一个要揭开哪个掩码”的决策视为机器人策略中至关重要的一部分。通过在数学上将学习过程分解为这两个截然不同的部分，他们发现机器人学习得更快，且犯错更少。

实验结果非常令人印象深刻。当他们用这种新方法测试困难的数学题和编程挑战时，机器人的解题能力显著提升。在一项名为 GSM8K 的著名数学测试中，他们的方法达到了 87.1% 的准确率；而在一项编程测试 MBPP 中，达到了 53.4%。这些数字高于以往任何尝试利用奖励来训练扩散模型的方法。这篇论文表明，通过关注“掩码”决策——即机器人揭示其思路的顺序——我们可以解锁这些快速、灵活的 AI 模型的新高度，使它们不仅速度更快，而且更加聪明。

技术摘要：面向扩散语言模型的掩码感知策略梯度

问题陈述

强化学习（RL）通过基于可处理对数似然值的策略梯度优化，成功提升了自回归大语言模型（LLMs）的推理能力。然而，将这些方法扩展到**掩码扩散语言模型（MDLMs）**仍面临挑战。与自回归模型从左至右顺序生成 Token 不同，MDLM 通过迭代地从全掩码序列中取消掩码（unmasking）来生成文本。

核心难点在于 MDLM 对数似限估计的不可处理性。现有方法通常使用证据下界（ELBO）变体或仅对生成轨迹中的 Token 预测部分进行建模来近似该似然值。这些方法忽略了 MDLM 生成过程中的一个关键方面：位置取消掩码的顺序。在 MDLM 中，生成轨迹涉及两个截然不同的决策：（1）预测在掩码位置放置什么 Token，以及（2）决定保留哪些位置不被掩码（以及哪些重新变为掩码）。此前忽略第二个决策的方法无法利用生成过程中可用于策略优化的完整信号。

方法论

作者将 MDLM 生成形式化为一个两阶段动作马尔可夫决策过程（MDP）。在每个去噪步骤 $t$ ，模型执行以下操作：

Token 预测： 为当前所有掩码位置预测 Token。
位置选择（重掩码）： 决定哪些预测的位置将被揭示，以及哪些将返回 [MASK] 状态。

为了实现对该过程的策略梯度优化，作者引入了概率性重掩码（Probabilistic Remasking）。标准的 MDLM 推理通常使用确定性的、贪婪的 Top- $K$ 策略，根据置信度得分选择要取消掩码的位置。这种贪密选择是不可微的，从而阻碍了对取消掩码顺序的基于梯度的优化。所提出的方法使用其变体取代了这一过程：

模型不再确定性地选择前 $K$ 个位置，而是从 Plackett–Luce 分布中采样一个位置子集 $U_t$ 。
该分布的概率直接由模型的 Token 对数似然（logits）推导而来，并通过温度参数 $\tau$ 进行缩放。
这种形式化创建了一个关于位置子集的微分分布，使得取消掩码决策的对数概率能够被纳入轨迹似然中。

策略梯度分解

通过将全轨迹似然 $\pi_\theta(\hat{z}|c)$ 建模为 Token 预测概率与取消掩码概率的乘积，策略梯度自然分解为两项：
$\nabla_\theta J(\theta) = \mathbb{E}_{\hat{z}} \left[ R(c, z_T) \sum_{t=1}^T \left( \underbrace{\nabla_\theta \log \pi_\theta(\hat{z}_t | c, z_{t-1})}_{\text{Token 梯度}} + \underbrace{\nabla_\theta \log p_{\text{unmask}}(U_t | \hat{z}_t, z_{t-1}, c)}_{\text{Masking 梯度}} \right) \right]$

Token 项： 强化预测 Token 的正确性。
Masking 项： 强化对取消掩码顺序的选择。
效率： Masking 项是从模型现有的用于 Token 预测的 logits 中计算得出的，无需额外的参数、架构更改或额外的向前传播过程。

作者使用 组序列策略优化（GSPO） 算法实现此方法，分别对 Token 项和位置项的归一化重要性比率进行裁剪。为了处理评估全轨迹的计算成本，他们采用了 StepMerge 近似法，将去噪步骤分组为段，以在减少向前传播次数的同时保留轨迹结构。

核心贡献

将 MDLM 形式化为两阶段 MDP： 论文提供了一个理论框架，证明了 MDLM 生成涉及 Token 选择和位置选择，且策略梯度据此进行分解。
概率性重掩码： 推导出了一个基于模型自身 logits 的可微取消掩码分布（Plackett–Luce），从而能够在不修改架构的情况下优化取消掩码顺序。
理论洞察： 证明了忽略策略梯度中的位置组件会导致无法捕捉到能提升预期回报的方向，即使 Token 概率保持不变也是如此。
最先进性能（SOTA）： 实证验证表明，联合优化这两项可以比现有的基于轨迹和基于 ELBO 的方法获得一致的改进。

实验结果

该方法在数学推理（GSM8K, MATH500）和代码生成（HumanEval, MBPP）基准测试上进行了评估，使用 LLaDA-8B-Instruct 作为基础模型。

性能增益： 所提方法在所有基准测试中均取得了新的 SOTA 结果。
- GSM8K： 准确率为 87.1%（生成长度为 128），比之前的最优方法（StepMerge）提升了 2.5%。
- MATH500： 准确率为 53.4%，提升了 4.0%。
- HumanEval： 准确率为 47.1%，提升了 2.2%。
- MBPP： 准确率为 53.4%，提升了 2.5%。
消融研究：
- 块大小（Block Size）： 随着推理块大小增加（从 32 到全序列），所提方法与基准方法之间的性能差距扩大，这表明优化位置选择的收益随取消掩码决策空间的复杂性而增长。
- 效率： 与强大的基于 ELBO 的基准方法（SPG）相比，该方法在实际运行时间（wall-clock time）上收敛更快，在约 15 小时内达到了 SPG 的最终准确率（SPG 约为 18 小时），并实现了更高的最终准确率（GSM8K 上为 81.0% vs 78.5%）。
- 泛化性： 该方法泛化到了不同的基础模型（Dream-7B）和规划任务（Sudoku, Countdown），并始终优于基准方法。
与 DCoLT 的比较： 与使用独立位置选择头的 LLaDOU 相比，所提方法在显著降低训练成本（约 160 GPU 小时 vs 800 GPU 小时）的同时，实现了相当或更好的准确率，因为它不需要额外的模块。

重要性与主张

论文声称，MDLM 的灵活性带来了比自回归模型更丰富的轨迹结构，即决定何时取消掩码位置。通过将此决策视为策略的一部分，并通过可微概率框架对其进行优化，作者证明了联合优化 Token 预测和位置选择对于最大化扩散语言模型在 RL 环境下的性能至关重要。

作者强调，该方法在实践上是高效的，因为它利用现有的模型 logits，不需要新的参数或架构更改，使其能够兼容标准的 MDLM 训练基础设施（如 StepMerge）。结果表明，“掩码梯度”提供了一个有意义的训练信号，而此前仅关注 Token 似然的方法未能捕捉到这一点。