PaperlayerView exact revision on arXiv

Source-linked AI summary

wd1: Weighted Policy Optimization for Reasoning in Diffusion Language Models

Xiaohang Tang, Rares Dolga, Sangwoong Yoon, Ilija Bogunovic

arXiv:2507.08838v2cs.LGcs.AIstat.ML

TL;DR

Diffusion language models pose reinforcement-learning challenges because likelihood approximation can amplify error, increase variance, and add computational overhead. wd1 replaces policy-ratio estimation with a weighted log-likelihood objective, while wd1++ extends it to denoising-stepwise optimization; both improve reported reasoning performance with lower training cost.

  • Problem

    In dLLM reinforcement learning, intractable likelihoods require approximation, which can exponentially amplify errors, increase policy-ratio variance, and add computational overhead.

  • Method

    wd1 uses a weighted log-likelihood objective that avoids explicit policy ratios and requires only one likelihood approximation; its weighting increases higher-advantage completions and decreases lower-advantage ones.

  • Results

    wd1 outperforms d1 on Sudoku and Countdown, while wd1++ achieves state-of-the-art performance of 44.2% on MATH500 and 84.5% on GSM8K with only 20 training steps.

  • Takeaways & Limitations

    The results establish wd1 as a more scalable and efficient approach for fine-tuning dLLMs, with a theoretical interpretation as energy-guided diffusion training and low-advantage sample unlearning.

  • Takeaways & Limitations

    The weighted log-likelihood objective may assign vanishing weights to low-advantage completions and reinforce suboptimal samples when all sampled rewards are identical and low.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Fine-tuning dLLMs with RL is difficult because intractable likelihoods make policy-ratio estimation costly and error-prone. wd1 replaces explicit ratios with weighted log-likelihood optimization, while wd1++ extends this approach to denoising steps and achieves strong reasoning results.

  • Motivation: Intractable dLLM likelihoods require approximations that can increase variance and computational overhead, especially for longer completions and more diffusion steps.GRPO estimates likelihoods for current, old, and reference policies at every training step.
  • wd1: wd1 reformulates policy optimization as a weighted log-likelihood objective requiring only a single likelihood approximation.The method avoids explicit policy ratios and targets lower bias, variance, and computational cost.
  • wd1: The advantage-dependent weight increases higher-advantage completions while decreasing the probability of lower-advantage ones.Its weight is defined as (−w+ + w−), with w+ proportional to exp(A) and w− proportional to exp(−A).
  • Theoretical interpretation: wd1 is interpreted as energy-guided discrete diffusion training combined with unlearning of low-advantage data.The advantage function guides generation toward higher-advantage completions.
  • Results: 76.4% on Sudoku (+58.8% over d1) and 51.2% on Countdown (+16% over d1) were achieved by wd1 on LLaDA-8B-Instruct without supervised fine-tuning.The experiments also report significantly less computational burden during RL training.
  • Results: 44.2% on MATH500 and 84.5% on GSM8K were achieved by wd1++ with only 20 training steps and 10× fewer rollouts than baseline methods.wd1++ uses intermediate completions from the decoding process and surpasses several concurrent RL methods for dLLMs.

2 PRELIMINARIES

