pith. sign in

arxiv: 2510.09541 · v3 · submitted 2025-10-10 · 💻 cs.CL · cs.AI

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
show 4 more authors
Feiyu Chen Tommi Jaakkola Yuandong Tian Bo Liu
This is my paper

Pith reviewed 2026-05-18 07:44 UTC · model grok-4.3

classification 💻 cs.CL cs.AI
keywords diffusion language modelspolicy gradientreinforcement learningmasked diffusionmodel alignmentlog-likelihood boundsparallel decoding
0
0 comments X

The pith

Sandwiched Policy Gradient uses upper and lower bounds on log-likelihood to reduce bias in RL for diffusion language models.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

Diffusion large language models can decode multiple tokens in parallel but their intractable log-likelihood blocks standard policy gradient methods for alignment. Prior approaches rely on one-sided approximations such as the ELBO that introduce bias into the gradient estimate. The paper introduces Sandwiched Policy Gradient, which combines an upper bound and a lower bound to bracket the true log-likelihood from both sides. Experiments on math and logic tasks show consistent gains over ELBO and one-step baselines. A sympathetic reader would care because tighter control of gradient bias could make RL training for parallel generative models more reliable and effective.

Core claim

The Sandwiched Policy Gradient estimator sandwiches the intractable log-likelihood of masked diffusion language models between a lower bound and an upper bound, yielding a policy gradient whose bias is smaller than that of ELBO-based or one-step surrogates; this estimator produces higher task accuracy when the model is optimized against downstream rewards on GSM8K, MATH500, Countdown, and Sudoku.

What carries the argument

Sandwiched Policy Gradient (SPG), which combines an upper bound and a lower bound on the true log-likelihood to form a policy gradient estimator.

If this is right

  • RL training of diffusion language models can be performed without relying solely on one-sided lower bounds.
  • Parallel decoding models become easier to align with task-specific rewards such as math problem accuracy.
  • The same bounding technique may extend to other generative architectures whose likelihood is intractable.
  • Training stability improves because the gradient variance is controlled from both sides of the likelihood.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • If the bounds can be tightened further without extra computation, SPG could become a default replacement for ELBO in any diffusion or masked generative model.
  • The method suggests that sandwich estimators might be useful for other latent-variable policies where only bounds on the evidence are available.
  • Applying SPG to larger models might reveal whether the bias reduction scales with model size or sequence length.

Load-bearing premise

The chosen upper and lower bounds remain close enough to the true log-likelihood that the resulting gradient estimate stays sufficiently unbiased for effective policy improvement.

What would settle it

Measure the gap between the SPG estimate and the exact log-likelihood on a small held-out set of sequences; if the gap remains large while the reported accuracy gains disappear, the sandwiching argument would be falsified.

read the original abstract

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.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 2 minor

Summary. The manuscript proposes Sandwiched Policy Gradient (SPG) for reinforcement learning alignment of masked diffusion language models (dLLMs). It replaces one-sided surrogates such as the ELBO with a pair of upper and lower bounds on the intractable log-likelihood to produce a policy gradient estimator with reduced bias. Experiments report accuracy gains of 3.6% on GSM8K, 2.6% on MATH500, 18.4% on Countdown, and 27.0% on Sudoku relative to prior RL methods for dLLMs.

Significance. If the sandwiched bounds are shown to be tight enough to measurably lower gradient bias in the masked diffusion setting, the method would offer a principled route to reward optimization for parallel-decoding language models. The reported task improvements suggest practical value for mathematical reasoning and constraint-satisfaction problems, provided the gains can be attributed to the bias-reduction mechanism rather than implementation details.

major comments (2)
  1. [§3.2] §3.2, Eq. (8)–(10): the upper and lower bounds on the log-likelihood are introduced, yet the manuscript supplies no quantitative evaluation of the bound gap (e.g., average or worst-case difference) across masking schedules or sequence lengths; without this, it remains unclear whether the sandwich controls gradient bias beyond what is already achieved by the ELBO baseline.
  2. [§4.3] §4.3, Table 2: the accuracy improvements are presented, but the section contains no direct measurements of gradient variance, bias estimates, or ablation on bound tightness; consequently the claim that SPG outperforms ELBO/one-step estimators specifically because of lower bias is not yet load-bearingly supported by the reported results.
minor comments (2)
  1. [§2] §2: the related-work discussion of diffusion-model RL omits several recent papers on variational bounds for non-autoregressive models; adding these would strengthen context.
  2. [Figure 3] Figure 3: the caption does not specify the exact masking ratio and reward scaling used for the plotted curves, making reproduction of the variance comparison difficult.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the detailed and constructive feedback. We address each major comment below and describe the revisions we will incorporate to strengthen the empirical support for the bias-reduction claims.

