Diffusion-State Policy Optimization for Masked Diffusion Language Models¶

会议: NeurIPS 2026
arXiv: 2602.06462
代码: https://github.com/dllm-reasoning/d1
领域: LLM 其他
关键词: 掩码扩散语言模型, 扩散状态策略优化, 细粒度信用分配, 强化学习, 策略梯度

一句话总结¶

针对掩码扩散语言模型仅依赖最终生成文本标量奖励导致中间去噪决策信用分配粗糙的问题，本文提出 DiSPO，通过在中间掩码状态复用 rollout 缓存 logits 进行同状态分支重采样并仅更新新填充 token，在不增加额外多步扩散 rollout 和优化步数的前提下显著提升了数学推理与规划能力。

研究背景与动机¶

掩码扩散语言模型（MDLM）通过在多个去噪步中迭代填补被遮蔽的 token 位置来生成文本，这种非自回归且任意顺序的解码机制赋予了模型极高的生成灵活性与任意时刻补全接口。然而，当将强化学习策略优化（如基于 diffu-GRPO 或 SPG 的强化学习对齐）应用于 MDLM 时，现存方法普遍沿袭自回归大模型的终止反馈设定——仅根据最终生成的完整序列计算一个标量奖励，并直接将这一序列级结果粗粒度地归因于整个去噪轨迹中的所有局部决策。

这种将整条轨迹坍缩为单次老虎机评估的做法在 MDLM 中带来了严重的信用分配问题。在实际 rollout 过程中，MDLM 在每个中间去噪步本就已经对当前所有掩码位置输出了细粒度的 token 分布与 logits，但基于最终结果的标量奖励根本无法辨别中间某一步的具体掩码填充到底是关键转折还是错误尝试。部分先验工作尝试通过对中间状态做奖励塑形（如 SAPO）或在步间比较不同完整轨迹（如 MDPO）来缓解反馈稀疏性，但这类方案要么需要从中间步启动昂贵的多步扩散继续采样，极大增加了 rollout 计算开销，要么未直接以单步状态条件下的填充动作为优化单元。

本文的切入角度是充分利用 MDLM 独特的即时补全特性：在去噪过程中的任意中间掩码序列天然构成了一个马尔可夫决策状态，而当前掩码位置的填补就是该状态下的动作。由于模型在采样时已计算并缓存了全部掩码 logits，我们完全可以在不增加任何额外多步扩散前向传播的前提下，在相同掩码状态下直接重采样备选动作。核心 idea：将中间掩码状态显式视为决策状态，通过复用采样缓存的 logits 生成同状态反事实分支补全，用基础终止奖励评估分支优劣并仅更新新填补的动作 token，构建零额外 rollout 开销的扩散状态策略优化插件（DiSPO）。

方法详解¶

整体框架¶

DiSPO 是一个即插即用的策略优化层，与任何基于终止反馈的基础优化器（如 diffu-GRPO、SPG）完全正交并解耦。其整体流水线分为三个串行协同阶段：首先在多步扩散采样过程中缓存中间掩码状态与对应 logits 并获得轨迹终止奖励；随后在选定的中间去噪步上，直接复用缓存 logits 进行局部同状态重采样生成若干反事实补全分支，利用原始终止奖励评估器打分并计算状态内相对优势；最后通过局部替代似然仅对新增填补的 token 位置计算策略梯度并反向传播，实现对中间去噪决策的精确信用分配。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["提示词输入与多步去噪采样<br/>缓存中间掩码状态与位置 Logits"] --> B["扩散状态与动作解耦<br/>固定掩码状态并定义动作空间"]
    B --> C["同状态缓存分支采样<br/>零额外前向生成反事实补全"]
    C --> D["状态内群组优势评估<br/>基础终止奖励打分与组内均值基线"]
    D --> E["局部替代似然梯度更新<br/>仅更新新增掩码动作 Token 权重"]
    E --> F["联合目标端到端优化<br/>混合基础终止损失与单步策略损失"]

关键设计¶

1. 扩散状态与动作解耦：将掩码序列转化为单步决策 Markov 状态