The paper introduces masked discrete diffusion language modeling and reviews policy optimization methods used for dLLMs. It emphasizes that approximated likelihood ratios can be biased, high-variance, and computationally expensive.

  • Diffusion language models: dLLM policies generate completions through iterative denoising rather than sequential autoregressive token generation.The generation policy is denoted πθ, with prompts q and completions o.
  • Diffusion language models: Masked diffusion progressively corrupts sequences with a mask token, with x0 denoting the fully denoised completion and xt the masked sequence at timestep t.The forward process is represented by pt|0(xt | x0).
  • Diffusion objectives: Denoising Cross Entropy approximates the marginal likelihood log πθ(x0) for supervised fine-tuning and reinforcement learning.Intermediate timesteps are sampled uniformly and masked sequences follow the predefined forward process.
  • Existing policy optimization: TRPO restricts policy updates with forward KL divergence, while PPO uses clipped policy ratios and reference-policy reverse-KL regularization.GRPO simplifies PPO by estimating advantages from normalized rewards within completion groups.
  • Policy optimization for dLLMs: Adapting GRPO to dLLMs is challenging because non-autoregressive denoising makes log πθ(o|q) intractable and requires approximation.Approximated likelihoods are used to compute policy ratios for importance sampling.
  • Policy optimization for dLLMs: Likelihood-ratio approximation errors can be exponentially amplified, while ELBO estimates may have high variance and d1 estimates may be biased.GRPO also applies the approximation separately to current, old, and reference policies, increasing computational overhead.

3 wd1: WEIGHTED POLICY OPTIMIZATION FOR DLLMS

wd1 reformulates dLLM policy optimization as a ratio-free weighted log-likelihood objective, reducing likelihood-estimation burdens while explicitly reinforcing beneficial completions and penalizing detrimental ones. wd1++ extends this approach to intermediate denoising completions.

  • wd1: wd1 eliminates likelihood-ratio estimation and requires only a single approximated current-policy likelihood, reducing computational burden, variance, and approximation error.The objective avoids both old-policy importance-sampling ratios and reference-policy regularization ratios.
  • wd1: Reverse-KL policy optimization yields a weighted log-likelihood objective whose weights are derived from group-relative advantages.The algorithm samples completions, computes group-relative advantages and normalized weights, then updates the policy using approximated completion log-likelihoods.
  • wd1: The weighted log-likelihood objective can still underuse low-advantage completions and reinforce uniformly low-reward samples.Exponential weighting may make negative samples receive negligible weights, while equal low rewards can cause all suboptimal completions to be reinforced.
  • wd1: wd1 explicitly reinforces positive samples and penalizes negative samples through complementary positive and negative weighting terms.The negative-sample penalty induces negative gradients, and optimization halts when all completions have identical advantages because w+ = w−.
  • wd1++: wd1++ incorporates intermediate clean completions from the denoising process, expanding the training group and using these completions to estimate advantages and weights.This addresses the unused intermediate completions produced during confidence-based remasking.

4 THEORETICAL INSIGHTS: ENERGY-GUIDED DIFFUSION SAMPLING

The paper interprets wd1 as training an energy-guided discrete diffusion model, with negative advantage defining the energy, and as unlearning low-advantage samples. This connects the weighted objective to established denoising score and cross-entropy training objectives.

  • Theoretical interpretation: wd1 is theoretically interpreted as energy-guided discrete diffusion training combined with negative sample unlearning.The energy guidance is based on the advantage function, while negative samples receive an unlearning interpretation.
  • Energy-guided sampling: Reverse-KL solution-policy sampling corresponds to energy-guided sampling with energy E(q, ·) = −Aπold(q, ·).The resulting target policy places greater sampling emphasis on completions with higher advantage.
  • Target diffusion process: The target policy uses the old policy’s forward diffusion process, with the forward-process assumption explicitly equating the two processes.The prompt is omitted in the diffusion notation used for this derivation.
  • Motivation: Training-free guided sampling would require estimating the posterior mean of exponential advantage, motivating direct training of the guided diffusion model instead.The proposed training objective avoids relying on that estimation procedure.
  • Weighted denoising objectives: Advantage-weighted denoising concrete score matching approximates the concrete score of the energy-guided discrete diffusion.The resulting AW-D-CSM objective can be converted into an advantage-weighted denoising cross-entropy objective for models such as LLaDA.
  • Theoretical equivalence: The weighted log-likelihood objective is equivalent to energy-guided diffusion training when DCE is used for likelihood approximation.The paper also interprets the additional negative-sample penalty in wd1 as data unlearning through ELBO minimization.

5 EXPERIMENTS

