Trust Region Masking for Long-Horizon LLM Reinforcement Learning

2025/12/28 by Yingru Li, Li, Yingru, Jiacai Liu +10
Computer Science · #Artificial Intelligence (cs.AI) #Domain Adaptation and Few-Shot Learning #FOS: Computer and information sciences #Information Theory (cs.IT) #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Multimodal Machine Learning Applications #Topic Modeling

paper · doi:10.48550/arxiv.2512.23075

openalex publication_date 2025/12/28 · openalex created_date 2025/12/31 · openalex updated_date 2026/07/28

Abstract

Policy gradient methods for Large Language Models optimize a policy π_θ via a surrogate objective computed from samples of a rollout policy π_roll. However, modern LLM-RL pipelines suffer from unavoidable implementation divergences -- backend discrepancies, Mixture-of-Experts routing discontinuities, and distributed training staleness -- causing off-policy mismatch (π_roll ≠ π_θ) and approximation errors between the surrogate and the true objective. We demonstrate that classical trust region bounds on this error scale as O(T²) with sequence length T, rendering them vacuous for long-horizon tasks. To address this, we derive a family of bounds -- both KL-based and TV-based -- including a Pinsker-Marginal bound (O(T^3/2)), a Mixed bound (O(T)), and an Adaptive bound that strictly generalizes the Pinsker-Marginal bound via per-position importance-ratio decomposition. Taking the minimum over all bounds yields the tightest known guarantee across all divergence regimes. Crucially, all bounds depend on the maximum token-level divergence D_KL^tok,max (or D_TV^tok,max), a sequence-level quantity that cannot be controlled by token-independent methods like PPO clipping. We propose Trust Region Masking (TRM), which masks entire sequences violating the trust region, enabling the first non-vacuous monotonic improvement guarantees for long-horizon LLM-RL.

Citations

• SimpleRL-Zoo: Investigating and Taming Zero Reinforcement Learning for Open Base Models in the Wild
• DAPO: An Open-Source LLM Reinforcement Learning System at Scale
• DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
• SWE-agent: Agent-Computer Interfaces Enable Automated Software Engineering
• Beyond Functional Correctness: Exploring Hallucinations in LLM-Generated Code
• BioXP-0.5B: Explainable Medical-AI via RL-GRPO
• SGLang: Efficient Execution of Structured Language Model Programs
• Efficient Memory Management for Large Language Model Serving with PagedAttention
• Let's Verify Step by Step
• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Fine-Tuning Language Models from Human Preferences
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures
• Proximal Policy Optimization Algorithms
• Constrained Policy Optimization
• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• Massively Parallel Methods for Deep Reinforcement Learning
• Scheduled Sampling for Sequence Prediction with Recurrent Neural Networks
• Trust Region Policy Optimization

Cited by

• Deterministic Inference across Tensor Parallel Sizes That Eliminates Training-Inference Mismatch

Related

• TROLL: Trust Regions improve Reinforcement Learning for Large Language Models
• Stabilizing Reinforcement Learning with LLMs: Formulation and Practices
• It's Not You, It's Clipping: A Soft Trust-Region via Probability Smoothing for LLM RL
• Dynamic Vocabulary Pruning: Stable LLM-RL by Taming the Tail
• Whatever Remains Must Be True: Filtering Drives Reasoning in LLMs, Shaping Diversity