PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking the Divergence Regularization in LLM RL

Jiarui Yao, Xiangxin Zhou, Penghui Qi, Wee Sun Lee, Liefeng Bo, Tianyu Pang

arXiv:2606.09821v1cs.LG

TL;DR

Off-policy LLM RL needs trust-region control, but ratio clipping can poorly reflect distributional shift in long-tailed vocabularies and DPPO discards harmful gradients beyond its boundary. DRPO replaces DPPO’s hard mask with a smooth, advantage-weighted quadratic regularizer while preserving absolute probability-shift geometry, and consistently enables stable, efficient training across six settings, matching or exceeding baseline accuracy.

  • Problem

    Ratio-based trust regions poorly reflect token-level distributional shift in long-tailed vocabularies, while DPPO discards gradients once harmful updates cross its boundary.

  • Method

    DRPO replaces DPPO’s hard mask with an advantage-weighted quadratic regularizer that preserves its absolute probability-shift trust-region geometry.

  • Results

    Across all six settings, DRPO consistently enables stable and efficient training, matching or exceeding the best evaluation accuracy achieved by baselines.

  • Takeaways & Limitations

    DRPO supports regularizer designs with bounded gradient weights, distribution-shift-aligned boundaries, and smooth corrective signals for LLM RL.

  • Takeaways & Limitations

    Because DRPO is advantage-weighted, its regularizer is not a pure divergence penalty.

Abstract

from arXiv · show

Reinforcement learning (RL) has become a key component of post-training large language models (LLMs). In practice, LLM RL is often off-policy because of training-inference mismatch and policy staleness, making trust-region control essential for stable optimization. Mainstream methods such as PPO and GRPO approximate this control with a ratio-clipping mechanism, but the importance ratio can be a poor proxy for distributional shift in long-tailed vocabularies. Recent work such as DPPO addresses this mismatch by replacing ratio-based clipping with a divergence-based mask, yielding a trust region defined by the sampled token's absolute probability shift. However, DPPO still relies on a hard mask: once a token crosses the trust-region boundary in a harmful direction, its gradient is discarded rather than corrected. To address this, we propose Divergence Regularized Policy Optimization (DRPO), which replaces the hard mask with a smooth advantage-weighted quadratic regularizer on policy shift. DRPO preserves the same trust-region geometry as DPPO while inducing bounded, continuous gradient weights that attenuate diverging updates and provide corrective signals beyond the boundary. Experiments across model scales, architectures, and precision settings show that DRPO improves the stability and efficiency of LLM RL training.

1 Introduction

LLM RL is typically off-policy, making trust-region control important, but ratio clipping poorly reflects distributional shift in long-tailed vocabularies. DRPO replaces DPPO’s hard divergence mask with a smooth, advantage-weighted regularizer that preserves Binary-TV geometry and supplies corrective gradients.

  • Motivation: Modern LLM RL is typically off-policy because inference and training engines differ numerically, creating a need for trust-region control.Rollouts are generated by inference engines, while optimization treats the LLM as an autoregressive token-level policy receiving scalar rewards.
  • Motivation: Ratio clipping poorly represents distributional shift: it over-constrains low-probability tokens and under-constrains high-probability tokens.A small probability increase can create a large ratio with little mass change, whereas a moderate ratio change can substantially shift high-probability mass.
  • Prior limitation: DPPO replaces ratio clipping with a Binary-TV divergence mask based on the sampled token’s absolute probability shift.The mask disables gradients when divergence exceeds a threshold and the update would increase it further.
  • Proposed method: DRPO preserves DPPO’s Binary-TV trust-region geometry while replacing its hard mask with an advantage-weighted ℓ2 regularizer.The construction rewrites the Binary-TV constraint as a token-adaptive ratio bound and applies the SPO construction.
  • Proposed method: DRPO uses continuous gradient weights that decay to zero at the boundary, become corrective beyond it, and amplify updates moving toward the behavior policy.The weights depend on each token’s Binary-TV shift and update direction; DPPO instead provides no corrective signal beyond the boundary and can change abruptly near it.
  • Design criteria: Ablations indicate that effective LLM RL regularizers should induce stable gradients and weight per-token penalties by absolute advantage.Standard KL or TV penalties can underperform because their gradients reintroduce ratio-based geometry; absolute-advantage weighting keeps the boundary independent of reward scale.

2 Background