针对终止反馈无法精准归因局部去噪选择的痛点，DiSPO 重新形式化了 MDLM 的决策过程。对于提示词 \(q\) 和行为模型参数 \(\theta_{\text{old}}\)，采样轨迹生成一系列部分遮蔽序列 \(x_{k,1}, x_{k,2}, \dots, x_{k,T}\)。将去噪步 \(t\) 的状态定义为 \(s_{k,t} = (q, x_{k,t})\)，未填充的掩码位置集合记为 \(M_{k,t} = \{i: x_{k,t,i} = \texttt{[MASK]}\}\)，已知上下文为 \(U_{k,t} = [L] \setminus M_{k,t}\)。此时的动作 \(a_{k,t}\) 被明确定义为在当前所有掩码位置上的联合填充 \(a_{k,t} = \{a_{k,t,i}\}_{i \in M_{k,t}}\)。根据条件独立性，动作策略分解为：

\[\pi_{\theta}(a_{k,t} \mid s_{k,t}) = \prod_{i \in M_{k,t}} \pi_{\theta}(a_{k,t,i} \mid s_{k,t}, i)\]

填充操作 \(o = \textsc{Fill}(x_{k,t}, a_{k,t})\) 确定性地将掩码位置替换为动作 token，而保持 \(U_{k,t}\) 上原有 token 完全不变。这一解耦使中间填充决策直接对应到一条完整文本产物，从而允许直接使用最终任务的黑盒奖励函数 \(R(q, o)\) 进行评分，无需人为构造启发式的中间奖励。

2. 同状态缓存分支采样与群组优势：零开销捕获局部动作偏好

为了在不引入额外昂贵的多步扩散生成的前提下建立对比信号，DiSPO 巧妙复用了主 rollout 过程中产生的计算中间件。在采样阶段，模型在每个状态 \(s_{k,t}\) 输出的掩码位置 logits 被完整缓存。在选定的状态子集 \((k,t) \in \mathcal{S}\) 上，算法直接从这些缓存 logits 中并发重采样 \(Z\) 个独立候选动作 \(a_{k,t,1}, \dots, a_{k,t,Z} \sim \pi_{\theta_{\text{old}}}(\cdot \mid s_{k,t})\)，并构造补全 \(o_{k,t,z} = \textsc{Fill}(x_{k,t}, a_{k,t,z})\)。由于所有分支均从同一个确定性状态 \(s_{k,t}\) 出发，外部环境与历史上下文严格对齐。系统利用原始终止奖励对各分支打分 \(R_{k,t,z} = R(q, o_{k,t,z})\)，并计算同状态组内平均基线及相对优势：

\[\bar{R}_{k,t} = \frac{1}{Z}\sum_{z=1}^{Z} R_{k,t,z}, \quad A_{k,t,z} = R_{k,t,z} - \bar{R}_{k,t}\]

该优势消除了特定提示词和上下文状态的固有难度偏差，精确刻画了“在当前掩码状态下，备选填充方案优于同批次平均水平的程度”。

3. 局部单步替代似然：仅更新动作 Token 的方差缩减机制

在对中间动作进行策略梯度更新时，若将梯度回传至整个序列的所有位置（包括已固定的已知 token），引入的无关随机噪声会导致极大的梯度方差。DiSPO 基于一步掩码提示重构技术构建了步级替代对数似然 \(\tilde{\pi}_{\theta}\)，并严格限制其求和范围仅覆盖动作位置 \(M_{k,t}\)：

\[\log \tilde{\pi}_{\theta}(a_{k,t,z} \mid s_{k,t}) := \mathbb{E}_{m}\left[\sum_{i \in M_{k,t}} \log \tilde{\pi}_{\theta}(a_{k,t,z,i} \mid s_{k,t}^{(m)}, i)\right]\]

单步策略损失形式化为重要性采样加权的优势乘积：

\[\mathcal{L}_{\text{step}}^{(k,t)}(\theta) = -\frac{1}{Z}\sum_{z=1}^{Z} \exp\left(\log \tilde{\pi}_{\theta}(a_{k,t,z} \mid s_{k,t}) - \log \tilde{\pi}_{\theta_{\text{old}}}(a_{k,t,z} \mid s_{k,t})\right) A_{k,t,z}\]

理论证明，这种 action-only 的局部更新将梯度方差按掩码占比 \(m/L\) 成比例缩减（命题 A.3），同时 \(Z\) 个分支的平均带来标准的 \(O(1/Z)\) 方差下降（命题 A.4），在理论与实现上兼顾了梯度的无偏性与极高数值稳定性。

4. 联合混合目标与后段偏置时间步采样：自适应调度关键决策节点

