💬 NLP

Diffusion-State Policy Optimization for Masked Diffusion Language Models

本文提出了扩散状态策略优化（DiSPO），这是一种即插即用方法，通过分支与评分机制直接优化中间令牌填充决策，从而在无需额外 rollout 或优化器步数的情况下提升掩码扩散语言模型的终端性能。

原作者： Daisuke Oba, Hiroki Furuta, Naoaki Okazaki

发布于 2026-05-20

📖 1 分钟阅读☕ 轻松阅读

原作者： Daisuke Oba, Hiroki Furuta, Naoaki Okazaki

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

以下是用通俗易懂的语言和生动的类比，对论文《Diffusion-State Policy Optimization for Masked Diffusion Language Models》（掩蔽扩散语言模型的扩散状态策略优化）进行的解读。

宏观视角：修复 AI 训练中的“盲区”

想象你正在教一名学生解一道复杂的谜题，比如数独。

旧方法（终端反馈）： 你让学生把整个谜题填完。等到最后，你再看结果。如果错了，你说“做得不好”，学生只能猜测自己究竟在哪一步填错了数字；如果对了，你说“做得好”。问题在于，学生根本不知道具体哪一步是失误，哪一步是神来之笔。他们得到的只是对整个过程的模糊“好”或“坏”的评价。
新方法（DISPO）： 这篇论文提出了一种名为DISPO（Diffusion-State Policy Optimization，扩散状态策略优化）的方法。DISPO 不会等到最后才评价，而是在学生填空中断的各个环节暂停。它会问：“如果你现在为这个特定的空位选择另一个数字，最终结果会更好吗？”它会瞬间测试几种替代方案，并针对该特定决策给出反馈。

核心问题：“粗糙的信用分配”

论文指出，当前的 AI 模型（特别是掩蔽扩散语言模型）就像“旧方法”中的学生。

这些模型通过反复填充被遮蔽（隐藏）的词语来生成文本。
目前，它们仅根据生成的最终句子获得奖励（分数）。
这被称为“粗糙的信用分配”。这就像只根据最终答案来批改数学试卷。如果答案错了，老师不知道学生是在第 1 步、第 5 步还是第 10 步犯了错。模型必须猜测是哪个中间步骤导致了错误。

解决方案：DISPO（“如果……会怎样”模拟器）

DISPO 作为一个“即插即用”的层，通过关注中间步骤来解决这个问题。以下是其工作原理，我们使用厨师类比来说明：

状态（锅）： 想象 AI 是一位正在炖汤的厨师。在某个时刻，锅里已经有一半的食材，但还有一些缺失（被遮蔽）。这就是“状态”。
动作（添加食材）： 厨师需要决定下一步添加什么。
分支（“如果……会怎样”测试）： 与其只添加一种食材并寄希望于好运，DISPO 会说：“让我们暂停时间。”
- 它保留当前的锅（状态）。
- 它快速模拟添加食材 A，看看汤会变成什么样。
- 它模拟添加食材 B，看看那锅汤会变成什么样。
- 它无需从头重新炖整锅汤。它利用“缓存日志”（一种思维捷径）来瞬间预测这些不同选择的结果。
奖励： 它比较结果。如果“食材 A"能做出味道更好的汤（更高的奖励），模型就会学会在这个特定时刻偏好 A 而不是 B。
更新： 模型仅针对那个特定的食材选择更新其“大脑”。它不需要重新学习整道食谱，只需修正那一个决策。

为何此法独特

论文强调了该方法的三个主要优势：

无需额外烹饪时间： 通常，为了测试不同的选择，AI 必须一遍又一遍地运行整个生成过程（这很慢）。DISPO 很聪明，因为它利用了模型在正常训练运行中已经生成的数据。它不需要额外的“ rollout"（额外的烹饪过程）。它只是重用了模型已经计算出的"logits"（模型内部的置信度分数）。
精准性： 它解决了“推诿责任”的问题。它不再因为一个坏词而惩罚整句话，而是精准定位哪个词的选择不够理想并进行修正。
即插即用： 你可以将这种方法附加到现有的训练方法（如diffu-GRPO或SPG）上，在不改变其核心结构的情况下使其更智能。

