💬 NLP

SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

该论文针对扩散语言模型因对数似然不可解而难以应用强化学习对齐的问题，提出了一种利用对数似然上下界构建的“夹心策略梯度”（SPG）方法，有效克服了现有单侧近似带来的偏差，并在 GSM8K、MATH500 等多个基准测试中显著提升了模型性能。

原作者： Chenyu Wang, Paria Rashidinejad, DiJia Su, Song Jiang, Sid Wang, Siyan Zhao, Cai Zhou, Shannon Zejiang Shen, Feiyu Chen, Tommi Jaakkola, Yuandong Tian, Bo Liu

发布于 2026-04-16

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Chenyu Wang, Paria Rashidinejad, DiJia Su, Song Jiang, Sid Wang, Siyan Zhao, Cai Zhou, Shannon Zejiang Shen, Feiyu Chen, Tommi Jaakkola, Yuandong Tian, Bo Liu

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

这篇论文介绍了一种名为 SPG（夹心策略梯度） 的新方法，旨在解决“扩散语言模型”（一种新型 AI）在通过“强化学习”（让 AI 从奖励中学习）变聪明时遇到的一个核心难题。

为了让你轻松理解，我们可以把整个过程想象成教一个正在玩“填字游戏”的 AI 如何解题。

1. 背景：两种不同的“填字”方式

首先，我们要了解两种 AI 写文章（或解题）的方式：

传统 AI（自回归模型）： 就像写日记。它必须一个字一个字地写，写完第一个才能写第二个。虽然慢，但它很清楚自己写了什么，也知道“这句话的总概率”是多少。
扩散模型（Diffusion Models）： 就像玩“大家来找茬”或“填字游戏”。它一开始把整篇文章都涂成黑块（全是 [mask]），然后一步步把黑块擦掉，猜出原本的字。
- 优点： 它可以并行工作，一下子擦掉好几个黑块，速度极快。
- 缺点： 因为它是一步步“猜”出来的，它很难算出“我最终猜出这句话的总概率”到底是多少（数学上叫“不可计算的对数似然”）。

2. 问题：当 AI 需要“奖惩”时，它迷路了

为了让 AI 更擅长做数学题或逻辑推理，我们需要用强化学习（RL）。这就好比给 AI 一个评分表：

如果它答对了，给奖励（+1 分）。
如果它答错了，给惩罚（-1 分）。

传统的做法（ELBO）：
以前的方法为了绕过“算不出总概率”的难题，只用了一个**“下限估计”**（ELBO）。

比喻： 想象你在教一个学生。
- 当他做对题时，你告诉他：“你现在的表现至少有 80 分，继续加油！”（这是下限，鼓励他）。
- 当他做错题时，你依然用这个逻辑：“你现在的表现至少有 80 分……"
- 问题出在哪？ 这里的逻辑是错的！如果你做错了，你希望分数越低越好（比如从 80 分降到 20 分），这样你才知道要改。但“下限估计”只能告诉你“至少 80 分”，它无法有效地告诉模型“你现在的表现太糟糕了，必须大幅降低概率”。这就导致模型在犯错时学不到东西，甚至越学越偏。

3. 解决方案：SPG（夹心策略梯度）

这篇论文提出的 SPG 方法，就像给 AI 老师配了一个**“三明治”评分系统**，专门用来处理“做对”和“做错”两种情况：

当 AI 做对题（奖励高）时：
- 策略： 使用**“下限估计”**（Lower Bound）。
- 比喻： “你做得很好，我们保证你的表现至少有这么好，请继续保持！”（最大化下限，鼓励它）。
当 AI 做错题（奖励低）时：
- 策略： 使用**“上限估计”**（Upper Bound）。
- 比喻： “你做得太差了，我们保证你的表现至多只有这么差，必须立刻改正！”（最小化上限，强力惩罚）。

为什么叫“夹心”？
因为真实的“总概率”就像夹心面包的馅料，它被夹在“下限”和“上限”这两片面包中间。

对于好答案，我们推“下限”往上走。
对于坏答案，我们推“上限”往下走。
这样，无论好坏，模型都能得到准确且方向正确的反馈，就像被两片面包紧紧夹住，被迫向正确的方向移动。

4. 关键技术：分块掩码（Block-wise Masking）

为了让这个“夹心”评分更准确，作者还设计了一个**“分块擦除”**的技巧。