LLM generation can be modeled as a token-level MDP with rewards assigned to sampled responses, while numerical mismatches and minibatch reuse make RL training off-policy. Trust-region methods address variance and instability, evolving from ratio-based PPO/GRPO/SPO constraints toward DPPO’s divergence-based control over absolute probability shifts.

  • LLM RL formulation: LLM generation is a token-level MDP in which a policy autoregressively samples responses and receives a scalar reward from a reward model or verifier.The state contains the prompt and generated prefix, and generation ends at [eos] or the token limit.
  • Off-policy training: Training-inference numerical discrepancies and reuse of rollout batches across mini-batches make LLM RL off-policy.The resulting data is sampled from a behavior policy µ.
  • Trust-region methods: TRPO reduces importance-sampling variance with a token-level surrogate and guarantees monotonic improvement within a KL- or TV-defined trust region.Its second-order optimization is computationally prohibitive at scale, motivating simpler alternatives.
  • PPO and GRPO: PPO approximates the trust region by deactivating gradients outside |r_t −1| ≤ϵ, while GRPO estimates advantages from groups of responses without a critic.PPO’s hard clipping can abruptly switch gradients to zero and provides no corrective signal after harmful boundary crossing.
  • Smooth regularization: SPO replaces PPO’s hard clip with a continuous gradient weight while preserving the same ratio-based trust-region boundary.Its per-token objective is a concave quadratic in the importance ratio.
  • DPPO: Ratio-based regions can misrepresent distributional shift in long-tailed vocabularies, so DPPO uses divergence-based masking that constrains sampled-token absolute probability shift.DPPO’s Binary-TV approximation collapses the distribution into the sampled token versus the rest and masks only harmful outward updates beyond threshold δ.

3 Method

DRPO derives a smooth, advantage-weighted quadratic regularizer from DPPO’s Binary-TV trust region. It preserves DPPO’s absolute-probability-shift geometry while continuously attenuating diverging updates and reversing gradients beyond the boundary.

  • DRPO regularization: DRPO replaces DPPO’s discontinuous binary mask with continuous gradient reweighting based on the magnitude and direction of each token’s probability shift.The regularizer preserves the same trust-region boundary while providing corrective signals outside it.
  • DRPO regularization: DRPO scales its quadratic regularizer by the behavior probability, changing the equilibrium from a fixed ratio shift to a fixed absolute probability shift.This factor yields the Binary-TV geometry required by DPPO rather than the ratio geometry of PPO and SPO.
  • Trust-region behavior: For diverging updates, DRPO attenuates gradients inside the trust region and reverses them outside it, directing policy updates back toward the boundary.For converging updates, the gradient is amplified to encourage movement toward the behavior policy.
  • Geometry and stability: DRPO penalizes advantage-weighted squared ℓ2 distance, whereas SPO uses an advantage-weighted Pearson-χ2 penalty dominated by low-probability tokens.DRPO assigns the same cost to the same absolute probability shift regardless of token probability at a fixed advantage value.
  • Geometry and stability: 7.8% of sampled tokens have µ(yt|st) ≤0.01, where ratio-based SPO weights can become unstable; DRPO’s absolute shift is bounded in [0, 1] with variance bounded by 1/4.DRPO therefore avoids unbounded low-tail weighting while retaining smooth divergence-based reweighting.

4 Experiments and Results

Across six experimental settings, DRPO consistently provides stable, efficient training and matches or exceeds baseline accuracy. Ablations further attribute its gains to smooth regularization, advantage weighting, and behavior beyond DPPO’s trust-region boundary.

  • Main results: DRPO consistently enables stable, efficient training across all six settings, matching or exceeding the best baseline evaluation accuracy.Experiments use Qwen3 and DeepSeek-R1-Distill-Qwen models with DAPO-derived math data, AIME 2024/2025 evaluation, and BF16 or challenging FP8 settings.
  • Main results: Ratio-based GRPO and SPO frequently train unstably, especially under low precision, where they often collapse before reaching reasonable performance.Even their strongest settings lag behind divergence-based methods in training efficiency and final accuracy.
  • Main results: Hard-mask methods such as GRPO and DPPO often underperform smooth regularization; DPPO can train stably yet converge more slowly and reach lower final accuracy than DRPO.The results support using a continuous gradient signal rather than discarding updates beyond a boundary.
  • Main results: The unregularized surrogate drops from 0.25 to 0.17 on Qwen3-4B-Base and loses performance in three of six settings, showing that a proper trust region remains necessary.Although the unregularized objective can perform strongly in some cases, it is unreliable across settings.
  • Ablations: Removing the absolute-advantage weight consistently degrades performance and destabilizes training, indicating that a stable trust-region boundary matters more than a pure divergence form.This ablation was conducted on Qwen3-30B-A3B-Base FP8-E2E and R1D for SPO and DRPO.
  • Ablations: Alternative forward-KL and TV penalties underperform DRPO, while Mask-DRPO performs comparably to DRPO when regularization is applied only outside DPPO’s trust region.The comparison supports DRPO’s smooth Binary-TV trust-region behavior and identifies corrective gradients beyond DPPO’s boundary as the main source of its gain.

