PaperlayerView exact revision on arXiv

Source-linked AI summary

SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

Chenyu Wang, Paria Rashidinejad, DiJia Su, Song Jiang, Sid Wang, Siyan Zhao, Cai Zhou, Shannon Zejiang Shen, Feiyu Chen, Tommi Jaakkola, Yuandong Tian, Bo Liu

arXiv:2510.09541v3cs.CLcs.AI

TL;DR

dLLM reinforcement learning is hindered by intractable log-likelihoods and biased one-sided surrogates. SPG sandwiches the likelihood with lower and upper bounds according to reward sign, using practical masking and mixture estimation. It outperforms prior diffusion-language-model RL methods across four reasoning benchmarks, with gains of up to 27.0%.

  • Problem

    dLLM reinforcement learning cannot directly use standard policy gradients because log-likelihood is intractable, while ELBO surrogates can bias updates and inadequately handle negative rewards.

  • Method

    SPG maximizes an ELBO for positive-reward sequences and minimizes an evidence upper bound for negative-reward sequences, supported by block-wise masking and bound mixtures.

  • Results

    SPG achieves state-of-the-art RL performance for dLLMs, improving accuracy by 3.6% on GSM8K, 2.6% on MATH500, 18.4% on Countdown, and 27.0% on Sudoku.

  • Takeaways & Limitations

    Using both likelihood bounds enables reward-aware policy optimization for dLLMs and supports strong results across mathematical and logical reasoning tasks.

  • Takeaways & Limitations

    The Monte Carlo estimate of the evidence upper bound is slightly biased because the logarithm is applied outside the expectation.

Abstract

from arXiv · show

Diffusion large language models (dLLMs) are emerging as an efficient alternative to autoregressive models due to their ability to decode multiple tokens in parallel. However, aligning dLLMs with human preferences or task-specific rewards via reinforcement learning (RL) is challenging because their intractable log-likelihood precludes the direct application of standard policy gradient methods. While prior work uses surrogates like the evidence lower bound (ELBO), these one-sided approximations can introduce significant policy gradient bias. To address this, we propose the Sandwiched Policy Gradient (SPG) that leverages both an upper and a lower bound of the true log-likelihood. Experiments show that SPG significantly outperforms baselines based on ELBO or one-step estimation. Specifically, SPG improves the accuracy over state-of-the-art RL methods for dLLMs by 3.6% in GSM8K, 2.6% in MATH500, 18.4% in Countdown and 27.0% in Sudoku.

1 Introduction

dLLMs offer lower-latency parallel decoding, but reinforcement-learning alignment is difficult because their intractable log-likelihood makes standard policy-gradient estimation unavailable. SPG addresses this with reward-aware upper and lower likelihood bounds and block-wise masking, achieving state-of-the-art results across reasoning and coding benchmarks.

  • dLLMs can decode multiple tokens in parallel, potentially reducing inference latency relative to autoregressive models.
  • RL alignment of dLLMs is challenging because their intractable log-likelihood is needed for accurate policy-gradient estimation.
  • SPG improves accuracy over state-of-the-art diffusion-language-model RL methods by 3.6% on GSM8K, 2.6% on MATH500, 18.4% on Countdown, and 27.0% on Sudoku.
  • SPG maximizes a lower likelihood bound for positive-reward responses and minimizes an upper bound for negative-reward responses.
  • Block-wise masking improves stability of the training objective’s estimation, while SPG also reports state-of-the-art results on four reasoning and two coding benchmarks.

2 Background

