SPG: Sandwiched Policy Gradient für maskierte Diffusions-Sprachmodelle

SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

October 10, 2025

papers.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

cs.AI

papers.abstract

Diffusion Large Language Models (dLLMs) etablieren sich als effiziente Alternative zu autoregressiven Modellen, da sie in der Lage sind, mehrere Tokens parallel zu dekodieren. Die Ausrichtung von dLLMs an menschlichen Präferenzen oder aufgabenbezogenen Belohnungen durch Reinforcement Learning (RL) gestaltet sich jedoch schwierig, da ihre nicht handhabbare Log-Likelihood die direkte Anwendung standardmäßiger Policy-Gradient-Methoden verhindert. Während frühere Arbeiten Surrogate wie die Evidence Lower Bound (ELBO) verwenden, können diese einseitigen Approximationen erhebliche Verzerrungen im Policy-Gradient verursachen. Um dies zu adressieren, schlagen wir den Sandwiched Policy Gradient (SPG) vor, der sowohl eine obere als auch eine untere Schranke der wahren Log-Likelihood nutzt. Experimente zeigen, dass SPG Baselines, die auf ELBO oder Ein-Schritt-Schätzung basieren, deutlich übertrifft. Insbesondere verbessert SPG die Genauigkeit gegenüber modernsten RL-Methoden für dLLMs um 3,6 % in GSM8K, 2,6 % in MATH500, 18,4 % in Countdown und 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.