Title: wd1:  Weighted Policy Optimization for Reasoning in Diffusion Language Models
Primary area: foundation or frontier models, including LLMs
Keywords: Diffusion Language Models, Reinforcement Learning, Reasoning

Abstract: 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 objective as a weighted log-likelihood, requiring only a single approximation for the current parametrized policy likelihood. Formally, we show that our proposed method is equivalent to training an energy-guided discrete diffusion model, thereby confirming its theoretical soundness. In experiments, wd1 outperforms diffusion-based GRPO (d1) while requiring lower computational cost, achieving up to a $+43\\%$ improvement in accuracy. Furthermore, we extend wd1 to denoising-stepwise weighted policy optimization (wd1++), which surpasses concurrent RL for dLLMs methods, attaining improvement $+6.2\\%$ on MATH500 ($44.2\\%$) and $+2.5\\%$ on GSM8K ($84.5\\%$) over d1 with only 20 training steps.

TL;DR: We propose a novel policy optimization method for dLLMs reasoning, reducing the error caused by log-likelihood approximation error.