Predictive Divergence Masks for LLM RL

Abstract: Reinforcement learning for 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.

• The paper introduces a closed-form predictive divergence mask that aligns training decisions with distributional trust regions, leading to enhanced stability in LLM RL.
• It derives a first-order expansion for KL divergence under policy-gradient updates and employs two lightweight tail estimators for top-K truncation.
• Empirical results show that the method lowers unsafe-keep rates and smooths reward curves, improving performance even in low-precision FP8 settings.

Predictive Divergence Masks: A Principled Direction Criterion for Divergence-Based Trust-Region LLM RL

Introduction

Trust-region methods are foundational in reinforcement learning for LLMs, where rollout data is typically generated by a behavior policy that diverges from the continually updated training policy. To maintain stability in this inherently off-policy regime, contemporary practice relies on masked or clipped objectives, such as Proximal Policy Optimization (PPO), often operationalized through sampled-token importance ratios. However, the adoption of distributional (divergence-based) trust-region masks—such as those in DPPO—has surfaced a methodological inconsistency: while the proximity to the trust region is measured by a policy divergence (e.g., KL), the decision to mask is still decided using the single-token ratio, thus breaking the alignment between trust-region definition and enforcement.

"Predictive Divergence Masks for LLM RL" (2607.10848) addresses this incongruity by introducing a predictive divergence mask in which the direction criterion is replaced by a closed-form estimate of the sign of the KL divergence change under the local policy-gradient update. Two lightweight estimators, tailored to the top-$K$ truncation typically exposed by inference engines, are presented: (1) an aggregated-tail and (2) a uniform-tail estimator. The mask is shown empirically to better align training decisions with the actual trend in divergence, resulting in improved training stability and effectiveness across various model scales and under stringent precision constraints.

Methodology and Theoretical Contributions

The central proposition of the paper is that the direction criterion in a masked objective should be defined by the distributional trust-region divergence, not a local importance ratio proxy. The authors formalize this by:

• Deriving an analytical, closed-form first-order expansion for the change in KL divergence under a policy-gradient update for softmax LLM policies.
• Decomposing this directional derivative into a local (sampled-token) and global (softmax normalization coupling) term. The ratio-based mask is shown to be equivalent to ignoring the global term, inherently disregarding the constraints imposed by the distributional trust region.
• Practical deployment with real LLM systems requires top-$K$ support: only the highest-probability log-probabilities and an aggregated tail are exposed. Two estimators are proposed for the unobserved tail's influence. Both are computationally negligible, and their difference is shown to be minor for $K \ll n$, where $n$ is the full vocabulary size.

Empirical Results

The method is evaluated on several Qwen3 models (Qwen3-4B-Base, Qwen3-8B-Base, Qwen3-30B-A3B-Base), and on DAPO-Math-17k train data with AIME24 and AIME25 as benchmarks. Top-$K$ KL masks with both the classic ratio-based and the proposed divergence-based direction criteria are compared, including in low-precision (FP8) settings to stress-test stability under severe training-inference mismatches.

Key qualitative and quantitative results include:

• Across all tested settings, the predictive divergence mask attains higher evaluation accuracy and improved RL stability compared to DPPO-TopK-KL with the ratio-based mask, with substantially less collapse in hard scenarios.

  Figure 1: Predictive divergence masks (both estimators) demonstrate superior stability and final accuracy compared to ratio-based DPPO and GRPO baselines at $\delta=0.15$.

• The improvement persists when the trust-region threshold is set more stringently ($\delta=0.05$), highlighting robustness to hyperparameter selection.

  Figure 2: Even under restrictive trust-region thresholds, predictive divergence masks maintain a performance margin over the baseline.

• Token-level analysis demonstrates that, for tokens where the ratio-based and divergence-based masks disagree, the divergence-based mask yields a lower unsafe-keep rate (fraction of updates that actually increase divergence) and retains more truly corrective updates.

  Figure 3: The divergence-based direction criterion results in a lower unsafe-keep rate among disagreement tokens compared to the ratio-based criterion.

  

  Figure 4: The divergence-based mask retains more corrective updates among the disagreement set, as indicated by a higher fraction of contracted divergences post-update.

• Training dynamics comparisons across all three model scales and both high and low-precision settings confirm improved PPO-KL regularity and a lower fraction of clipped tokens due to inappropriate mask decisions.

  Figure 5: Predictive mask yields smoother reward curves and lower clip fractions for Qwen3-4B-Base.

  

  Figure 6: For Qwen3-8B-Base, improvements mirror those for Qwen3-4B.

  

  Figure 7: Under FP8 E2E training, predictive mask maintains training regularity, whereas GRPO collapses.

  

  Figure 8: In the extreme FP8 rollout mismatch regime, predictive divergence mask secures stable training.

Implications and Outlook

The predictive divergence mask aligns the direction criterion with the same statistical divergence as the proximity criterion, yielding a self-consistent, distribution-aware trust-region mask. This closes the loophole whereby PPO/DPPO's token-level decisions could be misaligned with the real distributional shift, addressing a longstanding source of instability in practical off-policy LLM RL. The approach is computationally parsimonious, requiring only arithmetic on top-$K$ rollout statistics.

Practically, this resolves a critical mismatch in state-of-the-art LLM RL frameworks: by leveraging distributional information available from the top-$K$ statistics, RL fine-tuning is robustly stabilized against numerical and policy staleness-induced mismatches, with clear improvements in accuracy and reliability. Theoretical implications include a blueprint for extending other trust-region motivated objectives to similarly consistent, first-principles direction criteria—in particular, the extension to TV and other divergences is straightforward from the presented framework.

Given the mass deployment of off-policy RL fine-tuning in LLMs, the predictive divergence mask constitutes a strong argument for abandoning the sampled ratio as a universal proxy. Future directions include the exploration of higher-order corrections, the integration of token-level non-uniform trust-region constraints, and the extension to other generative modeling domains.

Conclusion

This work introduces and analyzes predictive divergence masks, a principled trust-region mask for LLM RL that leverages the closed-form first-order expansion of the policy divergence under gradient updates. This resolves the mismatch in current divergence-based RL methods between how the trust region is measured and how outward-pointing updates are defined. Empirical results on state-of-the-art LLMs confirm consistent training improvements across architectures, scales, and precision regimes. Distributionally aligned direction criteria as proposed herein are thus poised to become standard in high-stakes RL finetuning pipelines for LLMs.