PaperlayerView exact revision on arXiv

Source-linked AI summary

Predictive Divergence Masks for LLM RL

Xiangxin Zhou, Jiarui Yao, Penghui Qi, Bowen Ping, Jiaqi Tang, Haonan Wang, Tianyu Pang

arXiv:2607.10848v1cs.LG

TL;DR

PPO-style direction masks can disagree with the divergence used to define the trust region. This paper predicts divergence change for the next gradient step and shows predictive masks improve training across model scales and precision settings.

  • Problem

    PPO-style direction criteria use a sampled-token ratio that can disagree with the divergence defining the trust region.

  • Method

    The predictive divergence mask uses a closed-form directional derivative for softmax policies, with lightweight top-K estimators for truncated vocabulary distributions.

  • Results

    Across all four model–precision settings, predictive divergence masks improved over DPPO-TopK-KL while divergence-based methods remained stable.

  • Takeaways & Limitations

    When the trust region is defined by a divergence, tracking that divergence provides a better-aligned direction criterion for masking updates.

  • Takeaways & Limitations

    The predicted sign is a local first-order estimate because other tokens’ updates can affect the realized post-update divergence.

Abstract

from arXiv · show

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.

1 Introduction

LLM reinforcement learning optimizes token-level policies on rollouts that may come from a different behavior policy, motivating asymmetric trust-region masks. The paper proposes predictive divergence masks, which determine masking from whether the next gradient step is predicted to increase the policy divergence.

  • Motivation: LLM RL uses autoregressive token-level policies with sequence-level rewards, while rollout data may be generated by a numerically different behavior policy.This creates the off-policy setting that trust-region methods aim to stabilize.
  • Trust-region masks: Trust-region masks combine proximity and direction criteria, masking only updates that would further increase deviation from the behavior policy.Updates that reduce the relevant deviation remain active.
  • Problem: The sampled importance ratio is a poor proxy for distributional shift because it observes only the sampled token and ignores probability mass elsewhere in the vocabulary.DPPO therefore replaces ratio-based proximity with a divergence between behavior and training policies, typically estimated from top-K log-probabilities.
  • Contribution: Predictive divergence masks classify updates by whether the next gradient step will increase the divergence, using a closed-form first-order prediction for softmax LLM policies.The prediction’s coefficient is a simple function of current token probabilities and splits into two terms.

2 Background

LLM RL optimizes autoregressive policies token by token using responses sampled from a behavior policy. PPO uses a noisy sampled-token trust-region mask, while DPPO replaces its proximity test with divergence but retains PPO’s ratio-based direction criterion.

  • LLM RL as token-level policy optimization: LLM RL generates responses token by token from a behavior policy and optimizes expected response reward.The behavior policy is typically a slightly stale copy of the current policy served by an optimized inference engine.
  • PPO and the ratio-based trust region: PPO masks updates using proximity |r_t − 1| > ϵ and an asymmetric direction criterion that restrains only outward-pointing updates.Its proximity test uses the sampled token, so |r_t − 1|/2 is a one-sample Monte Carlo estimate of total variation distance.
  • DPPO and the divergence-based trust region: DPPO replaces PPO’s sampled-ratio proximity test with the per-token divergence criterion D_t > δ while retaining the asymmetric direction criterion.The divergence choice used throughout is the forward KL, D_KL(µ ∥ π), which controls total variation through Pinsker’s inequality.
  • DPPO and the divergence-based trust region: Top-K divergence restricts the support to behavior-policy top-K tokens plus the sampled token and aggregates the remaining vocabulary into a residual tail bucket.It lower-bounds true KL and is described as extremely tight because top-K typically captures more than 99% of probability mass.
  • DPPO and the divergence-based trust region: DPPO-TopK-KL is the primary baseline because DPPO upgrades only proximity, leaving PPO’s sampled-ratio direction criterion unchanged.The next section tests whether this criterion tracks the sign of divergence change and derives a divergence-based replacement.

3 Methodology