5 Closing Remarks · Appendix A. Related Work

The paper argues that regularizer design should prioritize the induced gradient geometry over the nominal divergence, especially under long-tailed vocabularies. DRPO replaces DPPO’s hard mask with a smooth quadratic regularizer that preserves trust-region geometry while providing corrective gradients beyond the boundary.

  • 5 Closing Remarks: Induced gradient form matters more than the nominal divergence used in the objective.The authors compare regularizers based on their optimization behavior rather than their stated divergence measure.
  • 5 Closing Remarks: Absolute-advantage terms are essential for maintaining a stable trust-region boundary.These terms prevent the regularizer from being a pure divergence while supporting stable boundary behavior.
  • 5 Closing Remarks: Objective-level regularizers can fail when their gradients induce undesirable trust-region geometry.The paper identifies gradient behavior as the key criterion for evaluating regularization choices.
  • 5 Closing Remarks: Ratio-based trust regions can produce high-variance, unbounded weights in long-tailed LLM vocabularies.The authors identify this as a common failure mode of ratio-based gradient design.
  • 5 Closing Remarks: Binary-TV’s absolute probability shift provides a better alternative to ratio-based trust-region control.The passage contrasts absolute probability shift with ratio-based weights under long-tailed vocabularies.
  • 5 Closing Remarks: DRPO replaces DPPO’s hard mask with a smooth quadratic regularizer while preserving the same trust-region geometry.Unlike DPPO, DRPO avoids abrupt gradient changes near the boundary and supplies corrective signals outside the trust region.

A.1 Traditional RL based on Trust Region Methods · A.2 RL for LLM Reasoning

Traditional trust-region RL progresses from computationally expensive second-order constraints to first-order clipping and smooth regularization. For LLM reasoning, off-policy training makes trust-region control essential, while DRPO combines divergence-based geometry with smooth, bounded corrective updates.

  • A.1 Traditional RL based on Trust Region Methods: TRPO guarantees monotonic improvement through KL-constrained second-order optimization, while CPO extends trust-region control to constrained MDPs.Both methods require second-order optimization, which is prohibitive at scale.
  • A.1 Traditional RL based on Trust Region Methods: PPO replaces explicit KL constraints with ratio clipping and first-order optimization, but clipping does not strictly bound likelihood ratios or enforce a defined divergence constraint.Truly PPO adds rollback clipping with KL-based triggering, while Trust Region-Guided PPO adapts clipping thresholds.
  • A.1 Traditional RL based on Trust Region Methods: SPO replaces PPO’s hard clipping with a smooth quadratic importance-ratio regularizer whose per-token optimum matches PPO’s clipping boundary.It also provides non-zero corrective gradients outside the trust region.
  • A.1 Traditional RL based on Trust Region Methods: DRPO changes SPO’s ratio-based geometry to divergence-based geometry, shifting the per-token optimum from |r_t − 1| = ϵ to |π(y_t|s_t) − μ(y_t|s_t)| = δ.The resulting boundary follows Binary-TV geometry while retaining SPO’s smooth gradient structure.
  • A.2 RL for LLM Reasoning: LLM RL is inherently off-policy because of training-inference mismatch and mini-batch policy staleness, making trust-region optimization essential for stable reasoning training.RL is described as a key technique for improving LLM reasoning.
  • A.2 RL for LLM Reasoning: LLM RL commonly uses PPO-style ratio clipping, with GRPO substituting group-relative advantages, while DAPO, CISPO, and M2PO modify clipping or importance-weight constraints.DAPO widens the upper clipping bound, CISPO uses truncated importance sampling, and M2PO constrains the second moment of importance weights.
  • A.2 RL for LLM Reasoning: Regularization-based approaches include online policy mirror descent, ratio-variance penalties, and per-token KL projections, but they retain ratio dependence, risk unbounded weights, or require per-token optimization.R2VPO can induce unbounded gradient weights for low-probability tokens, while TROLL solves an optimization problem for each token.
  • A.2 RL for LLM Reasoning: DPPO replaces ratio clipping with a TV- or KL-based binary mask, but its boundary causes abrupt gradient changes and supplies no corrective signal afterward.DRPO combines DPPO’s divergence-based geometry with a lightweight advantage-weighted ℓ2 regularizer, attenuating diverging updates, amplifying corrective ones, and keeping per-token weights bounded.

