arXiv:2507.08838cs.LGcs.AI2025-07中稿 · ICLR被引 63

提出无需计算策略比的优化方法,显著提升扩散语言模型推理能力

wd1: Weighted Policy Optimization for Reasoning in Diffusion Language Models

  • 将强化学习目标重构为加权对数似然,仅需一次策略似然近似
  • 在LLaDA-8B上实现59%准确率提升,计算成本更低
  • 适用于需要高效推理优化的研究者与工程师

通过强化学习提升基于扩散的大语言模型(dLLMs)的推理能力仍是开放问题。由于dLLMs似然函数难以计算,现有方法需在每步优化中近似当前、旧及参考策略的似然,带来额外计算开销,并导致强化学习目标中策略比估计的高方差与误差。为此,我们提出wd1,一种新型无比率策略优化方法,将强化学习目标重写为加权对数似然,仅需对当前参数化策略似然进行一次近似。我们形式化证明该方法可被解释为能量引导的离散扩散训练结合负样本遗忘,验证其理论合理性。在LLaDA-8B模型上的实验表明,wd1优于基于扩散的GRPO(d1),且计算成本更低,准确率最高提升59%。进一步扩展至wd1++,仅用20步强化学习即在MATH500上达到44.2%、GSM8K上达到84.5%的顶尖数学推理性能。

原文摘要 · Abstract (English)

Improving the reasoning capabilities of diffusion-based large language models (dLLMs) through reinforcement learning (RL) remains an open problem. The intractability of dLLMs likelihood function necessitates approximating the current, old, and reference policy likelihoods at each policy optimization step. This reliance introduces additional computational overhead, and can lead to large variance and estimation error in RL objective -- particularly in computing the policy ratio for importance sampling. To mitigate these issues, we introduce wd1, a novel ratio-free policy optimization approach that reformulates the RL objective as a weighted log-likelihood, requiring only a single approximation for the current parametrized policy likelihood. We formally show that our proposed method can be interpreted as energy-guided discrete diffusion training combined with negative sample unlearning, thereby confirming its theoretical soundness. In experiments on LLaDA-8B model, wd1 outperforms diffusion-based GRPO (d1) while requiring lower computational cost, achieving up to a $+59\%$ improvement in accuracy. Furthermore, we extend wd1 to denoising-stepwise weighted policy optimization (wd1++), achieving state-of-the-art math performance of $44.2\%$ on MATH500 and $84.5\%$ on GSM8K with only 20 RL training steps.

扩散模型强化学习推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。