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)は、複数のトークンを並列にデコードできる能力から、自己回帰モデルに代わる効率的な選択肢として注目を集めています。しかし、dLLMsを人間の嗜好やタスク固有の報酬に合わせて強化学習(RL)を用いて調整することは、その難解な対数尤度が標準的なポリシー勾配法の直接的な適用を妨げるため、困難を伴います。これまでの研究では、エビデンス下界(ELBO)のような代理指標が用いられてきましたが、これらの一方的な近似はポリシー勾配に大きなバイアスを導入する可能性があります。この問題に対処するため、我々は真の対数尤度の上界と下界の両方を活用する「サンドイッチポリシー勾配(SPG)」を提案します。実験結果から、SPGはELBOや1ステップ推定に基づくベースラインを大幅に上回ることが示されました。具体的には、SPGはdLLMs向けの最先端RL手法と比較して、GSM8Kで3.6%、MATH500で2.6%、Countdownで18.4%、Sudokuで27.0%の精度向上を達成しました。
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.