SPG: Gradiente di Politica a Sandwich per Modelli Linguistici di Diffusione con Mascheramento

SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

October 10, 2025

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

Abstract

I modelli linguistici di grandi dimensioni basati su diffusione (dLLM) stanno emergendo come un'alternativa efficiente ai modelli autoregressivi grazie alla loro capacità di decodificare più token in parallelo. Tuttavia, allineare i dLLM con le preferenze umane o con ricompense specifiche per il compito tramite apprendimento per rinforzo (RL) è impegnativo poiché la loro log-verosimiglianza intrattabile preclude l'applicazione diretta dei metodi standard del gradiente della politica. Sebbene lavori precedenti utilizzino approssimazioni unilaterali come il limite inferiore dell'evidenza (ELBO), queste approssimazioni possono introdurre un significativo bias nel gradiente della politica. Per affrontare questo problema, proponiamo il Gradiente della Politica a Sandwich (SPG) che sfrutta sia un limite superiore che un limite inferiore della vera log-verosimiglianza. Gli esperimenti dimostrano che SPG supera significativamente i metodi di riferimento basati su ELBO o stime a un passo. In particolare, SPG migliora l'accuratezza rispetto ai metodi RL all'avanguardia per i dLLM del 3,6% in GSM8K, del 2,6% in MATH500, del 18,4% in Countdown e del 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.