SPG : Gradient de Politique en Sandwich pour les Modèles de Langage à Diffusion Masquée

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

Les modèles de langage à grande échelle par diffusion (dLLM) émergent comme une alternative efficace aux modèles autorégressifs grâce à leur capacité à décoder plusieurs tokens en parallèle. Cependant, aligner les dLLM avec les préférences humaines ou les récompenses spécifiques à une tâche via l'apprentissage par renforcement (RL) est difficile car leur log-vraisemblance intraçable empêche l'application directe des méthodes classiques de gradient de politique. Bien que les travaux antérieurs utilisent des substituts comme la borne inférieure de l'évidence (ELBO), ces approximations unilatérales peuvent introduire un biais significatif dans le gradient de politique. Pour résoudre ce problème, nous proposons le Gradient de Politique Sandwich (SPG) qui exploite à la fois une borne supérieure et une borne inférieure de la vraie log-vraisemblance. Les expériences montrent que le SPG surpasse significativement les méthodes de référence basées sur l'ELBO ou l'estimation en une étape. Plus précisément, le SPG améliore la précision par rapport aux méthodes RL de pointe pour les dLLM de 3,6 % dans GSM8K, 2,6 % dans MATH500, 18,4 % dans Countdown et 27,0 % dans 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.