Technical Summary: Rethinking the Divergence Regularization in LLM RL

1. Problem Statement

Reinforcement learning (RL) is a cornerstone of post-training large language models (LLMs) to align with human preferences and enhance reasoning capabilities. However, modern LLM RL training is inherently off-policy due to mismatches between training and inference engines (numerical discrepancies) and the practice of splitting rollouts into multiple mini-batches. This necessitates Trust Region Policy Optimization (TRPO) to ensure stable updates by constraining the divergence between the current policy ($\pi$) and the behavior policy ($\mu$).

Current mainstream methods face two primary limitations in this context:

1. Ratio-Based Constraints (PPO, GRPO, SPO): These methods use the importance ratio $r_t = \pi(y_t|s_t)/\mu(y_t|s_t)$ to define the trust region. In LLMs with large, long-tailed vocabularies, the ratio is a poor proxy for distributional shift. A small probability increase on a low-probability token can yield an enormous ratio with negligible mass shift, while a moderate ratio change on a high-probability token can induce significant policy changes. This leads to over-constraining rare tokens and under-constraining common ones.
2. Hard Masking (DPPO): Recent work like Divergence Proximal Policy Optimization (DPPO) addresses the ratio issue by using a divergence-based mask (specifically Binary-TV, measuring absolute probability shift $|\pi - \mu|$). However, DPPO enforces the trust region via a hard binary mask. Once a token crosses the boundary in a harmful direction, its gradient is discarded (set to zero). This prevents further divergence but fails to provide a corrective signal to pull the policy back toward the trust region, leading to brittle optimization near the threshold.

2. Methodology: Divergence Regularized Policy Optimization (DRPO)

The authors propose DRPO, a method that replaces the hard mask of DPPO with a smooth, advantage-weighted quadratic regularizer on the policy shift, while preserving the same Binary-TV trust region geometry.

Core Formulation

DRPO derives its objective by rewriting the Binary-TV constraint $|\pi(y_t|s_t) - \mu(y_t|s_t)| \leq \delta$ as a token-adaptive ratio bound. Instead of clipping, it applies a smooth regularizer to the surrogate objective:

$$\mathcal{L}_{DRPO}(x, \pi) = \mathbb{E}_{y \sim \mu} \left[ \sum_{t=1}^{|y|} \left( r_t \hat{A}_t - \frac{|\hat{A}_t|}{2\delta} \mu(y_t|s_t) (r_t - 1)^2 \right) \right]$$

Where:

• $r_t$ is the importance ratio.
• $\hat{A}_t$ is the advantage estimate.
• $\delta$ is the trust region threshold.
• The term $\mu(y_t|s_t)$ scales the regularizer, converting the penalty from a ratio-based $\chi^2$ form (as in SPO) to an absolute probability shift form ($\ell_2^2$).

Gradient Behavior

The gradient of the DRPO objective introduces a continuous weight $w_t$ for each token's policy gradient:
$$w_t = 1 - \text{sign}(\hat{A}_t(r_t - 1)) \frac{|\pi(y_t|s_t) - \mu(y_t|s_t)|}{\delta}$$

This weight exhibits three key behaviors:

1. Inside the Trust Region: The weight is positive but smoothly attenuated as the shift approaches $\delta$, allowing updates to continue but with reduced magnitude.
2. Outside the Trust Region (Diverging): When the update moves further away from the behavior policy ($|\pi - \mu| > \delta$), the weight becomes negative. This reverses the gradient, providing a corrective signal that pulls the policy back toward the boundary, unlike the zero-gradient of hard masks.
3. Converging Updates: When the update moves back toward the behavior policy, the weight is amplified ($>1$), encouraging a smooth return to the trust region.

Theoretical Distinctions

• vs. SPO: SPO uses a $\chi^2$-type regularizer (dividing by $\mu$), which causes gradient weights to grow unbounded in the long-tail of low-probability tokens. DRPO uses an $\ell_2^2$-type regularizer (multiplying by $\mu$), ensuring gradient weights remain bounded regardless of token probability.
• vs. DPPO: DRPO preserves the Binary-TV geometry of DPPO but replaces the discontinuous mask with a smooth, differentiable regularizer.

3. Key Contributions

1. DRPO Algorithm: A novel policy optimization method that replaces hard divergence masks with a smooth, quadratic regularizer based on absolute probability shift (Binary-TV).
2. Bounded Gradient Weights: Theoretical and empirical demonstration that DRPO's gradient weights remain bounded in long-tailed vocabularies, avoiding the instability caused by ratio-based methods (PPO, SPO) where weights can explode for low-probability tokens.
3. Corrective Mechanism: Introduction of a continuous corrective signal for updates that exceed the trust region, preventing the "gradient death" observed in hard-mask approaches like DPPO.
4. Design Principles: Identification of three practical criteria for effective regularizers in LLM RL:
   • Induce a stable boundary aligned with distributional shift (not just ratios).
   • Keep per-token gradient weights bounded in long-tailed vocabularies.
   • Provide smooth corrective signals when the policy diverges.

4. Experimental Results

The authors evaluated DRPO across diverse settings, including:

• Models: Qwen3-4B, Qwen3-30B-A3B, Qwen3.5-35B-A3B, and DeepSeek-R1-Distill-Qwen-1.5B.
• Architectures: Dense and Mixture-of-Experts (MoE).
• Precision: BF16, FP8 (rollout only), and FP8 (end-to-end).
• Benchmarks: AIME 2024 and AIME 2025 (math reasoning tasks).

Key Findings:

• Stability: DRPO consistently enables stable training across all settings. In contrast, ratio-based methods (GRPO, SPO) frequently suffered from instability or collapse, particularly in low-precision (FP8) settings.
• Performance: DRPO matched or exceeded the best evaluation accuracy of all baselines. It outperformed hard-mask methods (DPPO) in convergence speed and final accuracy.
• Ablations:
  • Removing the absolute advantage weighting $|\hat{A}_t|$ from the regularizer led to performance drops and instability, confirming that the trust region boundary must be independent of reward scale.
  • Alternative divergence measures (KL, TV penalties) underperformed DRPO, suggesting that the specific geometry of the Binary-TV regularizer is crucial.
  • Applying the regularizer only outside the trust region (Mask-DRPO) yielded performance comparable to full DRPO, indicating that the primary gain comes from the corrective regularization outside the boundary.

5. Significance and Claims

The paper claims that the geometry of the gradient induced by a regularizer is more critical than the nominal divergence measure in the objective function. While prior works often focus on standard divergences (KL, JS), this work demonstrates that:

• Ratio-based constraints are fundamentally flawed for LLMs due to long-tailed vocabularies.
• Hard masks, while better than ratios, introduce brittleness by discarding corrective signals.
• DRPO offers a unified solution that combines the geometric correctness of divergence-based constraints (Binary-TV) with the optimization stability of smooth regularizers.

The authors position DRPO not just as a specific algorithm, but as an empirical lens for designing stable policy-optimization objectives, emphasizing that effective regularizers must provide bounded, continuous, and corrective gradient signals to handle the unique challenges of LLM RL.