在扩散去噪初期，序列大部分仍为掩码，单步填充容易产生极度不完整的无序乱码，导致奖励区分度低；而在去噪后期，上下文愈发清晰，中间决策对最终正确性的影响愈发决定性。实测表明，去噪前段（进度前 30%）仅有 4.5% 的状态在分支间产生奖励差异，而中后段则高达 18%–20% 以上。因此，DiSPO 采用后段偏置的多项式时间步采样分布 \(\omega(t) \propto (t/T)^4\)，将有限的更新预算集中在奖励区分度显著的高价值状态上。最终总损失将单步损失与任意基础优化器的终止损失 \(\mathcal{L}_{\text{base}}\) 线性相加：

\[\mathcal{L}(\theta) = \alpha_{\text{base}} \mathcal{L}_{\text{base}}(\theta) + \alpha_{\text{step}} \sum_{(k,t) \in \mathcal{S}} \mathcal{L}_{\text{step}}^{(k,t)}(\theta)\]

损失函数 / 训练策略¶

训练采用 LoRA 微调（rank=128, \(\alpha=64\)），配合 AdamW 优化器，学习率设为 \(3\times 10^{-6}\) 或 \(5\times 10^{-6}\)，梯度裁剪为 0.2。对于扩散时间步，每个轨迹仅抽取 1 个高收益步进行分支（\(|T_{\text{sub}}|=1\)），分支数默认设置为极轻量的 \(Z=2\)。基础权重设为 \(\alpha_{\text{base}}=1.0\)，单步权重 \(\alpha_{\text{step}}\) 在 0.1 到 0.9 间调优（默认推荐 0.5）。多步扩散采样步数与 RL 更新步数与 baseline 保持严格一致，主训耗时资源（rollout 与 optimizer 步）完全对齐。

实验关键数据¶

主实验¶

在 LLaDA-8B-Instruct 基础模型上，评估覆盖符号规划（Sudoku、Countdown）与数学推理（GSM8K、MATH500），在严格对齐扩散 rollout 计算与优化器更新步数的条件下，主实验对比结果如下（准确率 Exact-Match %，不同输出长度 \(N_{\text{gen}} \in \{128, 256, 512\}\) 的平均表现）：

基础优化器	方法	Sudoku (Avg)	Countdown (Avg)	GSM8K (Avg)	MATH500 (Avg)
无强化学习	LLaDA-8B-Instruct (Zero-shot)	8.0	18.7	74.5	31.5
diffu-GRPO 基线	diffu-GRPO [Zhao et al., 2025]	14.1	33.9	78.1	36.5
diffu-GRPO 增强	DiSPO (diffu-GRPO)	31.4 (+17.3)	53.9 (+20.0)	79.9 (+1.8)	37.8 (+1.3)
SPG 基线	SPG [Wang et al., 2026]	26.2	67.2	79.9	37.4
SPG 增强	DiSPO (SPG)	43.5 (+17.3)	68.5 (+1.3)	80.8 (+0.9)	38.5 (+1.1)
参考组 (非对齐高算力)	SAPO (带多步重展 Rollout 塑形)	19.6	53.3	79.2	36.8

在监督微调版本 LLaDA-SFT 上，即使采用极保守的单步权重 \(\alpha_{\text{step}}=0.1\)，DiSPO 依然全面超越 SFT & diffu-GRPO（Sudoku 由 16.1% 提至 21.0%，Countdown 由 36.3% 提至 36.8%，GSM8K 由 78.8% 提至 80.2%，MATH500 由 37.5% 提至 38.1%）。

消融实验¶

在 \(4\times 4\) 数独规划任务上对核心设计进行系统的消融实验，验证单步估计器设计与优化配方的必要性（准确率 %）：

