SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models (2510.09541v1)

Abstract: Diffusion LLMs (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.

• The paper introduces SPG, a novel RL algorithm that sandwiches the true log-likelihood between ELBO and EUBO to effectively handle both positive and negative rewards.
• It employs a block-wise masking strategy and a mixture objective to reduce gradient variance, leading to significant accuracy improvements across benchmarks like GSM8K, MATH500, Countdown, and Sudoku.
• Empirical results demonstrate faster convergence and robust performance, thereby establishing a new, scalable standard for RL alignment in masked diffusion language models.

Sandwiched Policy Gradient for Masked Diffusion LLMs

Introduction

Diffusion LLMs (dLLMs) have emerged as a competitive alternative to autoregressive (AR) models, offering parallel decoding and reduced inference latency. However, reinforcement learning (RL) alignment for dLLMs is hindered by the intractability of their log-likelihood, which precludes direct application of standard policy gradient methods. Existing approaches rely on surrogates such as the evidence lower bound (ELBO), but these introduce significant bias, especially when negative rewards are present. The paper introduces Sandwiched Policy Gradient (SPG), a novel RL algorithm that leverages both lower and upper bounds of the log-likelihood to provide a more robust and less biased policy gradient for masked diffusion LLMs.

Masked Diffusion LLMs and RL Challenges

Masked diffusion LLMs (MDLMs) operate by progressively corrupting clean text via random masking and training a neural network to reverse this process. The forward process is parameterized by a strictly decreasing noise schedule, and the reverse process is learned by maximizing the ELBO of the log-likelihood. While this framework enables parallel token generation, it complicates RL-based alignment due to the intractability of the true log-likelihood $\log T_\theta(x|c)$.

RL for dLLMs typically seeks to maximize expected reward via policy gradient methods. However, substituting the intractable log-likelihood with its ELBO surrogate is only valid for non-negative rewards, limiting the ability to penalize undesirable outputs and introducing bias in policy optimization. This limitation is particularly acute for advanced RL algorithms that utilize relative or negative rewards.

Sandwiched Policy Gradient: Algorithmic Framework

SPG addresses the log-likelihood estimation challenge by "sandwiching" the true log-likelihood between a tractable lower bound (ELBO) and an upper bound (EUBO). The algorithm maximizes the lower bound for positive-reward sequences and minimizes the upper bound for negative-reward sequences, yielding a valid lower bound for the original RL objective:

$$J_{\text{SPG}}(\theta) = \mathbb{E}_g \left[ \sum_{j=1}^g \left( \mathbb{I}[A_j \geq 0] A_j \text{LELBO}(x_j|c;\theta) + \mathbb{I}[A_j < 0] A_j \text{LEUBO}(x_j|c;\theta) \right) \right]$$

where $A_j$ is the advantage for sample $x_j$.

The EUBO is derived via a Rényi variational bound, with a hyperparameter $\beta$ controlling tightness. In practice, both bounds are estimated via Monte Carlo sampling using a block-wise masking strategy, which aligns the data distribution between policy rollout and optimization.

To further stabilize training and reduce gradient variance, the paper introduces a mixture objective for negative advantage traces:

$$\text{LMix}(x|c;\theta) = w \cdot \text{LEUBO}(x|c;\theta) + (1-w) \cdot \text{LELBO}(x|c;\theta)$$

where $w \in (0,1)$ is a blend coefficient. The mixture approach provides confidence-aware weighting and provably reduces gradient variance compared to using either bound alone.

Implementation Details

Model and Training Setup

• Base Model: LLaDA-8B-Instruct, a state-of-the-art dLLM.
• Benchmarks: GSM8K, MATH500, Countdown, Sudoku.
• RL Fine-Tuning: LoRA adaptation (rank 128, scaling 64), AdamW optimizer, batch size 6 per GPU, gradient accumulation 2, learning rate $3 \times 10^{-6}$, gradient clipping 0.2.
• Rollout: Sequence length 256, 128 diffusion steps, block-wise semi-autoregressive decoding (block size 32), temperature 0.9 (Sudoku: 0.3).
• Monte Carlo Estimation: Number of completions per prompt $g=6$, number of samples $m=2$.

Block-Wise Masking

Block-wise masking divides the sequence into blocks, selects a random block for masking, and keeps earlier blocks clean while fully masking later blocks. Within the selected block, tokens are randomly masked. This strategy matches the semi-autoregressive generation process and improves stability and efficiency of policy optimization.

Pseudocode

for step in range(num_steps): c = sample_prompt() x_group = [model.generate(c) for _ in range(g)] rewards = [reward_fn(c, x) for x in x_group] advantages = compute_advantages(rewards) for _ in range(inner_updates): for x, A in zip(x_group, advantages): z_samples = blockwise_masking(x, m) if A >= 0: grad = A * estimate_LELBO(x, z_samples) else: grad = A * (w * estimate_LEUBO(x, z_samples) + (1-w) * estimate_LELBO(x, z_samples)) update_model(grad)

Empirical Results

SPG achieves state-of-the-art performance across all four benchmarks, with the mixture approach yielding the best results:

• GSM8K: +3.6% accuracy over previous SOTA
• MATH500: +2.6%
• Countdown: +18.4%
• Sudoku: +27.0%

SPG demonstrates faster convergence, higher reward levels, and superior robustness under various inference strategies. Ablation studies confirm the necessity of penalizing negative advantage traces and the superiority of block-wise masking over random masking. The mixture objective consistently outperforms single-bound approaches, both in accuracy and gradient stability.

Theoretical and Practical Implications

The sandwiched objective provides a principled solution to the log-likelihood estimation problem in RL for dLLMs, enabling effective learning from both positive and negative rewards. The block-wise masking strategy ensures alignment between training and inference distributions, which is critical for stable optimization in masked diffusion models. The mixture approach offers a practical trade-off between bias and variance, with theoretical guarantees on variance reduction.

Practically, SPG enables RL-based alignment of dLLMs for complex reasoning tasks, with significant improvements in accuracy and efficiency. The framework is compatible with large-scale models and can be integrated with existing RL algorithms for LLMs. The empirical results suggest that diffusion-based LLMs, when properly aligned via SPG, can match or exceed the performance of AR models on challenging benchmarks.

Future Directions

Potential avenues for future research include:

• Adaptive tuning of the mixture coefficient $w$ and upper bound tightness $\beta$ during training.
• Extension of SPG to multimodal diffusion models and other discrete generative domains.
• Investigation of alternative masking and decoding strategies for further efficiency gains.
• Integration with advanced RL algorithms (e.g., off-policy, meta-RL) and exploration of sample efficiency improvements.
• Analysis of generalization and robustness properties in open-ended reasoning and instruction-following tasks.

Conclusion

SPG provides a robust and theoretically principled RL algorithm for masked diffusion LLMs, resolving the intractable log-likelihood challenge via sandwiched variational bounds and block-wise masking. The approach yields substantial empirical gains on mathematical and logical reasoning tasks, with strong generalization and stability. SPG establishes a new standard for RL alignment in dLLMs and opens new directions for efficient, scalable, and robust LLM training.