💬 NLP

SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

This paper introduces Sandwiched Policy Gradient (SPG), a novel reinforcement learning method that utilizes both upper and lower bounds of the log-likelihood to overcome the intractability of diffusion language models, significantly outperforming existing baselines on mathematical and logical reasoning benchmarks.

Original authors: 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

Published 2026-04-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: 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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: A New Way to Teach AI to Think

Imagine you have a very smart student (the AI) who is great at writing stories, but they write them one word at a time, like a typist. This is how most AI models work today (called Autoregressive).

Recently, a new type of student has emerged called the Diffusion Model. Instead of typing word-by-word, this student looks at a blank page full of scribbles (noise) and tries to "clean it up" all at once to reveal the final story. This is much faster because they can fix many words simultaneously.

However, there's a problem: How do we teach this "clean-up" student to get better at math or logic?

In the past, we used a method called Reinforcement Learning (RL). Think of this as a teacher giving the student a grade (a reward) after they finish a test.

If the answer is right, the teacher says, "Great job! Do more of that!"
If the answer is wrong, the teacher says, "Bad job! Stop doing that!"

The problem with Diffusion Models is that the teacher can't easily calculate the grade for a specific answer while the student is still "cleaning up" the noise. It's like trying to grade a painting while the artist is still mixing the paint on the canvas. Because the teacher can't see the exact score, they have to guess. Previous methods made a "best guess" that was often wrong, leading the student to learn the wrong lessons.

The Solution: The "Sandwich" Strategy

The authors of this paper propose a new teaching method called SPG (Sandwiched Policy Gradient).

Imagine the teacher is trying to estimate the "true quality" of the student's answer, but they can't see it directly. Instead, they use a Sandwich:

The Bottom Bun (Lower Bound): When the student gets a good answer (high reward), the teacher uses a "safe, conservative estimate" (a Lower Bound). They say, "We know this is at least this good. Let's encourage it!"
The Top Bun (Upper Bound): When the student gets a bad answer (low reward), the teacher uses a "strict, pessimistic estimate" (an Upper Bound). They say, "We know this is definitely no better than this. Let's discourage it!"

By "sandwiching" the true, invisible quality between a safe floor and a strict ceiling, the teacher can guide the student much more accurately than before.

Old Method: The teacher only looked at the bottom bun for everything. This meant they were too nice to bad answers and didn't punish them enough.
SPG Method: The teacher uses the bottom bun for good answers and flips the sandwich to use the top bun for bad answers. This creates a perfect balance.

The Secret Sauce: Block-Wise Masking

There was another issue. When the teacher tried to guess the quality of the answer, they were looking at the student's work in a chaotic, random way. It was like asking a student to fix a sentence by randomly erasing words one by one, which confused the student.

The authors introduced a Block-Wise Masking strategy.

The Analogy: Imagine the student is writing a story. Instead of erasing random words from the whole page, the teacher asks them to focus on one paragraph at a time.
The teacher keeps the beginning of the story clean, erases the middle paragraph (the "block"), and keeps the end clean.
This helps the student understand the context better. It's like practicing a song by focusing on one specific measure of music rather than random notes scattered across the whole sheet. This makes the learning process much more stable and efficient.

The Results: Smarter, Faster, and More Accurate

The paper tested this new "Sandwich" teacher on four difficult challenges:

GSM8K & MATH500: Complex math word problems.
Countdown: A game where you have to use numbers to reach a target.
Sudoku: The classic number puzzle.

The Results:
The SPG method crushed the competition.

On Sudoku, it improved accuracy by a massive 27% compared to the previous best method.
On Countdown, it jumped 18% higher.
On math problems, it improved by 3-4%.

In the world of AI, a 3% improvement on math problems is huge. It means the AI is suddenly much more reliable at solving logic puzzles.

Summary

The Problem: Diffusion AI models are fast but hard to train with rewards because we can't easily calculate their "score."
The Fix: The SPG Sandwich. We use a "safe floor" to encourage good answers and a "strict ceiling" to punish bad answers, giving the AI a clearer signal on how to improve.
The Trick: We teach the AI in blocks (chunks) rather than random pieces, making the training process smoother.
The Outcome: The AI becomes significantly smarter at math and logic, proving that this new "Sandwich" approach is the future of training fast, parallel AI models.

In short, the authors found a way to give a fast-but-confused AI student a much clearer set of instructions, turning it into a math and logic genius.

1. Problem Statement

Diffusion Large Language Models (dLLMs), particularly Masked Diffusion Language Models (MDLMs), offer a promising alternative to autoregressive (AR) models by enabling parallel token decoding, which reduces inference latency. However, aligning dLLMs with human preferences or task-specific rewards via Reinforcement Learning (RL) remains a significant challenge.

The core difficulty lies in the intractability of the log-likelihood ( $\log \pi_\theta(x|c)$ ) for diffusion models. Standard policy gradient methods (e.g., PPO, GRPO) rely on computing this log-likelihood to update the policy.

Current Limitations: Prior works approximate the true log-likelihood using the Evidence Lower Bound (ELBO) or one-step estimations.
The Flaw: The ELBO is a lower bound ( $ELBO \leq \log \pi_\theta$ $E L B O \leq lo g π_{θ}$ ).
- For positive rewards (good outputs), maximizing the ELBO is a valid proxy for maximizing the true likelihood.
- For negative rewards (bad outputs), minimizing the ELBO does not guarantee minimizing the true log-likelihood. This leads to biased policy gradients, preventing the model from effectively learning from negative feedback and resulting in suboptimal performance.