Appendix B. Detailed Derivation of the Gradient of DRPO

This appendix begins the derivation of the gradient for the DRPO objective in Equation (8).

  • Appendix B. Detailed Derivation of the Gradient of DRPO: The appendix introduces a derivation of the gradient of the objective in Equation (8).The supplied passage signals that the derivation follows, but does not provide its intermediate steps or final expression.

Appendix C. Induced Trust Regions of Alternative Regularizers

The appendix shows that alternative regularizers can induce trust-region geometries unlike their nominal divergences suggest. KL and K3 remain ratio-based, TV is binary, whereas DRPO continuously weights absolute probability shifts and preserves Binary-TV geometry.

  • Motivation: Trust-region regularizers should be evaluated by their token-level gradient weights, because sampled-token updates in long-tailed vocabularies depend on induced gradient geometry.Loss-level similarity between objectives does not guarantee similar optimization behavior after importance sampling.
  • Advantage-weighted KL: Advantage-weighted KL induces ratio-based boundaries, with finite cancellation depending on π_t/µ_t rather than absolute probability shift.For positive-advantage tokens, the sampled forward-KL form has no finite rollback point; for negative-advantage tokens, the boundary scales with µ_t.
  • Advantage-weighted KL with K3: The K3 KL estimator reduces estimation variance but leaves the trust-region geometry ratio-based, so its corrective force remains calibrated in importance-ratio space.Changing the estimator improves numerical behavior without changing the relevant boundary geometry.
  • Advantage-weighted TV: Advantage-weighted TV produces a binary, piecewise-constant gradient weight that distinguishes only update direction, not the magnitude of |r_t −1|.Its absolute-value form also has a nondifferentiable kink at r_t = 1, preventing gradual attenuation within the trust region.
  • Summary: DRPO’s weight depends continuously on absolute probability shift, preserving DPPO’s Binary-TV geometry while replacing its hard mask with a corrective smooth update.The appendix reports that ratio-based or binary-gradient penalties are less stable than the Binary-TV quadratic penalty.

Appendix D. More Experimental Details · D.1 Comparing with KL Regularization

Appendix D details the experimental setup and compares DRPO with a pure KL regularizer. DRPO consistently outperforms the KL alternative across six experiments, with the gap attributed to advantage scaling and ratio-based trust-region effects.

  • Appendix D. More Experimental Details: vLLM serves as the inference backend to speed up rollout, and math-verify2 verifies solutions in math reasoning tasks.These components complement the VeRL and Megatron implementation described in the experimental setup.
  • Appendix D. More Experimental Details: The experiments examine advantage scaling, different divergences, and binary approximation across several objective-function variants.Typical training runs use 800 steps for Qwen3-4B-Base, 300 for Qwen3-30B-A3B-Base, 110 for Qwen3.5-35B-A3B-Base, and 3000 for R1D.
  • D.1 Comparing with KL Regularization: The comparison adds a pure KL regularizer without advantage weighting to the baselines from Section 4.1.This method follows the KL formulation associated with Schulman et al. (2015) and Schulman et al. (2017).
  • D.1 Comparing with KL Regularization: The KL comparison uses the same hyperparameter δ = 12.5 as the tuned setting.The passage refers to Figure 8 for hyperparameter-tuning results.
  • D.1 Comparing with KL Regularization: DRPO consistently outperforms the pure KL regularizer across all six experiments.Figure 6 reports training dynamics for DRPO and direct KL penalization without the advantage weight | ˆAt|.
  • D.1 Comparing with KL Regularization: The performance gap reflects two issues: KL regularization lacks per-token advantage adaptation and still induces a ratio-based trust region.Removing | ˆAt| degrades performance because the token-wise optimum depends on current advantage magnitude; the passage also identifies ratio-based trust-region behavior as a remaining limitation.