Masked diffusion language models generate text by reversing progressive token masking, and RL treats this generation process as a policy optimized by rewards. Existing ELBO-based likelihood surrogates are problematic because the lower bound does not support effective learning from negative rewards.

  • MDLMs progressively replace clean tokens with [mask] and train a policy to reverse this corruption by predicting original tokens.
  • RL models dLLM generation as a policy that acts on partially masked sequences and receives a scalar reward for each complete response.
  • The policy’s log-likelihood is intractable, so prior work substitutes the ELBO or one-step estimates for the true likelihood.
  • Because ELBO ≤ log πθ, minimizing ELBO for negative rewards does not guarantee reducing true likelihood and can bias the learned policy.
  • SPG’s algorithm samples grouped completions, computes rewards and advantages, estimates bounds with masked samples, and updates the policy using the resulting gradient.

3 Sandwiched Policy Gradient with Evidence Bounds

SPG constructs a valid policy-gradient proxy by using an ELBO for positive advantages and an evidence upper bound for negative advantages. Monte Carlo estimation, block-wise masking, and a mixture of bounds improve practical stability, with the mixture offering lower gradient variance under the stated condition.

  • 3.1 A Lower Bound Objective for Policy Optimization: SPG uses the ELBO for positive-advantage samples and a tractable evidence upper bound for negative-advantage samples, yielding a lower bound on the original objective.
  • 3.2 A Tractable Evidence Upper Bound: The evidence upper bound is derived from a Rényi variational bound, with β ≥ 1 controlling tightness and values closer to 1 producing a tighter bound.
  • 3.2 A Tractable Evidence Upper Bound: Monte Carlo estimation of the upper bound is biased because the logarithm lies outside the expectation, so SPG retains a tighter but slightly biased formulation.
  • 3.3 Practical Considerations: Block-wise masking selects a random generation block, keeps earlier blocks clean, fully masks later blocks, and lightly masks prompt or clean blocks with probability 0.15.
  • 3.3 Practical Considerations: The negative-trace mixture combines the upper and lower bounds because the upper bound supplies stronger correction while ELBO estimation is easier and more stable with few samples.
  • 3.3 Practical Considerations: When the stated variance condition holds, the optimal mixture gradient has strictly smaller coordinate-wise variance than either bound alone.

4 Experiments

SPG is evaluated against recent diffusion-language-model RL baselines on reasoning and coding tasks, with ablations examining its likelihood estimation, masking, hyperparameters, and training dynamics. Across the main benchmarks, SPG w/ Mixture delivers the strongest reported results, while component studies support the roles of mixture estimation and block-wise masking.

  • 4.1 Experimental Setup and Main Results: SPG w/ Mixture converges faster and reaches higher rewards than D1, WD1, and UniGRPO during RL training.Reward curves report means and standard deviations over rolling windows of 50 steps.
  • 4.1 Experimental Setup and Main Results: 3.6%, 2.6%, 18%, and 27% test-accuracy gains over the previous state of the art are reported on GSM8K, MATH500, Countdown, and Sudoku, respectively, at generation length 256.SPG w/ Mixture is the strongest variant across the four reasoning benchmarks.
  • 4.1 Experimental Setup and Main Results: SPG consistently improves accuracy over baselines on HumanEval and MBPP across different generation lengths.The coding evaluation uses LLaDA-8B-Instruct fine-tuned on KodCode-Light-RL-10K.
  • 4.2 Ablations and Further Analysis: Removing the negative-advantage loss substantially reduces performance, while Mixture and EUBO generally outperform ELBO for negative advantage traces.EUBO is an exception on MATH500 in the reported comparison.
  • 4.2 Ablations and Further Analysis: Block-wise masking consistently outperforms random masking in Monte Carlo estimation, supporting alignment between rollout and optimization input distributions.The comparison is reported across the evaluated benchmarks and likelihood estimators.

5 Conclusion