消融维度	具体配置	\(N_{\text{gen}}=128\)	\(N_{\text{gen}}=256\)	\(N_{\text{gen}}=512\)	平均准确率 (Avg)	机制说明
梯度范围	全序列更新 (All tokens)	22.7	21.4	20.6	21.6	包含无关位置引入方差，效果显著下降
梯度范围	仅动作位置 (Action only)	26.6	23.3	20.0	23.3	符合理论预期，方差严格控制
分支数 \(Z\)	\(Z=2\) (默认设置)	26.6	23.3	20.0	23.3	极低显存与计算开销下即可提供稳定增益
分支数 \(Z\)	\(Z=3\)	25.1	22.8	21.0	23.0	边际收益趋缓，无单调提升
分支数 \(Z\)	\(Z=4\)	26.6	22.1	19.6	22.8	验证了仅需小分支即可满足需求
时间步采样	均匀采样 (Uniform)	25.1	23.0	17.5	21.9	早期无序状态稀释了有效监督信号
时间步采样	早期偏置 (Early-focused)	24.2	18.1	16.0	19.4	大量采样缺乏区分度的初期掩码，表现最差
时间步采样	后段偏置 (Late-focused, \(k=4\))	26.6	23.3	20.0	23.3	聚焦高信噪比决策关键步，增益最稳健
损失组合	仅基础损失 \(\mathcal{L}_{\text{base}}\)	18.4	12.9	11.0	14.1	缺乏中间单步信用分配信号
损失组合	仅单步损失 \(\mathcal{L}_{\text{step}}\)	24.9	18.4	17.7	20.3	局部对齐效果优异，但脱离全局序列目标约束
损失组合	双目标联合 (\(\mathcal{L}_{\text{base}} + \mathcal{L}_{\text{step}}\))	26.6	23.3	20.0	23.3	互补驱动，达到最佳整体规划效果

关键发现¶

有效抑制过早承诺错误（Premature Commitment）：数独生成中的早期约束违背分析表明，diffu-GRPO 倾向于在去噪初期就填入错误数字从而导致后续求解完全崩溃，而 DiSPO 显著将首次违背约束的时间步推后，证明单步策略优化有效指导了模型在中间阶段谨慎生成符合全局约束的合理局部决策。
规划任务收益远高于纯数学问答：在逻辑严密、多步依赖极强的数独和 Countdown 任务上，DiSPO 带来了超过 17–20 个百分点的飞跃；而在 GSM8K 等自然语言容错度较高的任务中，增益为 1–2 个百分点，表明细粒度状态信用分配对于多阶段强约束搜索尤为关键。
算力开销优势极其显著：相比 SAPO 等依赖多步二次扩散前向来做奖励塑形的方法，DiSPO 额外扩散前向步数严格为 0，即便在未高度优化的原型代码吞吐量受限下，其实际壁钟训练曲线仍快速追平并超越基线。

亮点与洞察¶

复用 Rollout 缓存实现零扩散开销分支：充分利用了 MDLM 在前向去噪时已经计算了全部位置 logits 的天然特性，仅通过缓存读取即刻生成多路 counterfactual 反事实动作，避开了传统强化学习中探索新动作必须重新推演多步轨迹的高昂代价。
严格证明了混合目标的理论自洽性：证明了组内相对基线下单步损失函数的期望梯度与单步期望回报目标完全对齐（差一个确定性系数 \((Z-1)/Z\)），并且与序列级策略梯度在线性期望下统一为一个规范的联合优化问题，为扩散模型的分步 RL 提供了严谨的数学基础。
通用即插即用生态兼容性：DiSPO 仅定义了单步状态维度的优势形式，底层不管是结合 diffu-GRPO 还是更先进的 SPG 或未来的 GDPO，均能无缝嵌入，具备极强的基座兼容性。

局限与展望¶

工程实现层面的壁钟时间开销：尽管算法层面的扩散 rollout 与参数更新次数为 0 增量，但当前开源实现的单步替代似然计算与奖励打分较为碎片化，多次未合并的 kernel 启动导致实际训练速度略有放缓（吞吐量约为基线的 0.4 倍），后续需通过算子融合、分支批处理（batching）和状态表征复用予以优化。
单步反事实补全的独立性假定：直接使用一次性全部填充 \(\textsc{Fill}(x_t, a_t)\) 并以此打分评估中间状态动作，在去噪非常早期的掩码密集状态下，可能因为上下文过少而导致补全分布略显生硬，未来可探索结合轻量级混合去噪步数。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙利用 MDLM 即时补全与 logits 缓存机制定义单步决策 MDP，思路清晰且极具扩散模型领域特色。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖多基线（diffu-GRPO, SPG, SAPO）、多任务（数学与复杂符号规划）、方差理论验证与早期违背深度诊断。
写作质量: ⭐⭐⭐⭐⭐ 动机阐述自然连贯，数学证明完备，计算开销拆解诚实且透明。
价值: ⭐⭐⭐⭐⭐ 为非自回归扩散语言模型的高效细粒度后训练提供了极具通用性的基础设计范式。