The method replaces PPO’s ratio-based direction test with a first-order prediction of whether the next policy-gradient step increases the divergence defining the trust region. It derives this prediction for softmax policies and estimates it from rollout engines’ top-K probabilities and tail masses.

  • Divergence-based direction criterion: The predictive divergence mask blocks tokens outside the trust region when the directional derivative predicts that the next update will increase divergence, and keeps them when it predicts a decrease.The criterion is the sign of the divergence’s first-order change, replacing the ratio-based direction criterion.
  • Closed-form prediction: For softmax policies, the directional derivative follows the sampled-token logit direction 1[i = k] −πi, scaled by the advantage and importance ratio.Because the importance ratio is positive, the step direction’s sign is determined by sign(Â), and the resulting first-order divergence-change sign is sign(Â · Ḋ).
  • Closed-form prediction: The divergence derivative includes both the sampled-token local probability gap and a global normalization-coupling term that the single sampled ratio cannot observe.The ratio-based criterion is equivalent to retaining only the local term, whereas the divergence-based criterion retains both terms.
  • Top-K estimation: Two top-K estimators recover the missing tail contribution by either aggregating all unseen probability mass into one token or spreading it uniformly across unseen vocabulary entries.The aggregated-tail estimator may overestimate the tail contribution, while the uniform-tail contribution is suppressed by 1/(n −m) and is negligible for large vocabularies.
  • Top-K estimation: Both estimators use only rollout-provided top-K probabilities and tail masses, requiring no additional forward or backward pass, finite-difference evaluation, or new hyperparameter.When retained top-K tokens capture most probability mass, the aggregated-tail and uniform-tail masks should behave nearly identically.
  • Prediction scope: The predictive sign is a token-level first-order local estimate, so the realized divergence change may also reflect gradient contributions from other tokens in the batch.The masks differ when the distribution-wide correction reverses the sampled-token ratio’s predicted direction; otherwise, their local signs agree.

4 Experiments and Results

Experiments across four Qwen3 model–precision settings show predictive divergence masks improve stability and performance over DPPO-TopK-KL, while better aligning token-level decisions with realized divergence changes. This advantage persists under a tighter trust-region threshold, though all methods perform worse.

  • Models, Data, and Benchmarks: Experiments use about 13k filtered DAPO-Math-17k training samples, AIME24/AIME25 avg@16 evaluation, and three Qwen3 base models across four model–precision settings.Training runs use 32 NVIDIA H20 GPUs; Qwen3-30B-A3B-Base includes FP8 rollout and FP8 end-to-end settings.
  • Main results: At δ = 0.15, GRPO clip-higher is unstable and collapses for both Qwen3-30B-A3B-Base settings, whereas divergence-based methods remain stable and predictive masks outperform DPPO-TopK-KL.The comparison isolates the direction criterion because DPPO-TopK-KL shares the same top-K KL proximity criterion.
  • Threshold sensitivity: At δ = 0.05, all methods perform worse, but both predictive divergence masks still outperform DPPO-TopK-KL, indicating greater robustness to the trust-region threshold.Aggregated-tail and uniform-tail masks behave similarly, consistent with a small tail correction and limited sensitivity to the tail model.
  • Token-level direction analysis: The two direction criteria disagree on only about 82 tokens per seed, so token-level differences are concentrated in a small disagreement set.Each seed averages 1435 tokens outside the trust region, and the token-level analysis uses 61 independent seeds.
  • Token-level direction analysis: 34.2% vs. 36.9%: the divergence-based criterion reduces unsafe keeps by 2.7 percentage points and has the lower rate on 38 of 61 seeds.Unsafe keeps are disagreement tokens retained despite a subsequent divergence increase, ΔD > 0.
  • Token-level direction analysis: 65.1% vs. 62.9%: divergence-based direction keeps a higher fraction of updates that contract the divergence, supporting better alignment with realized divergence changes.This pooled corrective statistic is complementary to the per-seed unsafe-keep comparison.

5 Closing Remarks · Appendix A. Related Work