The paper presents SPG as an RL algorithm for masked diffusion language models and situates it within diffusion-language-model and RL research. Its method combines bound-based likelihood optimization with masking-based estimation, while the supplied conclusion emphasizes improved performance over baselines.

  • 5 Conclusion: SPG maximizes a tractable lower bound for positive-reward sequences and minimizes an upper bound for negative-reward sequences.The approach is intended to address intractable dLLM log-likelihoods in policy-gradient training.
  • 5 Conclusion: SPG also uses block-wise masking for Monte Carlo estimation to improve optimization stability and efficiency.The broader masked-diffusion setup models sequences through progressively corrupted and denoised states.
  • 5 Conclusion: Experiments on four mathematical and logical reasoning benchmarks report significant improvements over baselines and state-of-the-art performance.The supplied conclusion summarizes the empirical scope without listing benchmark-specific values.

C Detailed Overview on RL Algorithms for dLLMs

Existing dLLM RL methods differ in policy objectives, likelihood estimators, and masking strategies, while SPG combines reward-dependent variational bounds with block-wise masking. Compared with trajectory-level updates, SPG achieves stronger results on the supplied benchmarks with lower computational requirements.

  • Baseline methods: Earlier baselines estimate likelihoods with ELBO and random masking, one-step prediction, or modified GRPO-style objectives.LLaDA-1.5 and UniGRPO use ELBO-based estimation, whereas D1 and WD1 use one-step prediction.
  • SPG: SPG uses ELBO for positive traces and EUBO or an ELBO–EUBO mixture for negative traces, with block-wise masking and prompt masking during likelihood estimation.Its policy-gradient objective uses group-relative advantage.
  • Scope: Learning an optimal inference-time unmasking order is complementary to RL algorithm improvements and is left for future work.LLaDOU adds a decoding-order prediction module, which SPG does not incorporate.
  • Comparison with trajectory-level updates: SPG consistently outperforms the StepWise trajectory-level method on MATH500 and Countdown.StepWise backpropagates through every diffusion timestep, increasing computational cost; the supplied table compares the two methods on these benchmarks.

D Evidence Upper Bound for dLLMs

This section derives an evidence upper bound for masked diffusion log-likelihoods from a Rényi variational bound and extends the result to sequences and continuous time. The resulting estimator is tighter than an unbiased relaxation but can have biased Monte Carlo estimates.

  • Derivation: Theorem 1 establishes an evidence upper bound for masked diffusion sequences under a monotonic denoising schedule and β ≥ 1.The derivation begins from a Rényi variational bound and treats categorical sequences through diffusion transitions.
  • Sequence bound: For sequence length n, the upper bound satisfies log πθ(x1:n) ≤ LEUBO(x1:n; θ), with zero constant term when the rescaled hyperparameter satisfies β̃ ≥ n.For β̃ < n, an additional constant term remains.
  • Continuous-time extension: The construction extends to continuous time by dropping the θ-independent constant and taking the limit T →∞.This parallels the derivation used for the lower-bound objective.
  • Caveat: Because the logarithm lies outside the expectation, Monte Carlo estimates of the upper bound are generally biased.An inequality-based unbiased relaxation is possible, but it loosens the bound and performs worse in practice because the gap to the true likelihood can increase.

E.1 Proof of Proposition 1

The proposition shows that mixing lower- and upper-bound gradients can reduce coordinate-wise variance. Under a non-degeneracy condition, the variance is strictly convex in the mixture coefficient and has a unique minimizer.

  • Gradient form: The mixture gradient uses a convex combination of lower- and upper-bound weighting terms multiplying the score-function gradient.This form makes ω the control parameter for blending the two estimators.
  • Variance minimization: If Var((ρβ − w(t, zt))∂θk log πθ(x | zt, c)) > 0, the mixture-gradient variance is a strictly convex quadratic in ω with a unique minimizer ω⋆.The minimizer is defined coordinate-wise for the mixture objective.
  • Proposition 1: At the optimal coefficient, the mixture has strictly smaller coordinate-wise variance than either the lower-bound or upper-bound gradient.The result applies in the non-degenerate case where the relevant variance is positive.

E.2 Additional Comparison Between the Mixture Loss and the Lower and Upper Bounds