read point-by-point responses

  1. Referee: [§3.2] §3.2, Eq. (8)–(10): the upper and lower bounds on the log-likelihood are introduced, yet the manuscript supplies no quantitative evaluation of the bound gap (e.g., average or worst-case difference) across masking schedules or sequence lengths; without this, it remains unclear whether the sandwich controls gradient bias beyond what is already achieved by the ELBO baseline.

    Authors: We agree that a direct quantitative assessment of the bound gap is needed to demonstrate the tightness of the sandwiched bounds. In the revised manuscript we will add a new subsection (or appendix) to §3.2 that reports the average and worst-case gap between the upper and lower bounds, computed across a range of masking schedules and sequence lengths on the GSM8K and MATH500 validation sets. These measurements will be compared against the ELBO gap to show the additional control provided by the sandwich. revision: yes

  2. Referee: [§4.3] §4.3, Table 2: the accuracy improvements are presented, but the section contains no direct measurements of gradient variance, bias estimates, or ablation on bound tightness; consequently the claim that SPG outperforms ELBO/one-step estimators specifically because of lower bias is not yet load-bearingly supported by the reported results.

    Authors: We acknowledge that the current experimental section lacks direct gradient-bias or variance measurements and an explicit ablation on bound tightness. We will expand §4.3 with (i) proxy estimates of gradient bias obtained by comparing the SPG and ELBO estimators against a high-sample Monte-Carlo reference gradient on a subset of training steps, (ii) reported gradient variance for each estimator, and (iii) an ablation that varies the tightness of the bounds (by adjusting the number of auxiliary samples) and shows the resulting effect on downstream accuracy. These additions will provide more direct evidence linking the observed gains to bias reduction. revision: yes

Circularity Check

0 steps flagged

No circularity: SPG is a proposed estimator validated by external task accuracy

full rationale

The abstract and method description introduce SPG as a new sandwiched upper/lower bound estimator for policy gradients in dLLMs, contrasting it with ELBO and one-step baselines. No equations, derivations, or self-citations are visible that reduce the claimed bias reduction or accuracy lifts (3.6% GSM8K etc.) to a fitted parameter renamed as prediction or to a self-referential definition. The performance results are presented as empirical outcomes on held-out tasks, providing independent falsifiability. The central premise—that the sandwich yields lower-bias gradients—rests on the explicit construction of the bounds rather than on any tautological reduction to inputs. This is the common case of an algorithmic proposal with external benchmarks, so the derivation chain is self-contained.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

The abstract provides no information on free parameters, axioms, or invented entities. No specific numbers, assumptions, or new constructs are described.

pith-pipeline@v0.9.0 · 5736 in / 1200 out tokens · 41032 ms · 2026-05-18T07:44:17.532333+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

What do these tags mean?
matches
The paper's claim is directly supported by a theorem in the formal canon.
supports
The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends
The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses
The paper appears to rely on the theorem as machinery.
contradicts
The paper's claim conflicts with a theorem or certificate in the canon.
unclear
Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Forward citations