D.2 Extended Ablations on Advantage Weighting

Extended ablations show that weighting either KL or TV regularization by |Â_t| consistently improves performance. Without this weight, the trust-region boundary depends on advantage magnitude, harming update calibration, training stability, and final accuracy.

  • Advantage Weighting: Adding |Â_t| clearly improves performance for both KL and TV penalties across all evaluated settings.Figure 7 compares each penalty with and without the advantage weight.
  • Advantage Weighting: Without |Â_t|, the effective trust-region boundary becomes entangled with advantage magnitude instead of remaining stable.This makes updates overly restrictive for small-advantage tokens and too loose for large-advantage ones.
  • Advantage Weighting: Noisy token-level advantages further worsen the resulting mismatch in training stability and final accuracy.The analysis links this behavior to Section 3.1 and the ablations in Section 4.2.

D.3 Hyperparameter Tuning of Advantage-Weighted KL Regularizer · D.4 Hyperparameter Tuning of DPPO Baseline · D.5 Hyperparameter Tuning of DRPO

The hyperparameter studies show that DRPO remains stronger than an advantage-weighted KL baseline and a tuned DPPO baseline, while requiring less fine-grained tuning. DRPO is also relatively robust across threshold choices, with only a minor performance drop when δ decreases from 12.5 to 2.5.

  • D.3 Hyperparameter Tuning of Advantage-Weighted KL Regularizer: Under R1D, the study sweeps hyperparameters for the advantage-weighted KL regularizer to test whether the KL baseline was simply under-tuned.The sweep uses the regularizer from Equation (13).
  • D.3 Hyperparameter Tuning of Advantage-Weighted KL Regularizer: Figure 8 shows that DRPO remains stronger across the full sweep, even when the KL baseline uses the same advantage weight | ˆAt|.This result addresses the possibility that KL underperformance reflects insufficient tuning.
  • D.3 Hyperparameter Tuning of Advantage-Weighted KL Regularizer: The robustness gap is attributed to DRPO’s Binary-TV trust-region geometry versus the KL penalty’s ratio-based geometry.The paper states that Binary-TV more faithfully reflects the intended divergence constraint and supports more reliable optimization behavior.
  • D.4 Hyperparameter Tuning of DPPO Baseline: DPPO requires finer threshold tuning: ε = 0.15 works best on Qwen3-30B-A3B-Base, yet remains worse than DRPO, while ε = 0.6 works best on FP8-E2E with similar performance to DRPO.Unlike DPPO, DRPO provides corrective gradients for tokens outside the trust region.
  • D.4 Hyperparameter Tuning of DPPO Baseline: The DPPO sweep supports δ = 12.5 as a relatively universal DRPO hyperparameter.The passage contrasts DPPO’s setting-dependent best thresholds with DRPO’s relatively universal δ = 12.5.
  • D.5 Hyperparameter Tuning of DRPO: Reducing DRPO’s threshold from 12.5 to 2.5 causes only a minor performance drop, indicating robustness across a broad hyperparameter range.The sensitivity test evaluates two choices of δ.

D.6 Mask Ablation with Alternative Divergence Penalties

The ablation shows that DRPO’s main benefit comes from correcting tokens that cross DPPO’s trust-region boundary, while the choice of divergence penalty determines calibration and stability. The Binary-TV quadratic penalty used by DRPO performs best and most stably among the tested alternatives.

  • Trust-region behavior: Inside the trust region, DRPO matches DPPO’s active gradient; outside it, DRPO restores the corrective penalty gradient discarded by DPPO’s hard mask.This preserves the unregularized token surrogate within the region while providing correction beyond its boundary.
  • Penalty comparison: The framework instantiates three divergence-penalty objectives for the mask ablation.The passage introduces these objectives after defining the trust-region indicator.
  • Penalty comparison: The squared log-ratio term in Equation (25) is exactly the penalty used in online policy mirror descent by the Kimi series.This connects one tested penalty to prior online policy mirror descent formulations.
  • Trust-region behavior: Applying the DRPO penalty only outside DPPO’s trust region performs close to applying it everywhere, confirming that boundary-crossing tokens drive the main gain.The comparison is reported in Figure 11.
  • Penalty comparison: Binary-TV quadratic regularization gives the best and most stable behavior, whereas ratio-space and KL-type penalties are harder to calibrate.Its gradient directly follows absolute probability shift.
Loading 2606.09821v1…