Trust Region Masking for Long-Horizon LLM Reinforcement Learning

Authors: Yingru Li, Jiacai Liu, Jiawei Xu, Yuxuan Tong, Ziniu Li, Baoxiang Wang

📄 arXiv: https://arxiv.org/abs/2512.23075

📊 Slides: trust_region_masking_slides_yingru.pdf

🔧 PR #4544: feat: trust region sequence masking

Abstract

Policy gradient methods for large language models optimize a surrogate objective computed from samples of a rollout policy $\pi_{\mathrm{roll}}$. When $\pi_{\mathrm{roll}} \neq \pi_{\theta}$, there is approximation error between the surrogate and the true objective. Prior work has shown that this off-policy mismatch is unavoidable in modern LLM-RL due to implementation divergence, mixture-of-experts routing discontinuities, and distributed training staleness.

Classical trust region bounds on the resulting error scale as $O(T^2)$ with sequence length $T$, rendering them vacuous for long-horizon tasks.

The Long-Horizon Problem

Given that $\pi_{\mathrm{roll}} \neq \pi_{\theta}$ is unavoidable, approximation error becomes a central concern. Classical error bounds scale as $O(T^2)$ with sequence length $T$. For modern LLMs generating responses of $T=4096$ tokens, these bounds become vacuous: even with small per-token divergence ($D_{\mathrm{KL}}^{\mathrm{tok,max}} = 10^{-4}$), the classical bound yields an error of $\approx 1677$, far exceeding any plausible improvement.

Off-Policy Mismatch Sources

Three factors contribute to the unavoidable mismatch between $\pi_{\mathrm{roll}}$ and $\pi_{\theta}$:

1. Implementation divergence: Different numerical implementations for inference (vLLM, SGLang) versus training (Megatron-LM, PyTorch FSDP) produce different logits from identical weights.

2. MoE routing discontinuities: In mixture-of-experts models, small numerical differences can trigger different expert selections, causing discrete jumps in token probabilities.

3. Distributed staleness: Asynchronous training pipelines create lag between rollout generation and gradient computation, so training occurs with updated weights $\pi_{\theta}$ while rollouts were generated with stale weights $\pi_{\mathrm{roll}}$.

Our Contributions

We derive two tighter bounds:

Crucially, both bounds depend on $D_{\mathrm{KL}}^{\mathrm{tok,max}}$ — the maximum token-level KL divergence across all positions in a sequence. This is inherently a sequence-level quantity: it requires examining the entire trajectory to compute, and therefore cannot be controlled by token-independent methods like PPO clipping.

Why Token-Level Methods Fail

PPO Clipping

PPO clipping attempts to control divergence by clipping the importance ratio. However, this suffers from gradient leakage — clipping only affects the gradient magnitude, not the fundamental approximation error.

Token Masking

Token-level masking excludes individual tokens that violate the trust region. However, this creates a theoretical problem: the masked gradient is no longer an unbiased estimator of the policy gradient.

The Fundamental Dilemma

The root cause is that $D_{\mathrm{KL}}^{\mathrm{tok,max}}$ is a sequence-level quantity that cannot be decomposed into independent token-level constraints. The only solution is to operate at the sequence level.

Solution: Trust Region Masking (TRM)

We propose Trust Region Masking (TRM), which excludes entire sequences from gradient computation if any token violates the trust region.

Why Sequence Masking Works

By masking at the sequence level, we ensure that:

1. The remaining samples come from policies that are provably close
2. The gradient estimator remains unbiased over the unmasked samples
3. The trust region bound applies to all included sequences

Masking Criterion

A sequence is masked if: $$\max_{t} D_{\mathrm{KL}}(\pi_{\mathrm{roll}}(\cdot|c_t) | \pi_{\theta}(\cdot|c_t)) > \epsilon$$

where $\epsilon$ is the trust region threshold.

Exact Computation

The rigorous guarantee requires exact KL computation with stored logits from the rollout policy.

Sample-Based Approximation

In practice, storing full logits may be expensive. The paper proposes sample-based approximations using importance ratios $\rho_t = \frac{\pi_{\theta}(y_t|c_t)}{\pi_{\mathrm{roll}}(y_t|c_t)}$:

The $k_3$ Estimator (for average-based filtering)

$$k_3(\rho) = \rho - 1 - \log \rho$$

Properties of $k_3$:

• Non-negative: $\rho - 1 - \log \rho \geq 0$ for all $\rho > 0$
• Unbiased: $\mathbb{E}{y \sim \pi{\mathrm{roll}}}[k_3(\rho)] = D_{\mathrm{KL}}$
• Ideal for average-based filtering since $(1/T)\sum_t k_3(\rho_t)$ converges to true average KL

The $|\log \rho|$ Estimator (for max-based filtering)

For the max criterion, we need a symmetric detector since both $\rho \gg 1$ and $\rho \ll 1$ indicate large divergence: $$|\log(100)| = |\log(0.01)| = 4.6$$

In contrast, $k_3$ is asymmetric: $k_3(100) = 94.4$ but $k_3(0.01) = 3.6$ (26× difference).

Caveat

Neither sample-based method provides a rigorous bound on $D_{\mathrm{KL}}^{\mathrm{tok,max}}$ — both are approximate detectors based on single samples per context.

Key Result

TRM provides the first non-vacuous monotonic improvement guarantees for long-horizon LLM-RL.

Paper Structure

1. Introduction: Off-Policy Mismatch, Long-Horizon Problem, Contributions
2. Background: Autoregressive Language Generation, Optimization Problem, Surrogate Objective, Divergence Measures
3. Theoretical Analysis: Performance Difference Identity, Key Lemmas, Classical TRPO Bound, New Bounds
4. Why Token-Level Methods Fail: PPO Clipping, Token Masking, Fundamental Dilemma
5. Solution: Trust Region Masking
6. Discussion: Limitations, Practical Considerations, Future Work
7. Conclusion

Citation

@article{li2025trustregion,
  title   = {Trust Region Masking for Long-Horizon LLM Reinforcement Learning},
  author  = {Li, Yingru and Liu, Jiacai and Xu, Jiawei and Tong, Yuxuan and Li, Ziniu and Wang, Baoxiang},
  journal = {arXiv preprint arXiv:2512.23075},
  year    = {2025},
  url     = {https://arxiv.org/abs/2512.23075}
}