Experiments evaluate wd1 and wd1++ on reasoning benchmarks, showing improved accuracy alongside lower training costs and benefits from negative-sample weighting.

  • Experimental Setup: wd1 is evaluated on GSM8K, MATH, Sudoku, and Countdown using LLaDA-8B-Instruct with LoRA fine-tuning.The setup uses d1 as the main baseline and follows specified dataset splits and decoding procedures.
  • Superior Reasoning Ability: wd1 surpasses d1 by 43% in Sudoku test accuracy and achieves up to 25% improvement on Countdown with maximum length 256.Relative to base LLaDA, gains reach 54% on Sudoku and 42% on Countdown.
  • Reduced Training Cost: wd1 reduces training cost by removing SFT and requiring less RL cost per step than d1.The comparison includes SFT cost, average training time, FLOPs, and theoretical NFEs per training step.
  • Ablation Study: Removing negative-sample reinforcement significantly hurts performance, while balanced positive and negative weighting is most effective.Negative-only training also causes pronounced deterioration relative to the default combined weighting.

6 RELATED WORK

Related work covers reinforcement learning for diffusion and autoregressive language models, weighted regression, and ratio-free policy optimization.

  • RL for Diffusion-based LLM: Diffusion-language-model RL includes denoising-trajectory methods, direct predictive-distribution optimization, and concrete-score-based policy objectives.These approaches differ in whether they backpropagate through denoising steps or avoid likelihood approximation through score estimation.
  • RL for AR Models: Autoregressive RL precedents include GRPO, rejection sampling fine-tuning, and methods that penalize likelihood for negative-reward samples.These works provide related weighted-likelihood and positive/negative-sample perspectives.
  • RL via Weighted Regression: Weighted regression has been studied in advantage-weighted regression, diffusion policies for continuous control, and reward-based fine-tuning of autoregressive language models.The paper places wd1 within this broader family of likelihood-based approaches.
  • “Ratio-Free” Policy Optimization: Ratio-free policy optimization avoids importance sampling and reference-model regularization, a property valuable for diffusion models with approximate likelihoods.Vanilla policy gradient methods and RLOO are cited as inherently ratio-free examples.

7 CONCLUSION

The paper introduces wd1 as an efficient weighted policy optimization method for dLLM reasoning and reports improved accuracy with reduced RL-training cost.

  • Conclusion: wd1 minimizes reliance on likelihood approximation to mitigate bias from approximation errors in policy ratios.Its objective is derived as a weighted log-likelihood approximation to reverse-KL-constrained policy optimization.
  • Conclusion: wd1 surpasses d1 by up to 16% in accuracy on reasoning benchmarks without supervised fine-tuning.The paper also reports notable computational-efficiency improvements during RL training.
  • Conclusion: The results support wd1 as a more scalable and efficient approach for fine-tuning diffusion language models.This conclusion combines the reported accuracy and RL-training efficiency findings.

8 ETHICS AND REPRODUCIBILITY STATEMENT

The paper reports no ethics concerns and describes implementation, dataset, reward, hyperparameter, and theoretical-proof details for reproducibility.

  • Ethics: The authors report no question or concern regarding the Code of Ethics.This statement is part of the ethics and reproducibility declaration.
  • Reproducibility: Implementation details are provided in Section 5, with additional dataset, reward-function, and hyperparameter details in Section B.Theoretical results are proved in Section A.

A.1 OBJECTIVE ESTIMATION ERROR DUE TO LIKELIHOOD APPROXIMATION