The mixture objective combines lower- and upper-bound behaviors to produce confidence-aware, non-collapsing weights and lower, more stable gradient norms. Toy landscapes further show that the two bounds emphasize different likelihood components depending on optimization direction.

  • Gradient stability: SPG with the mixture achieves lower and more stable gradient norms than alternative likelihood estimators during training.Figure 7 reports rolling-window means and standard deviations over 50 steps.
  • Mixture weighting: The mixture maintains per-sample weights by flooring very small EUBO weights while capping large weights through uncertainty-aware interpolation.This avoids near-zero EUBO updates and controls weighting sharpness with β and blending with ω.
  • Toy example: Figure 8 contrasts −LELBO and −LEUBO landscapes, with the upper-bound landscape flatter in low-value regions and sharper in high-value regions.The corresponding optimization behavior follows from whether the objective is maximizing or minimizing the likelihood surrogate.
  • Toy example: For negative advantage, EUBO preferentially minimizes the larger likelihood, producing a more balanced decrease than LELBO in the described extreme case.LELBO can diverge when either likelihood component reaches zero while the other remains insufficiently reduced.
  • Toy example: For positive advantage, LELBO increases the smaller likelihood effectively, whereas EUBO focuses on the larger likelihood and is less efficient for this objective.The toy example uses a two-token sequence with likelihood components represented by a and b.

F Additional Discussions and Future Directions on the RL Algorithm

The discussion examines unresolved bias and stabilization issues in reinforcement learning for diffusion language models. It also reports that reward shifting underperforms SPG and identifies correct RL stabilization and advantage estimation as future directions.

  • RL Stabilization: Correctly adapting importance sampling, clipping, and KL regularization to dLLMs remains unresolved because their log-likelihood is intractable.The authors retain the simplest SPG formulation and leave principled KL and trust-region optimization for future work.
  • Advantage Estimation: Group-relative advantage calculation can introduce bias, motivating RLOO or other methods as directions for reducing policy-update bias.The discussion notes that RLOO uses a leave-one-out reward expectation and remains unbiased in off-policy settings.
  • The Reward Shifting Alternate: Reward shifting is theoretically applicable when episodic rewards are terminal-only or KL regularization is used, but the method breaks down with partial rewards.
  • The Reward Shifting Alternate: Reward shifting underperforms SPG under both REINFORCE and GRPO objectives on MATH500 and Countdown.The reward-shifted objective normalizes each trace’s advantage using the batch minimum while estimating log-likelihood with the ELBO.

G.1 Datasets and Reward Functions

The experiments use established reasoning and coding benchmarks with task-specific rewards for correctness, formatting, code tests, and safety. Training follows prior work’s settings while using fixed rollout and optimization configurations across tasks.

  • Dataset and reward setup: The evaluation follows D1 and WD1 reward functions and train-test splits, except for Sudoku.Rewards encourage correctness and proper formatting with task-dependent granularity.
  • GSM8K: GSM8K combines XML structure, soft and strict formatting, integer-answer, and correctness rewards.The correctness component awards +2.0 when the answer matches ground truth.
  • MATH500 and Countdown: MATH500 uses format and correctness rewards, while Countdown rewards exact target-reaching expressions more strongly than partially correct number usage.Countdown assigns +1.0 for exact solutions, +0.1 for correct numbers without the target, and +0.0 otherwise.
  • Sudoku: Sudoku training uses 694,006 puzzles from 200 solutions, while evaluation samples puzzles from the remaining 88 solutions to reduce leakage and memorization.Because zero-shot performance was below 7%, all Sudoku experiments use 3-shot prompting with disjoint solution examples.
  • Coding: Coding uses structure and correctness rewards together with a safe-code guarantee that assigns zero reward when blocked modules are used.Correctness is measured by the fraction of unit tests passed.
  • Training configuration: Rollouts use 256-token sequences, 128 diffusion steps, block size 32, temperature 0.9 except for Sudoku, six completions per prompt, and two Monte Carlo samples.Training runs for 6000 steps on GSM8K, Countdown, and coding, 4000 on MATH500, and 2500 on Sudoku.