以前的做法： 随机擦除几个字。这就像在填字游戏里，随机把几个字涂黑，AI 可能会感到困惑，因为上下文不连贯。
SPG 的做法： 把文章分成几块（比如每 32 个字一块）。在训练时，它保留前面的块是干净的，把后面的块全涂黑，只随机擦除当前正在处理的那一块。
比喻： 这就像老师教学生解题时，只遮住当前正在思考的那一行，而前面的步骤都清晰可见。这样 AI 在练习“猜字”时，看到的上下文更符合它实际生成时的样子，训练起来更稳定、更高效。

5. 成果：AI 变聪明了多少？

实验结果显示，使用 SPG 方法的 AI 在数学和逻辑推理任务上表现惊人：

GSM8K（小学数学题）： 准确率提升了 3.6%。
MATH500（高中数学题）： 提升了 2.6%。
Countdown（数字倒计时游戏）： 提升了 18.4%（这是一个巨大的飞跃！）。
Sudoku（数独）： 提升了 27.0%（几乎翻倍！）。

总结

简单来说，这篇论文解决了一个难题：如何让那种“一边擦除黑块一边猜字”的 AI，在通过“奖惩机制”学习时，不会因为算不清总分而学偏？

作者想出了一个**“夹心”策略**：

做对了，用下限鼓励它；
做错了，用上限惩罚它。
再配合**“分块擦除”**的练习方法，让 AI 在数学和逻辑推理上变得更强、更准。

这就好比给 AI 老师配了一副**“双向眼镜”**，既能看到它的高光时刻，也能精准地指出它的低谷时刻，从而让它进化得更快。

这是一份关于论文《SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models》（SPG：用于掩码扩散语言模型的夹心策略梯度）的详细技术总结。

1. 研究背景与问题 (Problem)

背景：
扩散大语言模型（dLLMs），特别是掩码扩散语言模型（MDLM），因其能够并行解码多个 token，在推理延迟和效率上展现出比自回归（AR）模型更大的潜力。然而，将 dLLMs 与人类偏好或特定任务奖励（如数学推理）对齐，通常需要通过强化学习（RL）进行后训练。

核心挑战：
在 dLLMs 上应用标准 RL（如策略梯度）面临一个根本性难题：不可解的对数似然（Intractable Log-Likelihood）。

标准策略梯度公式 $\nabla_\theta J(\theta) = E[R \nabla_\theta \log \pi_\theta(x|c)]$ 依赖于计算生成序列的真实对数似然 $\log \pi_\theta(x|c)$ 。
对于扩散模型，该值无法直接计算。
现有方案的缺陷： prior work（如 D1, UniGRPO）通常使用证据下界（ELBO）或单步估计作为真实似然的代理。
- ELBO 的局限性： ELBO 仅是真实似然的下界（ $ELBO \le \log \pi_\theta$ ）。
- 梯度偏差： 当奖励 $R$ 为正时，最大化 ELBO 是合理的；但当奖励 $R$ 为负（即需要惩罚低质量输出）时，最小化 ELBO 并不能保证最小化真实的对数似然。这导致策略梯度存在显著偏差，模型无法有效从负反馈中学习，限制了 RL 在 dLLMs 上的性能。

2. 方法论 (Methodology)

为了解决上述偏差问题，作者提出了 SPG (Sandwiched Policy Gradient，夹心策略梯度) 算法。其核心思想是利用似然的上下界来“夹住”不可解的真实对数似然，从而构建一个无偏或低偏差的优化目标。

2.1 夹心策略梯度 (Sandwiched Policy Gradient)

SPG 根据样本的相对优势（Advantage, $A_j$ ）动态选择优化目标：

正优势样本 ( $A_j \ge 0$ )： 最大化 证据下界 (ELBO, $L_{ELBO}$ )。这鼓励模型生成高质量序列。
负优势样本 ( $A_j < 0$ )： 最小化 证据上界 (EUBO, $L_{EUBO}$ )。
- 这是本文的关键创新。通过最小化上界，可以确保真实对数似然 $\log \pi_\theta$ 被有效降低（因为 $\log \pi_\theta \le L_{EUBO}$ ，最小化上界有助于降低下界）。
- 这解决了传统方法中“最小化下界无法有效惩罚负样本”的问题。

优化目标函数定义为：
$J_{SPG}(\theta) = E \left[ \frac{1}{g} \sum_{j=1}^g \left( \mathbb{1}_{A_j \ge 0} \cdot A_j L_{ELBO}(x_j) + \mathbb{1}_{A_j < 0} \cdot A_j L_{EUBO}(x_j) \right) \right]$

2.2 证据上界 (Evidence Upper Bound, EUBO)

