SPG: Сэндвич-метод градиента политики для маскированных диффузионных языковых моделей

SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

October 10, 2025

Авторы: 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

Аннотация

Диффузионные большие языковые модели (dLLMs) становятся эффективной альтернативой авторегрессивным моделям благодаря их способности декодировать несколько токенов параллельно. Однако согласование dLLM с человеческими предпочтениями или задачами, специфичными для вознаграждений, с помощью обучения с подкреплением (RL), является сложной задачей, поскольку их невычислимая логарифмическая вероятность исключает прямое применение стандартных методов градиента политики. Хотя предыдущие работы используют суррогаты, такие как нижняя граница доказательства (ELBO), эти односторонние аппроксимации могут вносить значительную погрешность в градиент политики. Для решения этой проблемы мы предлагаем метод "Сэндвич-градиент политики" (SPG), который использует как верхнюю, так и нижнюю границу истинной логарифмической вероятности. Эксперименты показывают, что SPG значительно превосходит базовые методы, основанные на ELBO или одношаговой оценке. В частности, SPG улучшает точность по сравнению с современными методами RL для dLLM на 3,6% в GSM8K, 2,6% в MATH500, 18,4% в Countdown и 27,0% в 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.