扩散状态策略优化用于掩码扩散语言模型
Diffusion-State Policy Optimization for Masked Diffusion Language Models
AI总结:
本文提出Diffusion-State Policy Optimization(DiSPO),一种用于掩码扩散语言模型的插件信用分配层,通过直接优化中间填充决策来改进生成过程,实验表明其在数学和规划基准测试中优于现有基线方法。
摘要:
掩码扩散语言模型通过迭代填充掩码标记来生成文本,但仅对最终完成结果的终端奖励对中间填充决策的信用分配过于粗糙。我们提出Diffusion-State Policy Optimization(DiSPO),一种插件信用分配层,直接优化中间填充决策。在选定的中间掩码状态下,DiSPO通过从滚出缓存的logits中重新采样当前掩码位置,评估由此产生的完成结果,并仅更新新填充的标记,无需额外的多步扩散滚出或优化器步骤。我们为分支完成形式化了一个固定状态目标,并推导出一个策略梯度估计器,该估计器重用与终端反馈策略优化相同的滚出。在LLaDA-8B-Instruct上的实验表明,DiSPO在匹配的滚出计算和优化器步骤下,一致提高了终端反馈基线,包括diffu-GRPO和SPG,在数学和规划基准测试中。我们的项目页面可在https://daioba.github.io/dispo上找到。
英文摘要:
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 .
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。