SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

Chenyu Wang, Paria Rashidinejad, DiJia Su, Song Jiang, Sid Wang +7 more
10/10/2025
cs.CLcs.AI

No official code found. Best implementation: facebookresearch/SPG (65 stars, MIT, last commit 2025-10-29).

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.

View on arXivView PDF

Code Implementations(1)

facebookresearch/SPGCommunity71%
MIT

Code for paper "SPG Sandwiched Policy Gradient for Masked Diffusion Language Models"

656Shell, PythonSep 26, 202511 months agoMIT

Similar Papers

Simple Policy Gradients for Reasoning with Diffusion Language Models

91% match1 repo62025

wd1: Weighted Policy Optimization for Reasoning in Diffusion Language Models

91% match2025

Improving Reasoning for Diffusion Language Models via Group Diffusion Policy Optimization

89% match2025

d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning

89% match1 repo4592025

Learning Unmasking Policies for Diffusion Language Models

89% match2025

Boundary-Guided Policy Optimization for Memory-efficient RL of Diffusion Large Language Models

89% match1 repo102025

Cite this paper

@article{wang2025spg,
  title  = {SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models},
  author = {Chenyu Wang and Paria Rashidinejad and DiJia Su and Song Jiang and Sid Wang and Siyan Zhao and Cai Zhou and Shannon Zejiang Shen and Feiyu Chen and Tommi Jaakkola and Yuandong Tian and Bo Liu},
  year   = {2025},
  eprint = {2510.09541},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2510.09541v2}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2510.09541.svg)](https://codeofpaper.com/paper/2510.09541)

Discussion