Predictive Divergence Masks for LLM RL

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.

한 줄 요약

PPO 기반 방향 기준의 한계를 보완한 predictive divergence mask를 제안하여 LLM RL의 훈련 안정성과 효과성을 개선한다.

핵심 기여도

PPO의 방향 기준을 divergence 기반으로 재설계하여, 실제 divergence 변화와 더 잘 일치시킴.
softmax 정책에 대해 1차 변화량을 닫힌 형태로 유도.
top-K 기반 rollout 엔진에 맞춘 두 가지 가벼운 추정기 제안 (aggregated-tail, uniform-tail).
61개 시드에서 평균적으로 2.7% 포인트 낮은 unsafe-keep 비율을 기록함.

핵심 아이디어

PPO는 trust region 안정화를 위해 단일 샘플 기반의 importance ratio를 사용하지만, 이는 전체 분포 변화를 반영하지 못한다. DPPO는 proximity 기준을 divergence로 대체했으나, 방향 기준은 여전히 PPO의 ratio 기반을 따르고 있다. 이는 두 기준 간 일관성을 해친다. 본 연구는 방향 기준을 divergence의 1차 변화량을 예측하는 방식으로 재설계하여, 실제 divergence 변화와의 일치도를 높였다. 이는 softmax 정책에서 닫힌 형태로 유도 가능하며, top-K 기반 rollout 엔진에 맞춰 가벼운 추정기를 제안함으로써 실용성을 확보했다.

기술적 접근법

**Proximity criterion**: KL divergence $ D_{\mathrm{KL}}^{\mathrm{TopK}}(\mu \|\pi) $를 사용.
**Direction criterion**: 샘플 기반 $ \operatorname{sign}(\hat{A}_t (r_t - 1)) $ 대신, divergence의 1차 변화량 $ \dot{D} $를 기반으로 결정.
**Predictive divergence mask**: $ \hat{A}_t \cdot \dot{D} \leq 0 $일 때 업데이트 유지.
**Top-K 추정기**: aggregated-tail, uniform-tail 방식으로 미관측 토큰 영역을 근사.
**실험 설정**: 1024개 시퀀스로 구성된 rollout, 61개 시드, 각 시드당 평균 1435개의 trust region 밖 토큰 분석.

주요 결과

**Unsafe-keep rate**: divergence 기준은 34.2% (PPO 기준 36.9%)로 2.7% 포인트 낮음.
**Corrective update rate**: divergence 기준은 65.1% (PPO 기준 62.9%)로 2.2% 포인트 높음.
**Disagreement set**: 시드당 약 82개 토큰에서 두 기준이 다른 결정을 내림.
**Top-K KL divergence**: 전체 KL의 99% 이상을 포착하며, 실용적 rollout 엔진과 호환됨.

의의 및 한계

**의의**: PPO의 단일 샘플 기반 방향 기준의 한계를 보완하고, 전체 분포 변화를 반영한 trust region 제어를 가능하게 함.
**실용성**: rollout 엔진에서 제공하는 top-K 정보만으로도 유의미한 성능 개선을 달성.
**한계**: disagreement set이 작아서 전체 훈련 효과는 크지 않으며, 모델 규모와 정밀도 설정에 따라 결과가 다를 수 있음.

실용적 활용

LLM의 off-policy RL 훈련에서 trust region 안정화가 필요한 상황, 특히 rollout 엔진이 top-K만 제공하는 실제 배포 환경에서 유용. 대규모 모델의 정밀도 저하를 방지하면서 훈련 효과성을 유지할 수 있음.