2. Methodology: Sandwiched Policy Gradient (SPG)

The authors propose SPG, a novel RL algorithm that addresses the intractability issue by "sandwiching" the true log-likelihood between a tractable lower bound and a tractable upper bound, depending on the reward signal.

A. The Sandwiched Objective

Instead of using a single surrogate for all samples, SPG adapts the optimization direction based on the advantage ( $A_j$ ) of the generated sequence:

Positive Advantage ( $A_j \geq 0$ ): The model maximizes the Evidence Lower Bound (ELBO). This encourages the generation of high-reward sequences.
Negative Advantage ( $A_j < 0$ ): The model minimizes a tractable Evidence Upper Bound (EUBO). This ensures that low-reward sequences are effectively penalized by reducing their true log-likelihood, which the ELBO fails to do.

The objective function is formulated as:
$J_{SPG}(\theta) = \mathbb{E} \left[ \frac{1}{g} \sum_{j=1}^g \left( \mathbb{1}_{A_j \geq 0} \cdot A_j \mathcal{L}_{ELBO}(x_j) + \mathbb{1}_{A_j < 0} \cdot A_j \mathcal{L}_{EUBO}(x_j) \right) \right]$

B. Evidence Upper Bound (EUBO) Derivation

The paper derives a tractable EUBO for masked diffusion models based on the Rényi variational bound.

The bound involves a hyperparameter $\beta \geq 1$ that controls tightness.
The formulation places the logarithm outside the expectation, which introduces a slight bias in Monte Carlo estimation but provides a tighter bound than unbiased looser alternatives.
The gradient of the EUBO naturally down-weights uncertain tokens and up-weights confident ones, acting as a strong correction signal for negative traces.

C. Mixture Strategy for Stability

To balance the stability of the ELBO (easier to estimate) and the effectiveness of the EUBO (better for penalization), SPG introduces a Mixture Loss for negative advantage traces:
$\tilde{\mathcal{L}}_{Mix} = \omega \cdot \tilde{\mathcal{L}}_{EUBO} + (1-\omega) \cdot \mathcal{L}_{ELBO}$

Theoretical Benefit: Proposition 1 proves that this mixture strictly reduces the variance of the policy gradient compared to using either bound alone.
Mechanism: The mixture creates a "confidence-aware" weighting that prevents vanishing gradients (a risk with EUBO alone) while avoiding the systematic bias of ELBO alone.

D. Block-Wise Masking Strategy

To improve the stability of the Monte Carlo estimation required for these bounds, the authors propose a block-wise masking strategy for generating perturbed samples during training:

Instead of random masking across the whole sequence, the sequence is divided into blocks.
A random block is selected to be partially masked, while preceding blocks remain clean and following blocks are fully masked.
This aligns the data distribution during the policy rollout (which often uses block-wise semi-autoregressive decoding) with the distribution used for likelihood estimation, significantly improving training stability.

3. Key Contributions

SPG Algorithm: A new policy gradient algorithm that resolves the bias in dLLM RL by using a lower bound for positive rewards and an upper bound for negative rewards.
Theoretical Derivation: A rigorous derivation of the Evidence Upper Bound (EUBO) for masked diffusion models and a proof that mixing EUBO and ELBO reduces gradient variance.
Block-Wise Masking: A practical technique to align training and inference distributions, enhancing optimization stability.
State-of-the-Art Performance: Demonstrated superior results across mathematical reasoning, logical reasoning, and coding benchmarks.

4. Experimental Results

The authors evaluated SPG on LLaDA-8B-Instruct across four reasoning benchmarks and two coding benchmarks, comparing it against baselines like D1, WD1, UniGRPO, and LLaDA-1.5.

Key Performance Gains (at 256 sequence length):

GSM8K (Math): Improved accuracy by 3.6% over the previous SOTA (UniGRPO).
MATH500 (Math): Improved by 2.6%.
Countdown (Logic): Improved by 18.4%.
Sudoku (Logic): Improved by 27.0% (a massive gain, highlighting the difficulty of the task for previous methods).
Coding (HumanEval/MBPP): Improved by 1.9% and 4.7% respectively.

Ablation Studies:

Mixture vs. Pure Bounds: The mixture approach ( $\omega \in (0,1)$ ) consistently outperformed using only EUBO or only ELBO, confirming the variance reduction theory.
Masking Strategy: Block-wise masking significantly outperformed random masking, particularly on the Countdown task (+23.9% gain).
Robustness: SPG maintained superior performance across various inference strategies (different block sizes, random vs. confidence-based unmasking), demonstrating strong generalization.

5. Significance

This paper represents a critical step forward in the development of Diffusion Language Models.

Solving the RL Bottleneck: It provides the first principled solution to the intractable log-likelihood problem in dLLM reinforcement learning, enabling effective learning from both positive and negative feedback.
Parity with Autoregressive Models: By achieving SOTA results on complex reasoning tasks, SPG demonstrates that dLLMs can be competitive with, and in some logical tasks superior to, autoregressive models when properly aligned.
Efficiency: The method achieves these gains without requiring trajectory-level backpropagation (which is computationally expensive), making it scalable for large models.

In summary, SPG establishes a new standard for training diffusion-based LLMs, proving that with the correct theoretical bounds and estimation strategies, these models can be effectively optimized for complex reasoning and alignment tasks.