改进掩码扩散语言模型的中间决策优化，提升生成质量。
Diffusion-State Policy Optimization for Masked Diffusion Language Models
- 在中间掩码状态分支重采样，直接优化填词决策。
- 相同计算量下，数学与规划任务准确率显著提升。
- 无需额外扩散推演，可作为通用插件集成到现有模型。
掩码扩散语言模型通过迭代填词生成文本，但仅在最终结果上给予奖励，导致中间填词决策的信用分配粗糙。我们提出扩散状态策略优化（DiSPO），一种可插入的信用分配模块，直接优化中间填词决策。在选定的中间掩码状态，DiSPO通过从回滚缓存的逻辑值中重采样当前掩码位置，评估生成结果，并仅更新新填入的词元，无需额外多步扩散推演或优化器步骤。我们为分支生成定义了固定状态目标，并推导出复用终端反馈策略优化相同回滚路径的策略梯度估计器。在 LLaDA-8B-Instruct 上的实验表明，DiSPO 在匹配推演计算量和优化步数条件下，持续优于 diffu-GRPO 与 SPG 等基线模型，在数学与规划基准测试中表现更优，支持其作为掩码扩散策略优化的通用插件。项目页面见 https://daioba.github.io/dispo。
原文摘要 · Abstract (English)
Masked diffusion language models generate text through iterative masked-token filling, but terminal-only rewards on final completions provide coarse credit assignment for the intermediate filling decisions that shape the generation process. We propose Diffusion-State Policy Optimization (DiSPO), a plug-in credit-assignment layer that directly optimizes intermediate filling decisions. At selected intermediate masked states, DiSPO branches by resampling the currently masked positions from rollout-cached logits, scores the resulting completions, and updates only the newly filled tokens, requiring no additional multi-step diffusion rollouts or optimizer steps. We formalize a fixed-state objective for branched completions and derive a policy-gradient estimator that reuses the same rollouts as terminal-feedback policy optimization. Experiments on LLaDA-8B-Instruct show that DiSPO consistently improves terminal-feedback baselines, including diffu-GRPO and SPG, on math and planning benchmarks under matched rollout compute and optimizer steps, supporting its use as a general plug-in for masked diffusion policy optimization. Our project page is available at https://daioba.github.io/dispo .
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。