Search papers, labs, and topics across Lattice.
This paper introduces predictive divergence masks to enhance reinforcement learning stability in large language models (LLMs) by refining the direction criterion of the PPO-style approach. By replacing the traditional ratio-based direction criterion with a divergence-based prediction, the authors demonstrate that this method better aligns with the actual changes in policy divergence. The results show significant improvements in RL training performance across various model scales and precision settings, indicating a more effective approach to managing off-policy updates.
Predictive divergence masks can significantly enhance RL training stability in LLMs by aligning direction criteria with actual divergence changes.
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.
