Predictive Divergence Masks for LLM RL
Xiangxin Zhou, Jiarui Yao, Penghui Qi, Bowen Ping, Jiaqi Tang, Haonan Wang, Tianyu Pang
Predictive divergence masks align RL update directions with distributional shifts, improving LLM training stability.
How can we improve the stability of LLM reinforcement learning by replacing ratio-based trust-region masks with divergence-based predictive masks?
Reinforcement learning for LLMs relies on trust-region masks to prevent unstable updates, but current methods use a single-sample importance ratio to decide both the proximity to the behavior policy and the direction of the update. This ratio is a poor proxy for distributional shift, causing the direction criterion to conflict with the divergence-based proximity criterion. The predictive divergence mask replaces this ratio-based check with the directional derivative of the divergence itself. This approach uses a closed-form coefficient that accounts for both the sampled token and the global softmax normalization coupling across the vocabulary. Across model scales and precision settings, this method improves training stability and effectiveness compared to standard divergence-based baselines.
Paper Primer
Standard trust-region methods like PPO use the importance ratio $r_t$ to decide if an update is "outward-pointing." While recent methods like DPPO upgrade the proximity criterion to a full distributional divergence, they still rely on the single-sample ratio to decide the update direction, creating a logical inconsistency between the two criteria.
The predictive divergence mask is a drop-in replacement: it calculates the first-order change of the divergence along the gradient step and masks tokens where this change is positive. The core move is a closed-form decomposition of the divergence's directional derivative: it combines a local term (the sampled token gap) with a global term (the softmax normalization coupling) that captures how the entire distribution shifts.
The predictive divergence mask reduces the "unsafe-keep" rate of updates compared to ratio-based direction criteria.
Token-level analysis across 61 independent seeds shows a 2.7 percentage-point reduction in the fraction of disagreement tokens that push the policy further outside the trust region.
Because production rollout engines only expose top-K probabilities, the authors provide two lightweight estimators—aggregated-tail and uniform-tail—to approximate the global coupling term. Both estimators require no additional forward or backward passes, adding negligible computational overhead.
Why does the ratio-based direction criterion fail when paired with a divergence-based proximity criterion?
The ratio-based criterion only observes the sampled token, whereas the divergence-based proximity criterion measures the shift of the entire probability distribution. The two can disagree because the sampled token's movement does not account for how the softmax normalizer shifts the remaining probability mass across the vocabulary.
Is this method sensitive to the specific way the unseen tail of the vocabulary is modeled?
No. The aggregated-tail and uniform-tail estimators perform nearly identically in practice because the top-K tokens capture over 99% of the probability mass, making the tail's contribution to the directional derivative negligible.
This paper demonstrates that when a trust region is defined by a distributional divergence, the direction criterion must be derived from that same divergence to maintain mathematical consistency.
Researchers using divergence-based trust regions (like DPPO) should replace ratio-based direction checks with this predictive derivative to improve training stability without additional computational cost.
Introduction and Motivation
We expose the inconsistency of ratio‑based direction masks and motivate a divergence‑based alternative.
LLM reinforcement learning relies on trust‑region masks that prune token updates. The dominant PPO‑style mask decides both whether a token lies outside the trust region (proximity) and whether its update would push the policy farther away (direction) using a sampled‑token importance ratio. This ratio‑based direction criterion is a single‑sample proxy that can disagree with the actual change in the divergence that defines the proximity region, creating an inconsistency that harms training.
The mask’s direction test looks only at the sampled token, ignoring how the whole softmax distribution shifts, so it can mistakenly block updates that would actually reduce the policy‑behavior divergence.
The ratio‑based direction mask is a brittle proxy that can mask the wrong tokens, motivating a divergence‑based direction criterion.
LLM RL and Trust-Region Baselines
Background clarifies token‑level RL, PPO masking, and the DPPO baseline.
Token‑level reinforcement learning treats a language model as a sequential policy $\pi$ that emits tokens conditioned on the growing context sₜ. The standard RL pipeline samples responses with a stale rollout policy $\mu$ and uses per‑token importance ratios to correct the policy gradient.
The mask keeps updates that move the policy toward the behavior policy while discarding outward‑pointing updates that would push the policy beyond a trusted region.
DPPO replaces the noisy single‑token proximity test with a full‑distribution divergence, keeping the same asymmetric direction test as PPO.
The Predictive Divergence Mask
We replace the ratio‑based direction mask with a divergence‑based Predictive Divergence Mask computed via a directional derivative.
The ratio‑based direction criterion looks only at the sampled token’s importance ratio, so under a divergence‑based trust region it can mis‑signal whether an update pushes the policy outside the region.
Instead of asking “does this single sampled ratio increase?” we ask “does the KL‑divergence between the target policy $\mu$ and the current policy $\pi$ increase if we take the gradient step?” – the mask blocks a token update whenever the first‑order change of the divergence is positive.
How does the Predictive Divergence Mask differ from the traditional ratio‑based direction criterion?
The ratio‑based mask uses only the local term $\text{sign}(\pi_k-\mu_k)$, ignoring how changing $\pi_k$ reshapes the entire distribution. The Predictive Divergence Mask adds the global coupling $\sum_i \pi_i(\mu_i-\pi_i)$, so it reacts to the full KL‑divergence change rather than a single‑sample proxy.
Local term: $\pi_k-\mu_k = 0.4 - 0.35 = 0.05$.
Aggregated‑tail term: $\pi_{\text{tail}}(\mu_{\text{tail}}-\pi_{\text{tail}})=0.3\,(0.4-0.3)=0.03$.
Aggregated‑tail $\dot{D}_{\text{agg}} = 0.05 + (\pi_2(\mu_2-\pi_2)) + 0.03 = 0.05 + 0.3\,(0.25-0.3) + 0.03 = 0.05 - 0.015 + 0.03 = 0.065$.
Uniform‑tail term: $\frac{0.3\,(0.4-0.3)}{n-m}= \frac{0.03}{5-2}=0.01$.
Uniform‑tail $\dot{D}_{\text{uni}} = 0.05 - 0.015 + 0.01 = 0.045$.
Both $\dot{D}$ values are positive → the Predictive Divergence Mask blocks the update.
The global tail contribution is tiny (0.03 vs 0.01) because the vocabulary is large relative to $K$, confirming that the two estimators yield the same blocking decision in practice.
Obtain the top‑K probabilities $\{\pi_i\}_{i\in K}$ and the tail mass $\pi_{\text{tail}}$ from the current policy.
Retrieve the corresponding target probabilities $\{\mu_i\}_{i\in K}$ and $\mu_{\text{tail}}$ (e.g., from a reference rollout).
Identify the sampled token $k$ and compute the local gap $(\pi_k-\mu_k)$.
Compute the global coupling term using the chosen estimator:
Form $\dot{D}= (\pi_k-\mu_k) +$ (global term).
Set the mask: block the token if $\dot{D}>0$, otherwise keep it.
Directional Derivative Criterion
Introduces the predictive divergence mask built from directional derivatives.
Ratio‑based token masks rely on noisy importance‑ratio estimates, which can mislead policy updates. The predictive divergence mask replaces that proxy with a first‑order estimate of how each token’s divergence will change.
It predicts the sign of the KL‑divergence change for a token if the policy were nudged infinitesimally along that token’s direction.
How does this directional‑derivative signal differ from the ordinary gradient used in PPO?
The ordinary gradient aggregates over the whole policy distribution, yielding a single update direction. The directional derivative isolates the change along the specific sampled‑token direction, providing a token‑level sign that can be masked independently.
Token A: $\dot{D}_t^{(A)} \le 0$ ⇒ $M^{\text{pred}}_t = 0$.
Token B: $\dot{D}_t^{(B)} > 0$ and $D_t > \delta$ ⇒ $M^{\text{pred}}_t = 1$.
Token C: $\dot{D}_t^{(C)} > 0$ and $D_t > \delta$ ⇒ $M^{\text{pred}}_t = 1$.
The mask activates only for tokens whose divergence is already significant and whose first‑order prediction is positive, focusing learning on promising directions.
Remark 3: the mask is a token‑level first‑order approximation. The actual parameter update aggregates gradients from all tokens in the batch, so the realized divergence shift at one token can also be affected by other tokens’ updates.
Empirical Evaluation
Predictive divergence masks boost stability and accuracy across model scales, outperforming baselines.
Standard LLM RL uses noisy importance‑ratio proxies for trust‑region masks. We replace them with predictive divergence masks derived from directional derivatives.
Predictive divergence masks keep training stable and improve average accuracy across all four model‑precision settings.
GRPO clip‑higher collapses in the two 30B‑A3B settings while the divergence‑based methods remain stable and achieve higher accuracy.
**Figure 1.** Evaluation accuracy (avg@16, averaged over AIME24 and AIME25) over training at $\delta = 0.15$ on the four settings: Qwen3-4B-Base, Qwen3-8B-Base, Qwen3-30B-A3B-Base with FP8 E2E training, and Qwen3-30B-A3B-Base with FP8 Rollout. The two predictive divergence masks use the aggregated-tail and uniform-tail estimators for the top-$K$ KL directional-derivative coefficient.
**Figure 2.** Evaluation accuracy (avg@16, averaged over AIME24 and AIME25) across the four settings with divergence threshold $\delta = 0.05$.
**Figure 3.** **Unsafe-keep rate on disagreement tokens** ($\delta = 0.15$, 61 seeds). We consider tokens outside the trust region ($D > \delta$) on which the ratio-based and divergence-based direction criteria make opposite keep/clip decisions. For each criterion, among the disagreement tokens it keeps, we report the fraction whose divergence increases after the update ($\Delta D > 0$). Lower is better. (a) Mean unsafe-keep rate across seeds. The divergence-based direction criterion is lower than the ratio-based criterion (34.2% vs. 36.9%). Error bars show standard errors. (b) Per-seed paired comparison, with each line corresponding to one seed and colored by which criterion has the lower unsafe-keep rate.
**Figure 4.** Realized divergence change of kept disagreement tokens ($\delta = 0.15$, 61 seeds). For each direction criterion, we collect the disagreement tokens it keeps and measure $\Delta D = D_{post} - D_{pre}$ after one actual update. Tokens from all seeds are combined in this figure. (a) Empirical CDF of $\Delta D$. The shaded region $\Delta D < 0$ indicates kept updates that contracted the divergence. (b) Fraction of kept updates with $\Delta D < 0$. The divergence-based direction criterion keeps a larger fraction of contracting updates than the ratio-based criterion (65.1% vs. 62.9%).
**Figure 5.** Training dynamics for Qwen3-4B-Base at $K = 20$ and $\delta = 0.15$. AIME24 and AIME25 are shown separately. The remaining panels report reward, response length, PPO-KL, and clip fraction.
**Figure 6.** Training dynamics for Qwen3-8B-Base at $K = 20$ and $\delta = 0.15$, with the same metrics and methods as Figure 5.
**Figure 7.** Training dynamics for Qwen3-30B-A3B-Base under FP8 E2E training at $K = 20$ and $\delta = 0.15$.
**Figure 8.** Training dynamics for Qwen3-30B-A3B-Base with FP8 Rollout at $K = 20$ and $\delta = 0.15$.
Softmax Sensitivity Analysis
We ablate the clipping rule into separate proximity and direction criteria and derive softmax sensitivity.
Separating the PPO clipping rule into a proximity mask and a direction mask lets us test whether each component is essential: the proximity mask enforces a trust‑region bound, while the direction mask decides the sign of the update using the sampled advantage.
The clipping rule can be split: a proximity criterion that limits how far the new policy may drift, and a direction criterion that determines whether the update pushes the policy outward or inward.
How does the direction criterion here differ from PPO’s standard sampled‑ratio direction?
Standard PPO uses the sign of the advantage times the sampled importance ratio (r − 1) to decide update direction. Our criterion replaces the ratio with the sign of the directional derivative of the divergence, aligning the update direction with the same metric that defines the trust‑region mask.
Directional Derivative Derivation
Derives the directional derivative of the forward KL divergence used for the predictive divergence mask.
The directional derivative quantifies how a scalar perturbation $\eta$ changes a quantity of interest—in this case the log‑policy and the forward KL divergence.
Binary KL Equivalence
Binary KL reduces the predictive direction to the DPPO ratio criterion, revealing a degenerate case.
DPPO also considers a binary KL approximation that collapses the vocabulary into the sampled token and its complement. In this binary support the two policies are defined as $\mu_{\text{bin}} = (\mu_k,\,1-\mu_k)$ and $\pi_{\text{bin}} = (\pi_k,\,1-\pi_k)$.
The binary KL proximity criterion is $D_{KL}(\mu_{\text{bin}}\|\pi_{\text{bin}})=\mu_k \log\frac{\mu_k}{\pi_k}+(1-\mu_k)\log\frac{1-\mu_k}{1-\pi_k}$. Applying Eq. (7) yields $\dot{D}_{\text{bin}} = 2(1-\pi_k)(\pi_k-\mu_k)$, whose sign equals $\operatorname{sign}(\pi_k-\mu_k)=\operatorname{sign}(r-1)$ with $r=\pi_k/\mu_k$.
Thus the predictive mask with the aggregated‑tail estimator matches the DPPO‑Binary‑KL direction, a degenerate case where the binary approximation removes distributional degrees of freedom. By contrast, the top‑K KL setting retains additional tokens, so the global correction can differ from the sampled‑token ratio direction.
TV Divergence Masks
Derives TV‑based predictive divergence masks and their directional‑derivative formulas for top‑K estimators.
This appendix applies the first‑order derivation used for KL to the total‑variation (TV) divergence, yielding the TV‑based predictive divergence mask.
TV is nondifferentiable when a probability gap is exactly zero; we adopt the standard subgradient rule $\operatorname{sign}(0)=0$.
Here $s_i=\operatorname{sign}(\mu_i-\pi_i)$ simply records the sign of the probability gap for each coordinate.
We now specialize the TV derivative to the two top‑K tail‑handling schemes.
Aggregated‑tail estimator: the support consists of the retained set $K$ and a single tail bucket that aggregates all unseen mass.
The tail sign factor is $s_{\text{tail}}=\operatorname{sign}(\mu_{\text{tail}}-\pi_{\text{tail}})$.
Uniform‑tail estimator: the residual mass is spread uniformly over the $n-m$ unseen tokens, where $m=|K|$ and $n$ is the vocabulary size.
For any unseen token $j\notin K$ we approximate $\pi_j \approx \pi_{\text{tail}}/(n-m)$ and $\mu_j \approx \mu_{\text{tail}}/(n-m)$.
The masking rule itself is unchanged; only the proximity measure and the direction coefficient are swapped for their TV versions.
Experimental Details
Provides implementation, hyperparameter settings, and training dynamics for reproducibility.
We built the training pipeline on VeRL, leveraging Megatron‑LM for model training and vLLM for rollout generation.
All runs default to BF16; for the 30‑B model we also test FP8 in the rollout engine alone and FP8 end‑to‑end.
Each rollout batch forms a single PPO epoch; DPPO‑TopK‑KL and the predictive runs use a top‑K support of $K = 20$, with $\delta = 0.15$ for the main comparison and $\delta = 0.05$ for the tight‑threshold variant.
Training‑time diagnostics (Figures 5–8) show that the predictive divergence masks match DPPO‑TopK‑KL in PPO‑KL magnitude while consistently yielding lower clip‑fraction values.
**Table 1.** Training hyperparameters. Both Qwen3-30B-A3B-Base precision settings use the values in the final column.