SPG: Gesandwicht Beleidsgradiënt voor Gemaskeerde Diffusie Taalmodellen

SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

October 10, 2025

Auteurs: 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

cs.AI

Samenvatting

Diffusion large language models (dLLMs) doen zich gelden als een efficiënt alternatief voor autoregressieve modellen vanwege hun vermogen om meerdere tokens parallel te decoderen. Het afstemmen van dLLMs op menselijke voorkeuren of taakspecifieke beloningen via reinforcement learning (RL) is echter uitdagend, omdat hun onberekenbare log-waarschijnlijkheid de directe toepassing van standaard policy gradient-methoden verhindert. Hoewel eerder werk surrogaten zoals de evidence lower bound (ELBO) gebruikt, kunnen deze eenzijdige benaderingen aanzienlijke policy gradient-bias introduceren. Om dit aan te pakken, stellen we de Sandwiched Policy Gradient (SPG) voor, die zowel een boven- als een ondergrens van de werkelijke log-waarschijnlijkheid benut. Experimenten tonen aan dat SPG aanzienlijk beter presteert dan baseline-methoden gebaseerd op ELBO of éénstapschatting. Specifiek verbetert SPG de nauwkeurigheid ten opzichte van state-of-the-art RL-methoden voor dLLMs met 3,6% in GSM8K, 2,6% in MATH500, 18,4% in Countdown en 27,0% in Sudoku.

English

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.