The paper proposes a predictive divergence mask that replaces PPO’s sampled-ratio direction test with a first-order prediction of divergence change. Related work situates this approach within trust-region optimization, LLM reasoning RL, and divergence-based policy-drift masks.

  • 5 Closing Remarks: Divergence-based masks such as DPPO improve PPO’s proximity criterion but retain its sampled-ratio direction criterion.The paper identifies this inherited direction test as the mismatch it addresses.
  • 5 Closing Remarks: The divergence-based direction criterion keeps a larger fraction of contracting updates than the comparison criterion in Figure 4.Figure 4 measures ΔD = Dpost − Dpre after one actual update for kept disagreement tokens across 61 seeds at δ = 0.15.
  • 5 Closing Remarks: The predictive divergence mask decides whether to mask an update using the predicted first-order change of the divergence itself.For softmax policies, the prediction decomposes into a local ratio-based term and a global term capturing distribution-wide coupling.
  • 5 Closing Remarks: The divergence-based direction criterion is better aligned with realized divergence change and improves training stability and effectiveness across model scales and precision settings.The method remains a local first-order approximation because actual updates aggregate gradients across tokens and may introduce token interactions.
  • Appendix A. Related Work: TRPO and CPO enforce explicit KL constraints, whereas PPO approximates the trust-region constraint with a cheap clipped surrogate.Mirror-descent and proximal formulations make the constraint explicit again, while SPO replaces the hard...
  • Appendix A. Related Work: PPO-style optimization underpins RLHF and modern reasoning training, while GRPO removes the critic using group-relative advantages.Recent follow-ups refine advantage estimation, clipping ranges, or importance-sampling correction; DAPO decouples th...
  • Appendix A. Related Work: DPPO replaces sampled ratio deviation with an explicit distributional divergence mask, including the top-K KL realization used by this work.Complementary research studies non-uniform token-level trust regions that vary constraints across tokens.

Appendix B. Derivation of the Directional Derivative · B.1 Softmax sensitivity · B.2 Directional derivative of the divergence

Appendix B derives Proposition 1 for generic softmax supports, including the reduced support used in the top-K construction with an aggregate tail bucket. It connects softmax sensitivity to the directional derivative of the divergence along the policy-gradient-induced logit direction.

  • Appendix B. Derivation of the Directional Derivative: The derivation applies to both the full vocabulary and reduced top-K supports when the tail is represented as an aggregate bucket.This establishes the scope of Proposition 1's generic softmax-support formulation.
  • B.1 Softmax sensitivity: B.1 introduces first-order softmax sensitivity under a fixed logit perturbation direction v.The logits are perturbed as z_i(η) = z_i + ηv_i.
  • B.2 Directional derivative of the divergence: The proof uses the unit logit direction induced by the policy-gradient step for the per-token surrogate r ˆA.Its logit gradient is given as ˆA r ∇z log π_k.
  • B.2 Directional derivative of the divergence: Because r is positive, the sign of the realized step is determined by the advantage ˆA.The derivative is therefore computed along the corresponding unit direction.
  • B.2 Directional derivative of the divergence: The divergence derivative is developed by taking averages under the current policy and rewriting the forward KL.The derivation explicitly introduces both the policy average and the forward-KL representation.
  • B.2 Directional derivative of the divergence: Applying the softmax-sensitivity result to the η-dependent KL term produces the intermediate derivative expression.The first KL term is independent of η, while the second is handled using Lemma 4.
  • B.2 Directional derivative of the divergence: The two terms in Eq. (24) are identified separately before being substituted back into the equation.This decomposition supplies the components needed for the final directional-derivative expression.
  • B.2 Directional derivative of the divergence: Substitution yields Eq. (7), completing the proof of Proposition 1.The appendix closes after establishing the stated directional-derivative result.

Appendix C. Equivalence under Binary-KL Approximation

Under the binary-KL approximation, the aggregated-tail predictive direction reduces exactly to the ratio-based direction criterion. Consequently, the predictive mask is identical to DPPO-Binary-KL, unlike top-K KL settings that retain additional tokens.

  • Binary-KL equivalence: The aggregated-tail predictive direction reduces exactly to the ratio-based direction criterion when the vocabulary is collapsed into the sampled token and its complement.This binary support contains the sampled token k and a tail bucket.
  • Binary-KL equivalence: Under the same binary KL proximity criterion, the predictive mask with the aggregated-tail estimator is identical to DPPO-Binary-KL.The equivalence follows from applying the predictive-direction expression to the two-atom support under nondegenerate softmax probabilities.
  • Limitation: The binary approximation removes the distributional degrees of freedom used by the divergence-based direction criterion, whereas top-K KL retains additional tokens.Therefore, the top-K setting can produce a global correction that differs from the sampled-token ratio direction.