作者基于 Rényi 变分界推导了适用于掩码扩散模型的 EUBO。

公式形式涉及对期望外部取对数（ $\log E[\dots]$ ），这通常会导致蒙特卡洛估计的偏差。
为了平衡紧度与计算可行性，作者引入了超参数 $\beta$ 来控制界限的紧度。

2.3 混合估计与方差降低 (Mixture Estimation)

直接最小化 EUBO 在蒙特卡洛估计中可能不稳定且计算成本高。为此，作者提出了混合策略：

对于负优势样本，使用 $L_{Mix} = \omega \cdot L_{EUBO} + (1-\omega) \cdot L_{ELBO}$ 。
理论优势： 证明了最优的混合系数 $\omega^*$ 可以严格降低梯度的方差，且优于单独使用上界或下界。
置信度感知加权： 混合梯度根据模型预测的置信度自动调整权重，对不确定的 token 给予较小权重，对高置信度 token 给予较大权重，防止梯度消失。

2.4 块状掩码策略 (Block-wise Masking Strategy)

为了在策略 rollout 和优化过程中更好地对齐数据分布，作者提出了一种新的蒙特卡洛采样策略：

不再使用完全随机的掩码，而是将序列分为若干块（Block）。
随机选择一个块进行掩码，保持前面的块干净（Clean），后面的块完全掩码。
这种策略更符合 dLLMs 在实际生成时（通常是半自回归、分块去噪）的分布，提高了训练的稳定性和效率。

3. 主要贡献 (Key Contributions)

提出 SPG 算法： 首个针对扩散语言模型设计的、基于“夹心”上下界的策略梯度算法，有效解决了因对数似然不可解导致的 RL 梯度偏差问题。
理论推导与混合优化： 推导了适用于 MDLM 的证据上界（EUBO），并提出了上下界混合估计方法，从理论上证明了其能降低梯度方差并提高训练稳定性。
块状掩码技术： 提出了一种与 dLLMs 生成机制更匹配的块状掩码策略，显著提升了蒙特卡洛估计的稳定性。
SOTA 性能： 在多个基准测试中取得了当前扩散模型 RL 训练的最优结果。

4. 实验结果 (Results)

作者在四个数学/逻辑推理基准（GSM8K, MATH500, Countdown, Sudoku）和两个代码基准（HumanEval, MBPP）上进行了广泛实验。

主要发现：

推理任务表现卓越：
- GSM8K: 相比 SOTA (UniGRPO)，准确率提升 3.6% (86.1% vs 82.5%)。
- MATH500: 提升 2.6% (40.0% vs 37.4%)。
- Countdown: 提升 18.4% (70.7% vs 52.3%)。
- Sudoku: 提升 27.0% (94.0% vs 67.0%)。
代码任务： 在 HumanEval 和 MBPP 上也分别提升了 1.9% 和 4.7%。
消融实验验证：
- 移除负优势惩罚（即不使用 EUBO）会导致性能大幅下降，证明了负反馈学习的重要性。
- 混合策略（Mixture）优于单独使用 ELBO 或 EUBO，且梯度更稳定。
- 块状掩码策略显著优于随机掩码（特别是在 Countdown 任务上提升巨大）。
泛化性： SPG 在不同推理策略（如不同的去噪块大小、随机去噪 vs 置信度去噪）下均表现出鲁棒性，优于基线模型。

5. 意义与影响 (Significance)

理论突破： 填补了扩散语言模型在强化学习对齐方面的理论空白。之前的方法由于依赖 ELBO 近似，在处理负奖励时存在本质缺陷。SPG 通过引入上界最小化，提供了一个 principled（有原则的）解决方案。
性能提升： 证明了通过改进优化目标，dLLMs 在复杂推理任务（如数学解题、数独）上的能力可以大幅超越现有的自回归模型或早期的扩散 RL 方法。
通用性： 提出的“夹心”思想和块状掩码策略不仅适用于当前的 MDLM，也为未来设计更高效的扩散模型 RL 算法提供了新的范式。
效率与稳定： 尽管引入了蒙特卡洛估计，但通过混合策略和块状掩码，SPG 在保持计算成本与现有方法相当的同时，显著提高了训练的稳定性和收敛速度。

总结：
SPG 通过巧妙地利用似然的上界和下界，解决了扩散语言模型在强化学习中“负反馈学习难”的核心痛点。其提出的算法不仅在数学和逻辑推理任务上刷新了记录，也为扩散模型在需要高可靠性推理的领域（如代码生成、科学计算）的落地应用奠定了坚实基础。