SPG: Gradiente de Política Intercalada para Modelos de Lenguaje de Difusión Enmascarada

SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

October 10, 2025

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

Resumen

Los modelos de lenguaje de gran escala basados en difusión (dLLMs) están surgiendo como una alternativa eficiente a los modelos autorregresivos debido a su capacidad para decodificar múltiples tokens en paralelo. Sin embargo, alinear los dLLMs con las preferencias humanas o recompensas específicas de tareas mediante aprendizaje por refuerzo (RL) es un desafío, ya que su log-verosimilitud intratable impide la aplicación directa de los métodos estándar de gradiente de políticas. Si bien trabajos previos utilizan sustitutos como la cota inferior de la evidencia (ELBO), estas aproximaciones unilaterales pueden introducir un sesgo significativo en el gradiente de políticas. Para abordar esto, proponemos el Gradiente de Política Emparedado (SPG), que aprovecha tanto una cota superior como una cota inferior de la log-verosimilitud verdadera. Los experimentos muestran que SPG supera significativamente a los métodos de referencia basados en ELBO o estimación de un solo paso. En concreto, SPG mejora la precisión sobre los métodos de RL más avanzados para dLLMs en un 3.6% en GSM8K, 2.6% en MATH500, 18.4% en Countdown y 27.0% en 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.