확산 언어 모델을 위한 강화 학습 안정화

초록

그룹 상대 정책 최적화(GRPO)는 학습 후 자율 회귀(AR) 언어 모델에 매우 효과적이지만, 이를 확산 대형 언어 모델(dLLM)에 직접 적용하면 보상 붕괴가 자주 발생합니다. 우리는 두 가지 호환성 문제의 원인을 규명했습니다. 첫째, GRPO는 시퀀스 확률로 정의되는 중요도 비율에 의존하는데, 이는 dLLM에서는 계산이 불가능하며 추정(예: ELBO 기반 또는 평균장 우도 프록시)해야 하여 본질적으로 잡음이 많은 비율을 생성합니다. 둘째, 표준 GRPO의 공식화는 추정된 비율을 위해 설계되지 않았습니다: 조건부 클리핑이 모델-불특정 추정 잡음에 의해 비정상적으로 우회되어 그래디언트 스파이크를 생성하는 반면, 고정된 그룹 크기 정규화는 높은 분산의 비율 추정에서 그래디언트 크기 변동을 증폭시킵니다. 우리는 이러한 효과들이 정책 표류를 유발하고 비율 분산을 더욱 증가시키는 자기 강화 불안정성 순환을 형성함을 보여줍니다. 이 순환을 깨기 위해 dLLM에 맞게 조정된 GRPO의 재구성인 StableDRL을 제안합니다. 이는 (i) 이상치로 인한 스파이크를 억제하기 위한 무조건 클리핑과 (ii) 개별 샘플 그래디언트의 볼록 껍질 내에서 업데이트를 제한하기 위한 자기 정규화를 사용합니다. 또한 계단식 어텐션 메커니즘을 통해 StableDRL을 블록 단위 확산 모델로 확장합니다.

English

Group Relative Policy Optimization (GRPO) is highly effective for post-training autoregressive (AR) language models, yet its direct application to diffusion large language models (dLLMs) often triggers reward collapse. We identify two sources of incompatibility. First, GRPO relies on importance ratios defined by sequence probabilities, which are intractable in dLLMs and must be estimated (e.g., via ELBO-based or mean-field likelihood proxies), yielding inherently noisy ratios. Second, standard GRPO's formulation is not designed for estimated ratios: its conditional clipping can be anomalously bypassed by model-agnostic estimation noise, producing gradient spikes, while its fixed group-size normalization amplifies gradient-magnitude fluctuations under high-variance ratio estimates. We show these effects form a self-reinforcing instability loop that drives policy drift and further increases ratio variance. To break this loop, we propose StableDRL, a reformulation of GRPO tailored for dLLMs that uses (i) unconditional clipping to suppress outlier-induced spikes and (ii) self-normalization to constrain updates within the convex hull of per-sample gradients. We further extend StableDRL to block-wise diffusion models via a staircase attention mechanism.

확산 언어 모델을 위한 강화 학습 안정화

Stabilizing Reinforcement Learning for Diffusion Language Models

초록

Support