SPG: Gradiente de Política Sanduíche para Modelos de Linguagem de Difusão Mascarada

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

Resumo

Modelos de linguagem de grande escala baseados em difusão (dLLMs) estão surgindo como uma alternativa eficiente aos modelos autoregressivos devido à sua capacidade de decodificar múltiplos tokens em paralelo. No entanto, alinhar dLLMs com preferências humanas ou recompensas específicas de tarefas por meio de aprendizado por reforço (RL) é desafiador, pois sua log-verossimilhança intratável impede a aplicação direta de métodos padrão de gradiente de política. Embora trabalhos anteriores utilizem substitutos como o limite inferior da evidência (ELBO), essas aproximações unilaterais podem introduzir um viés significativo no gradiente de política. Para resolver isso, propomos o Gradiente de Política Sanduíche (SPG), que aproveita tanto um limite superior quanto um limite inferior da log-verossimilhança verdadeira. Experimentos mostram que o SPG supera significativamente as abordagens baseadas em ELBO ou estimação de um passo. Especificamente, o SPG melhora a precisão em relação aos métodos de RL state-of-the-art para dLLMs em 3,6% no GSM8K, 2,6% no MATH500, 18,4% no Countdown e 27,0% no 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.