LLM 강화학습을 위한 예측적 발산 마스크

초록

대규모 언어 모델(LLM)을 위한 강화학습은 일반적으로 오프-정책 업데이트를 안정화하기 위해 신뢰 영역 마스크(trust-region mask)에 의존한다. 지배적인 PPO 방식 접근법은 샘플링된 토큰 중요도 비율을 두 가지 기준에 사용한다: 정책이 행동 정책에서 너무 멀리 벗어났는지 확인하는 근접성 기준과, 업데이트가 정책을 더 멀리 밀어내는지 확인하는 방향성 기준이다. 최근 연구인 DPPO는 PPO의 비율 기반 테스트를 행동 정책과 학습 정책 간의 확률 발산으로 대체하여 근접성 기준을 개선했다. 그러나 그 방향성 기준은 여전히 PPO에서 계승되었다. 토큰은 샘플링된 토큰 중요도 비율이 1에서 멀어질 때만 마스킹될 수 있다. 우리는 이러한 비율 기반 방향성 기준이 단일 샘플 근사치(single-sample proxy)이며, 근접성 기준을 정의하는 발산의 변화와 부호가 일치하지 않을 수 있음을 관찰한다. 따라서 우리는 예측 발산 마스크(predictive divergence mask)를 제안하는데, 이는 다음 정책-경사 단계가 신뢰 영역에서 사용하는 동일한 발산을 증가시킬지 감소시킬지를 묻는다. LLM 강화학습에서 사용되는 이산 소프트맥스 정책의 경우, 이 예측을 폐쇄형(closed form)으로 유도한다. 프로덕션 롤아웃 엔진은 어휘의 절단된(top-K) 뷰만 노출하기 때문에, 이 예측을 위해 두 가지 경량 top-K 추정량을 개발한다. 상세 분석 결과, 발산 기반 방향성은 샘플링된 비율보다 발산의 실현된 변화와 더 잘 정렬되며, 결과적인 마스크는 모델 규모 및 정밀도 설정 전반에 걸쳐 강화학습 훈련을 개선한다.

English

Reinforcement learning for large language models (LLMs) typically relies on trust-region masks to stabilize off-policy updates. The dominant PPO-style approach uses the sampled-token importance ratio for two criteria: a proximity criterion, which asks whether the policy has moved too far from the behavior policy, and a direction criterion, which asks whether the update pushes it farther away. Recent work DPPO improves the proximity criterion by replacing PPO's ratio-based test with a probability divergence between the behavior and training policies. However, its direction criterion is still inherited from PPO. A token can be masked only when the sampled-token importance ratio moves away from one. We observe that this ratio-based direction criterion is a single-sample proxy that can disagree in sign with the change of the divergence that defines the proximity criterion. We therefore propose the predictive divergence mask, which asks whether the next policy-gradient step will increase or decrease the same divergence used by the trust region. For the discrete softmax policies used in LLM RL, we derive this prediction in closed form. Because production rollout engines expose only a truncated (top-K) view of the vocabulary, we develop two lightweight top-K estimators for this prediction. Detailed analysis shows the divergence-based direction is better aligned with the realized change of the divergence than the sampled ratio, and the resulting masks improve RL training across model scales and precision settings.