Diffusion-based GRPO can suffer exponentially amplified likelihood-approximation error, while wd1 has linear error because its objective avoids policy ratios and uses a weighted log-likelihood formulation.

  • Approximation error: Diffu-GRPO can suffer exponential objective error when approximated current and old policy log-likelihoods have bounded errors.The analysis denotes approximation errors ϵ and ϵ′ and considers the worst case without ratio clipping or regularization.
  • Approximation error: wd1 has only linear approximation error, contrasting with the exponential error analyzed for diffu-GRPO.The bound uses a constant C independent of ϵ and ϵ′.
  • Objective derivation: wd1 derives a weighted log-likelihood objective from reverse-KL-regularized policy optimization.The derivation uses the optimal policy form, expectation identities, group samples, and weight normalization.
  • Diffusion interpretation: The theoretical construction interprets wd1 as energy-guided discrete diffusion training over masked responses.The derivation defines intermediate masked-token marginals, concrete scores, and an advantage-weighted denoising objective.

B.1 DATASET, TRAINING AND EVALUATION PROTOCOL

The experiments train and evaluate wd1 and related methods across GSM8K, MATH500, Sudoku, Countdown, and arithmetic-code settings, comparing accuracy, training cost, and reward dynamics.

  • Datasets and evaluation: wd1 and reproduced d1 use GSM8K, MATH, Sudoku, and Countdown datasets, with task-specific train and evaluation splits.Sudoku uses 1M training puzzles and 256 synthetic evaluation puzzles; Countdown uses a 3-number subset and 256 synthetic questions.
  • Datasets and evaluation: wd1++ uses an OpenR1-derived dataset for MATH500 and the GSM8K train split, with Math-Verify parsing answers during full-parameter fine-tuning.The same answer-verification setup is used for wd1, wd1++, and MDPO.
  • Results and cost: wd1 outperforms d1 by a large margin on Sudoku and Countdown while achieving comparable performance on math problem-solving tasks.The comparison reports test accuracy across tasks and notes that wd1 was selected from fewer checkpoints.
  • Results and cost: wd1 without SFT shows a faster and more stable reward-learning process than wd1-SFT on Sudoku and Countdown, with negligible differences on GSM8K and MATH500.Reward dynamics are reported during training.

C.2 ADDITIONAL ABLATION STUDY

Additional analyses examine reward dynamics, completion length, weight mixing, the exponential-weight coefficient, and an AceCode transfer experiment.

  • Weight ablations: Equal positive and negative weight proportions are presented as the most robust design because imbalances can increase negative-sample likelihood or unlearn high-quality samples.The failure modes arise when sampled completions are uniformly low- or high-reward.
  • Reward and efficiency dynamics: wd1 without SFT demonstrates better rewards than the compared setting in Sudoku and Countdown.This comparison is reported in the reward-dynamics figure.
  • Weight ablations: ψ = 1 is used in all benchmark evaluations, while overly large values such as 10 can cause a performance drop.Small ψ values converge to similar rewards, whereas larger ψ values assign more extreme sample weights.
  • Reward and efficiency dynamics: wd1 produces smaller completion lengths and better token efficiency than d1 on GSM8K and MATH500.The comparison concerns completion-length dynamics during training.
  • Transfer experiment: wd1 achieves consistent improvements over the base model after 200 training steps on AceCode-87K.The experiment follows the Open-R1 implementation.

D LIMITATIONS

The paper identifies scope, reward-distribution, and likelihood-approximation limitations for wd1, while relating its negative-sample term to diffusion data unlearning.

  • Limitations: wd1 may lose effectiveness when all generations in a sampled group receive identical rewards.The paper notes that dataset difficulty can create this condition and suggests reward design and curriculum learning as mitigations.
  • Limitations: The current wd1 framework is restricted to text-based reasoning rather than multimodal or unified diffusion-based models.The paper identifies these extensions as future research directions.
  • Limitations: wd1 relies on the d1-based likelihood approximation, which is computationally efficient but introduces bias.ELBO-based alternatives can add computational overhead and exhibit high variance.
  • Theoretical interpretation: Under DCE, wd1’s negative-sample reinforcement term can be viewed as masked discrete-diffusion data unlearning.The analogy connects minimizing likelihood on lower-advantage samples with a masked diffusion extension of NegGrad.
Loading 2507.08838v2…