Cited by 10 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. Reinforcement Learning for Diffusion LLMs with Entropy-Guided Step Selection and Stepwise Advantages

    cs.LG 2026-03 unverdicted novelty 8.0

    Derives an exact unbiased policy gradient for RL post-training of diffusion LLMs via entropy-guided step selection and one-step denoising rewards, achieving state-of-the-art results on coding and logical reasoning benchmarks.

  2. Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models

    cs.LG 2026-05 conditional novelty 7.0

    TraFL applies trajectory flow balancing to post-train diffusion language models, preventing mode collapse and delivering consistent gains on reasoning tasks that hold under increased sampling.

  3. Relative Score Policy Optimization for Diffusion Language Models

    cs.CL 2026-05 unverdicted novelty 7.0

    RSPO interprets reward advantages as targets for relative log-ratios in dLLMs, calibrating noisy estimates to stabilize RLVR training and achieve strong gains on planning tasks with competitive math reasoning performance.

  4. ReflectDrive-2: Reinforcement-Learning-Aligned Self-Editing for Discrete Diffusion Driving

    cs.RO 2026-05 unverdicted novelty 7.0

    ReflectDrive-2 achieves 91.0 PDMS on NAVSIM with camera input by training a discrete diffusion model to self-edit trajectories via RL-aligned AutoEdit.

  5. Discrete Tilt Matching

    cs.LG 2026-04 unverdicted novelty 7.0

    Discrete Tilt Matching recasts dLLM fine-tuning as state-level matching of tilted local unmasking posteriors, producing a stable weighted cross-entropy loss that improves Sudoku and Countdown performance when applied ...

  6. Discrete Tilt Matching

    cs.LG 2026-04 unverdicted novelty 7.0

    DTM recasts dLLM fine-tuning as weighted cross-entropy matching of tilted local posteriors, with demonstrated gains on Sudoku and math tasks.

  7. ReflectDrive-2: Reinforcement-Learning-Aligned Self-Editing for Discrete Diffusion Driving

    cs.RO 2026-05 unverdicted novelty 6.0

    ReflectDrive-2 combines masked discrete diffusion with RL-aligned self-editing to generate and refine driving trajectories, reaching 91.0 PDMS on NAVSIM camera-only and 94.8 in best-of-6.

  8. Diffusion-State Policy Optimization for Masked Diffusion Language Models

    cs.CL 2026-02 unverdicted novelty 6.0

    DiSPO optimizes intermediate decisions in masked diffusion LMs by branching at selected masked states, resampling tokens, scoring completions, and updating only new tokens using a derived policy-gradient estimator tha...

  9. Diffusion-State Policy Optimization for Masked Diffusion Language Models

    cs.CL 2026-02 unverdicted novelty 6.0

    DiSPO is a plug-in credit-assignment method for masked diffusion LMs that optimizes intermediate filling decisions via branched completions from rollout-cached logits.

  10. LLaDA2.0: Scaling Up Diffusion Language Models to 100B

    cs.LG 2025-12 conditional novelty 6.0

    LLaDA2.0 scales discrete diffusion language models to 100B parameters via systematic conversion from autoregressive models using a 3-phase WSD training scheme and releases open-source 16B and 100B MoE variants.

Reference graph

Works this paper leans on

15 extracted references · 15 canonical work pages · cited by 7 Pith papers

  1. [1]

    The diffusion duality.arXiv preprint arXiv:2506.10892, 2025

    Accessed: 2025-01-24. William Peebles and Saining Xie. Scalable diffusion models with transformers. InProceedings of the IEEE/CVF international conference on computer vision, pages 4195–4205, 2023. Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Yourlanguagemodelissecre...

    work page arXiv 2025

  2. [2]

    The cost of the first 8 glasses, each costing $5. 29

    work page

  3. [3]

    First, let’s calculate the cost of the first 8 glasses:8×5 = 40dollars

    The cost of the next 8 glasses, each costing $3. First, let’s calculate the cost of the first 8 glasses:8×5 = 40dollars. Next, let’s calculate the cost of the next 8 glasses:8×3 = 24dollars. Now, we add these two amounts together:40 + 24 = 64dollars. Therefore, Kylar needs to pay $64 for them. </reasoning> <answer> 64 </answer> Question:Amy had two eyesha...

    work page

  4. [4]

    The first glass costs $5

    work page

  5. [5]

    Every second glass costs 60% of $5, which is0.6×5 = 3dollars

    work page

  6. [6]

    Kylar wants to buy 16 glasses, so he will have 8 full-priced glasses and 8 discounted glasses

    work page

  7. [7]

    The cost for 8 full-priced glasses is8×5 = 40dollars

    work page

  8. [8]

    The cost for 8 discounted glasses is8×3 = 24dollars

    work page

  9. [9]

    </reasoning> <answer> 64 </answer> Question:Amy had two eyeshadow palettes with four colors each and three makeup sets that came with six eyeshadow colors each

    Adding these the total cost is40 + 24 = 64dollars. </reasoning> <answer> 64 </answer> Question:Amy had two eyeshadow palettes with four colors each and three makeup sets that came with six eyeshadow colors each. Her sister steals one of the palettes. Amy uses up half of the colors from one makeup set. How many eyeshadow colors does she have left? Model:SP...

    work page

  10. [10]

    Amy initially has two eyeshadow palettes with four colors each:2×4 = 8colors

    work page

  11. [11]

    She also has three makeup sets with six colors each:3×6 = 18colors

    work page

  12. [12]

    Adding these together gives gives:8 + 18 = 26colors

    work page

  13. [13]

    Her sister steals one palette, so she loses:4colors

    work page

  14. [14]

    Amy uses up half of the colors from one makeup set:6 2 = 3colors

    work page

  15. [15]

    Subtracting these lost and used colors:26−4−3 = 19colors </reasoning> <answer> 19 </answer> 32

    work page