SPG: 마스크 확산 언어 모델을 위한 샌드위치 정책 그래디언트

초록

확산 기반 대형 언어 모델(dLLMs)은 다중 토큰을 병렬로 디코딩할 수 있는 능력으로 인해 자기회귀 모델의 효율적인 대안으로 떠오르고 있습니다. 그러나 dLLMs를 강화 학습(RL)을 통해 인간의 선호도나 작업 특화 보상과 정렬시키는 것은, 그들의 다루기 어려운 로그 가능도가 표준 정책 경사법의 직접적인 적용을 방해하기 때문에 어려운 과제입니다. 기존 연구에서는 증거 하한(ELBO)과 같은 대체 방법을 사용하지만, 이러한 일방적인 근사는 상당한 정책 경사 편향을 초래할 수 있습니다. 이를 해결하기 위해, 우리는 실제 로그 가능도의 상한과 하한을 모두 활용하는 샌드위치 정책 경사법(SPG)을 제안합니다. 실험 결과, SPG는 ELBO나 일단계 추정을 기반으로 한 기준선을 크게 능가하는 것으로 나타났습니다. 구체적으로, 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.

SPG: 마스크 확산 언어 모델을 위한 샌드위치 정책 그래디언트

SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

초록

Support