확산 언어 모델을 위한 효율적이고 안정적인 강화학습

Efficient and Stable Reinforcement Learning for Diffusion Language Models

강화학습(RL)은 확산 기반 대형 언어 모델(dLLMs)의 복잡한 추론 능력을 이끌어내는 데 필수적이다. 그러나 dLLM에 RL을 적용하는 것은 효율성과 안정성 측면에서 고유한 과제에 직면한다. 이러한 문제를 해결하기 위해, 우리는 dLLM을 위한 RL의 효율성과 안정성을 동시에 향상시키도록 설계된 프레임워크인 시공간 가지치기(STP)를 제안한다. STP는 (1) 정적 사전 정보(static priors)를 사용하여 탐색 공간을 제한하는 '공간적 가지치기'와 (2) 불필요한 후반부 정제 단계를 우회하는 '시간적 가지치기'를 통해 생성 과정의 중복성을 압축한다. 이론적 분석에 따르면 STP는 로그 우도 추정의 분산을 엄격히 감소시켜 더 안정적인 정책 업데이트를 보장한다. 광범위한 실험 결과, STP는 효율성과 정확도 모두에서 최신 베이스라인 모델들을 능가하는 것으로 나타났다. 코드는 https://github.com/Lolo1222/STP 에서 확인할 수 있다.

Reinforcement Learning (RL) is crucial for unlocking the complex reasoning capabilities of Diffusion-based Large Language Models (dLLMs). However, applying RL to dLLMs faces unique challenges in efficiency and stability. To address these challenges, we propose Spatio-Temporal Pruning (STP), a framework designed to simultaneously improve the efficiency and stability of RL for dLLMs. STP compresses the redundancy in the generative process through: (1) \textit{spatial pruning}, which constrains the exploration space using static priors; and (2) \textit{temporal pruning}, which bypasses redundant late-stage refinement steps. Our theoretical analysis demonstrates that STP strictly reduces the variance of the log-likelihood estimation, thereby ensuring more stable policy updates. Extensive experiments demonstrate that STP surpasses state-of-the-art baselines in both efficiency and accuracy. Our code is available at https://github.com/Lolo1222/STP.