SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models
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.
68edb2ecde1fee572713a753
Chenyu Wang
68edb2ecde1fee572713a754
Paria Rashidinejad
68edb2ecde1fee572713a755
DiJia Su
68edb2ecde1fee572713a756
Song Jiang
68edb2ecde1fee572713a757
Sid Wang
68edb2ecde1fee572713a758
Siyan Zhao
68edb2ecde1fee572713a759
Cai Zhou
68edb2ecde1fee572713a75a
Shannon Zejiang Shen
68edb2ecde1fee572713a75b
Feiyu Chen
68edb2ecde1fee572713a75c
Tommi Jaakkola
68edb2ecde1fee572713a75d
Yuandong Tian
68edb2ecde1fee572713a75e
Bo Liu
SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models
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.
The Sandwiched Policy Gradient method improves reinforcement learning for diffusion large language models by using both upper and lower bounds of log-likelihood, outperforming ELBO-based methods.
diffusion large language models dLLMs autoregressive models reinforcement learning RL log-likelihood policy gradient methods evidence lower bound ELBO Sandwiched Policy Gradient SPG GSM8K MATH500 Countdown Sudoku