Predictive Divergence Masks for LLM RL
2026-07-12 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study a way to make reinforcement learning safer and more stable when training large language models. They point out that the usual method, PPO, uses a ratio to decide if a training step is good but this can sometimes give misleading signals. To fix this, the authors propose a new method called the predictive divergence mask, which better predicts if the next training step will truly improve the model by looking at changes in probability divergence. They show how to calculate this prediction efficiently even when only partial information about possible words is available, and their method leads to better training results.
Reinforcement LearningLarge Language ModelsPPO (Proximal Policy Optimization)Importance RatioProbability DivergenceTrust RegionPolicy GradientSoftmax PolicyTop-K SamplingOff-policy Updates
Authors
Xiangxin Zhou, Jiarui Yao, Penghui Qi, Bowen Ping, Jiaqi Tang, Haonan Wang, Tianyu Pang
Abstract
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.