Appendix D. Predictive Divergence Masks under TV Divergence

This appendix extends predictive divergence masking from KL to total variation (TV) divergence. It derives the TV direction criterion and gives aggregated-tail and uniform-tail top-K estimators while retaining the same masking rule.

  • TV divergence extension: TV uses the same first-order principle as KL to derive a divergence-based direction criterion for a generic support.The paper focuses primarily on top-K forward KL because it matches DPPO-TopK-KL; TV is treated as an extension.
  • TV divergence extension: The TV directional derivative is defined with subgradient convention sign(0) = 0 when a probability gap is zero.Relative to KL, the additional factor s_i = sign(µ_i −π_i) arises from differentiating TV’s absolute value.
  • Top-K estimators: The aggregated-tail estimator represents retained tokens K plus one tail bucket for both TV proximity and the softmax coupling term.The unseen tail is treated as a single coordinate under this construction.
  • Top-K estimators: The uniform-tail estimator spreads residual mass over the n −m unseen tokens while leaving observed head terms unchanged.For j /∈K, it approximates π_j and µ_j by distributing π_tail and µ_tail across n −m coordinates.
  • Masking rule: With either estimator, the masking rule is unchanged, but the proximity measure and direction coefficient are replaced by their TV counterparts.The resulting predictive TV mask corresponds to using TV as the proximity criterion.

Appendix E. Additional Experimental Details · E.1 Training Details

Training uses VeRL with Megatron-LM and vLLM, typically in BF16, while also evaluating FP8 rollout-only and end-to-end settings for Qwen3-30B-A3B-Base. Experiments control shared optimization and rollout settings, use one PPO epoch per rollout batch, and set method-specific top-K and divergence thresholds.

  • E.1 Training Details: The training pipeline uses VeRL, with Megatron-LM for training and vLLM for rollout generation.These components define the implementation stack used in the experiments.
  • E.1 Training Details: Both training and rollout use BF16 by default.The default precision is shared across the training and rollout components.
  • E.1 Training Details: For Qwen3-30B-A3B-Base, the study evaluates FP8 Rollout and FP8 E2E precision settings.FP8 Rollout applies FP8 only in vLLM, whereas FP8 E2E applies it in both Megatron-LM and vLLM.
  • E.1 Training Details: All methods use the same learning rate, batching, rollout temperature, response budget, and group size.These shared settings keep the main training and rollout controls consistent across methods.
  • E.1 Training Details: Each rollout batch is used for one PPO epoch.The training schedule therefore performs a single PPO pass per collected rollout batch.
  • E.1 Training Details: DPPO-TopK-KL and predictive runs use a top-K support of K = 20.The top-K support is the shared truncation setting for these runs.
  • E.1 Training Details: The main comparison uses δ = 0.15, while the tight-threshold comparison uses δ = 0.05.These are the two divergence-threshold settings specified for the experiments.
  • E.1 Training Details: Table 1 summarizes the training hyperparameters, with both Qwen3-30B-A3B-Base precision settings using the final column.The table provides the experiment-level hyperparameter summary and assigns both precision variants to its final column.

E.2 Detailed training dynamics

Figures 5–8 provide detailed training dynamics across model scales and precision settings, tracking reward, task accuracy, response length, PPO-KL, and clip fraction. The predictive divergence masks occupy a similar PPO-KL range to DPPO-TopK-KL while generally producing lower clip fractions.

  • Training dynamics: Figures 5–8 track training reward, separate AIME24 and AIME25 accuracy, response length, PPO-KL, and clip fraction.These metrics are reported for Qwen3-4B-Base, Qwen3-8B-Base, and Qwen3-30B-A3B-Base under the specified training settings.
  • Training dynamics: DPPO-TopK-KL and the two predictive divergence masks occupy a similar PPO-KL range, while the predictive divergence masks generally exhibit lower clip fractions.The comparison is made across the detailed training dynamics shown in Figures 5–8.
  • Training dynamics: Figures 5 and 6 show Qwen3-4B-Base and Qwen3-8B-Base training dynamics at K = 20 and δ = 0.15.AIME24 and AIME25 accuracy are shown separately, alongside reward, response length, PPO-KL, and clip fraction.
Loading 2607.10848v1…