结果：在数学和逻辑方面表现更佳

研究人员在名为LLaDA-8B-Instruct的模型上测试了该方法。他们给模型布置了高难度任务，包括：

数学： 解决小学应用题（GSM8K）和竞赛级数学题（MATH500）。
规划： 解决数独谜题和“倒计时”数字游戏。

结果：
当将 DISPO 添加到标准训练方法中时，模型在这些任务上的表现显著提升。

它减少了“过早承诺”（过早填入数字导致陷入死胡同）。
它正确解决的谜题更多。
关键在于，它是在不增加主训练循环计算资源的情况下实现了这些提升。改进完全来自于更仔细地审视中间步骤。

总结比喻

可以把旧方法想象成一位高尔夫教练，他只在打完 18 洞的整轮比赛后才告诉你“好球”或“坏球”。你完全不知道第 3 洞的挥杆是否是问题所在。

DISPO则是一位站在你身边的教练，在每一个洞都陪着你。当你准备击球时，教练会说：“如果你向左瞄准 5 度，很可能会打上果岭；如果向右瞄准，就会打进沙坑。让我们试试向左。”它在球落地之前就针对决策提供反馈，帮助你为每一个特定时刻学习正确的策略，而不仅仅是关注最终得分。

技术摘要：面向掩码扩散语言模型的扩散状态策略优化

问题陈述

掩码扩散语言模型（MDLMs）通过在多个去噪步骤中迭代填充掩码位置来生成文本。虽然策略优化（PO）方法已成功提升了语言模型的推理与对齐能力，但大多数现有方法依赖于仅终端奖励。在这些方法中，一个标量奖励被分配给最终完成结果，随后该奖励被反向传播至所有中间生成决策。

这种方法导致了粗糙的信用分配。它将整个局部填充决策序列视为单一单元，未能利用在 rollout 过程中已产生的丰富中间状态 - 动作信息。具体而言，在任何中间去噪步骤，模型都会为所有当前掩码位置生成 token 级分布，然而标准的终端反馈 PO 忽略了这种粒度，无论各决策对最终结果的具体贡献如何，均赋予所有决策相同的结局信号。

方法论：扩散状态策略优化（DISPO）

作者提出了扩散状态策略优化（DISPO），这是一种即插即用的信用分配层，旨在直接优化中间填充决策，而无需额外的多步扩散 rollout。

核心概念

DISPO 将 MDLMs 的生成过程重构为状态 - 动作框架：

状态（ $s_{k,t}$ ）： 去噪步骤 $t$ 处的中间部分掩码序列。
动作（ $a_{k,t}$ ）： 在该状态下对当前掩码位置的联合填充。

算法机制

Rollout 与缓存： 在使用冻结的行为参数（ $\theta_{old}$ ）进行标准 rollout 期间，模型生成部分掩码序列的轨迹。关键在于，缓存了每个访问状态中所有掩码位置的 logits。
状态条件分支： 在选定的中间状态处，DISPO 通过利用 rollout 中缓存的 logits 对掩码位置重新采样替代填充，从而“分支”。这从同一中间状态生成了多个“反事实”完成结果（ $Z$ 个分支）。
评分与优势： 每个分支完成结果使用与标准 PO 相同的终端奖励函数（ $R$ ）进行评分。基于该特定状态下 $Z$ 个完成结果的奖励，计算每个分支的组相对优势。
Token 局部更新： 策略梯度被计算以仅更新当前状态下的新填充 token（即动作），将先前填充的 token 视为固定状态的一部分。
代理似然： 为了解决 MDLMs 中精确序列似然不可计算的问题，DISPO 利用可计算的单步掩码 token 代理（ $\tilde{\pi}_\theta$ ）。这使得能够专门针对被更新的掩码位置计算似然比。

