Predictive Divergence Masks for LLM RL

2026-07-12

AI Summary

This work addresses a critical inconsistency in existing reinforcement learning (RL) methods for large language models (LLMs), where the directionality criterion based on importance sampling ratios conflicts with the distributional divergence dynamics underlying trust-region constraints, often leading to unstable or inefficient training. To resolve this, the authors propose a predictive divergence masking mechanism that forecasts the direction of KL divergence change before policy gradient updates and dynamically masks token-level updates likely to increase divergence. This approach introduces, for the first time, divergence-change direction as a masking criterion—replacing conventional importance sampling ratios—and derives a closed-form solution for discrete softmax policies. A lightweight KL divergence estimator tailored for top-K vocabularies is also developed. Experiments demonstrate that the method significantly enhances RL training stability and performance across diverse model scales and precision settings, with predicted divergence directions closely aligning with empirical changes.

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.

Problem

• reinforcement learning
• large language models
• trust-region methods
• policy divergence
• importance ratio

Innovation

• Predictive Divergence Mask
• Trust Region
• Policy Divergence
• Top-K Estimation
• LLM Reinforcement Learning