H.1 Additional Evaluations to the Main Results

Additional evaluations reproduce the main comparison across four reasoning benchmarks and analyze completion length. SPG consistently outperforms the other evaluated models while adapting its effective sequence length to task difficulty.

  • Complete evaluation results: SPG consistently outperforms all other models across the four reasoning benchmarks in the complete evaluation results.The comparison includes reported and reproduced results for D1 and WD1; SPG uses a single checkpoint selected by average accuracy across generation lengths.
  • Dynamics of Completion Length: SPG adapts effective sequence length to task difficulty and makes effective use of the available generation length during RL training.The dynamics are reported over training and at the best checkpoint for the four reasoning benchmarks.

H.2 Additional Ablation Results

The ablations test masking, likelihood bounds, sampling, inference, fine-tuning, and computational cost. Results support block-wise masking and tighter bounds, show strong Pass@K performance, and indicate robustness across training choices.

  • Algorithm components: SPG w/ Mixture consistently outperforms other likelihood estimation methods for negative advantage traces, while block-wise masking improves over random masking on both benchmarks.These comparisons are reported in the log-likelihood and masking ablation tables.
  • Likelihood bounds: The tighter upper bound performs better than the looser unbiased bound, which may suffer from greater discrepancy from the true log-likelihood.For positive advantage traces, using the lower bound performs better than using the upper bound.
  • Pass@K performance: SPG achieves the best Pass@K performance across settings, including 55.6% Pass@4 on MATH500 and 76.6% on Countdown for SPG w/ Mixture.Evaluations use temperature 0.9 and generation length 256; Pass@K improves for all models as K increases from 1 to 4.
  • Full fine-tuning: Full fine-tuning yields comparable or better performance than LoRA for SPG w/ Mixture, supporting robustness across fine-tuning choices.The full-fine-tuning experiments use smaller learning rates than the LoRA experiments.
  • Sampling and masking: Prompt masking has minimal impact on SPG, while increasing Monte Carlo samples from m = 2 to m = 4 gives similar MATH performance and better Countdown performance.The m = 4 models also show lower gradient norms throughout training than the m = 2 models.
  • Computational cost: Block-wise masking adds no overhead relative to random masking, while Monte Carlo sampling adds computation relative to one-step estimation but rollout remains the primary cost.

H.3 Additional Evaluations

Additional evaluations examine SPG’s components, checkpoint-selection robustness, inference settings, computational cost, and qualitative reasoning behavior. Across validation-based evaluation and inference-strategy ablations, SPG consistently outperforms baselines, while examples show more concise and accurate reasoning.

  • Component Ablations: Ablations identify improved log-likelihood estimation and block-wise masking as important contributors to SPG’s performance.The component studies hold other settings constant, and block-wise masking also improves UniGRPO while remaining below SPG.
  • Ablations: SPG maintains superior performance across different inference strategies, while additional studies vary fine-tuning, prompt masking, and Monte Carlo sampling settings.The supplied ablation tables cover inference strategy, full fine-tuning, prompt masking, and the number of Monte Carlo samples.
  • Computational Cost: The evaluation also reports per-update computational time for each method on GSM8K using 8×A100 GPUs.The computational comparison measures average time in minutes per gradient update.
  • Validation-Based Checkpoint Selection: SPG consistently surpasses all baselines when checkpoints are selected using a validation split and evaluated across multiple generation lengths.The original test set is divided into validation and test halves, with average accuracy reported across generation lengths 128, 256, and 512.
  • Qualitative Analysis: Qualitative GSM8K examples indicate that SPG produces more concise and accurate reasoning than other models, although correct answers can still rely on flawed reasoning.The examples use generation length 256 and 128 diffusion steps; the paper notes that such artifacts are expected to diminish with a stronger base model.
Loading 2510.09541v3…