Rethinking the Trust Region in LLM Reinforcement Learning

Paper: 2602.04879 Authors: Penghui Qi, Xiangxin Zhou, Zichen Liu, Tianyu Pang, Chao Du, Min Lin, Wee Sun Lee Affiliations: Sea AI Lab, National University of Singapore Code: https://github.com/sail-sg/Stable-RL Categories: cs.LG

Abstract

Reinforcement learning (RL) has become a cornerstone for fine-tuning Large Language Models (LLMs), with Proximal Policy Optimization (PPO) serving as the de facto standard algorithm. Despite its ubiquity, we argue that the core ratio clipping mechanism in PPO is structurally ill-suited for the large vocabularies inherent to LLMs.

The Problem with PPO’s Ratio Clipping

PPO constrains policy updates based on the probability ratio of sampled tokens:

L_PPO = E_t[min(r_t * A_t, clip(r_t, 1-ε, 1+ε) * A_t)]

where r_t = π_θ(a_t|s_t) / π_old(a_t|s_t)

This is a noisy single-sample Monte Carlo estimate of true policy divergence.

Why This Fails for LLMs

Over-penalizing low-probability tokens:

• Increasing a rare token’s probability from 10^-5 to 10^-3 creates ratio = 100
• Triggers aggressive clipping even though actual divergence is negligible
• Slows down learning on important but rare tokens

Under-penalizing high-probability tokens:

• A drop from 0.99 to 0.8 has small ratio change
• But this is a catastrophic shift in probability mass
• Often remains unpenalized by clipping

Training-Inference Mismatch

Numerical discrepancies between training and inference engines make the probability ratio highly volatile for low-probability tokens, while TV divergence remains stable.

Key Contributions

• Theoretical Formulation: Policy improvement bounds tailored to finite-horizon, undiscounted LLM generation
• Stability Analysis: Identifies sources of training instability and the role of low-probability tokens in exploration
• DPPO Algorithm: Divergence Proximal Policy Optimization with principled trust region constraints

DPPO: The Solution

Replace heuristic clipping with direct divergence estimation:

Policy Improvement Bound for LLMs

For finite-horizon T with no discount (γ=1):

J(π) - J(μ) = L'_μ(π) - Δ(μ, π)

Where the penalty term depends on Total Variation (TV) divergence or KL divergence, not probability ratios.

Efficient Approximations

To avoid huge memory footprint, DPPO introduces:

Binary Divergence:

D_Binary = |p_θ(a*) - p_old(a*)|

Only track the sampled token’s probability change.

Top-K Divergence:

D_TopK = Σ_{i∈TopK} |p_θ(a_i) - p_old(a_i)|

Track top-K most probable tokens.

These capture essential distributional shifts with negligible overhead.

The DPPO Objective

L_DPPO = E_t[r_t * A_t]    subject to D_TV(μ || π) ≤ δ

Uses a principled mask based on direct approximation of policy divergence, ensuring updates stay within a theoretically grounded trust region.

Results

On AIME24 using Qwen3-30B-A3B-Base:

• DPPO significantly outperforms GRPO baselines
• Superior training efficiency and stability
• Works even without rollout routing replay (R3)

Why This Matters

Aspect	PPO	DPPO
Constraint	Probability ratio	Direct divergence
Low-prob tokens	Over-penalized	Properly handled
High-prob tokens	Under-constrained	Properly constrained
Theoretical basis	Heuristic	Trust region theory

Takeaways

1. Ratio clipping is flawed for LLMs: The large vocabulary makes single-sample ratio estimates unreliable
2. Direct divergence estimation: TV or KL divergence provides more stable and principled constraints
3. Efficient approximations exist: Binary and Top-K divergence capture essential shifts with minimal overhead
4. Low-probability tokens matter: They drive exploration and shouldn’t be over-penalized
5. Training stability: DPPO provides a more robust foundation for RL-based LLM fine-tuning

L_PPO = E_t[min(r_t * A_t, clip(r_t, 1-ε, 1+ε) * A_t)]

其中 r_t = π_θ(a_t|s_t) / π_old(a_t|s_t)

J(π) - J(μ) = L'_μ(π) - Δ(μ, π)

D_Binary = |p_θ(a*) - p_old(a*)|

D_TopK = Σ_{i∈TopK} |p_θ(a_i) - p_old(a_i)|

L_DPPO = E_t[r_t * A_t]    约束条件 D_TV(μ || π) ≤ δ