训练目标

DISPO 定义了一个固定状态期望回报目标。它可以通过加权和与基础终端反馈优化器（例如 diffu-GRPO 或 SPG）结合：
$\mathcal{L}(\theta) = \alpha_{base} \mathcal{L}_{base}(\theta) + \alpha_{step} \mathcal{L}_{step}(\theta)$
其中 $\mathcal{L}_{step}$ 是从分支完成结果中导出的状态级损失。理论分析证实，该组合目标为包含终端目标和时间步平均固定状态目标的混合目标提供了有效的策略梯度信号。

主要贡献

状态条件信用分配： 本文提出了一种方法，将中间掩码序列视为状态，将掩码填充视为动作，实现了标准终端反馈 PO 所缺乏的细粒度信用分配。
基于缓存 logits 的高效分支： 一项关键创新是利用 rollout 缓存的 logits 从同一状态生成替代完成结果。这使得能够在无需额外多步扩散前向传播的情况下进行同状态比较，从而保持主导训练成本（rollout 和优化器步骤）与基线相匹配。
理论形式化： 作者将 DISPO 形式化为固定状态目标并推导了策略梯度估计量。他们证明了步级损失是固定状态目标的无偏（缩放）估计量，并且将其与终端损失结合可以优化定义明确的混合目标。
方差降低策略： 该方法采用了受方差分析启发的两项设计选择：
- Token 局部更新： 将梯度限制在可操作（掩码）token 上，与更新所有 token 相比降低了方差。
- 同状态平均： 对 $Z$ 个分支进行平均，将方差降低了 $O(1/Z)$ 倍。

实验结果

作者在LLaDA-8B-Instruct上评估了 DISPO，将其与两个基础优化器diffu-GRPO和SPG集成。

基准测试： 实验在数学推理（GSM8K, MATH500）和符号规划（Sudoku, Countdown）上进行。
性能： 在匹配计算预算（相同数量的 rollout、组大小和优化器步骤）下，DISPO 在所有基准测试中一致地改进了两个基础优化器。
- 在规划任务中提升尤为显著（例如，使用 diffu-GRPO 时，Sudoku 准确率从约 14% 提升至约 31%）。
- 改进在不同生成预算（ $N_{gen} \in \{128, 256, 512\}$ ）下持续存在。
消融研究：
- 状态级与终端级： 单独的状态级目标优于仅终端基线，而两者的结合产生了最佳结果。
- 时机： 偏向后期的时间步采样（关注更接近完成的状态）优于均匀采样或偏向早期的采样，这表明接近最终完成的中间状态提供了更可靠的信用信号。
- 梯度方差： 实验证实，仅更新可操作 token 并增加分支大小 $Z$ 显著降低了梯度方差。
误差分析： 对 Sudoku 上“过早承诺”的分析显示，与基线相比，DISPO 将约束违规的首次发生时间推移到了更晚的去噪步骤，表明优化中间动作有助于避免早期且难以纠正的错误。

意义与主张

本文主张 DISPO 是面向掩码扩散策略优化的通用即插即用模块。其主要意义在于：

利用中间信息： 它证明了 MDLM 生成过程中固有的中间状态 - 动作信息可用于细粒度信用分配，而无需生成全新轨迹的计算开销。
互补性： DISPO 被呈现为对现有方法（如奖励塑形（例如 SAPO）或轨迹级目标）的补充。虽然 SAPO 基于中间状态塑造奖励，但 DISPO 在固定状态下直接优化填充动作。
效率： 通过重用缓存的 logits，DISPO 在实现这些改进的同时并未增加主导训练成本（多步 rollout），使其成为现有 MDLM 训练流程中的实用补充。

作者指出，虽然他们当前的实现由于重复的代理和奖励评估而产生了一些挂钟时间开销，但算法成本与基线相匹配，且性能提升具有鲁棒性。他们将 DISPO 定位为迈向基于扩散的语言模型中更高效、更有效的推理的一步。