Trust Region Masking for Long-Horizon LLM Reinforcement Learning
Yingru Li ⋅ Jiacai Liu ⋅ Jiawei Xu ⋅ Yuxuan Tong ⋅ Ziniu Li ⋅ Baoxiang Wang
Abstract
Policy gradient methods for Large Language Models optimize a policy $\pi_\theta$ via a surrogate objective computed from samples of a rollout policy $\pi_{roll}$. However, modern LLM-RL pipelines suffer from unavoidable implementation divergences, such as backend discrepancies, Mixture-of-Experts routing discontinuities, and distributed training staleness. These factors cause an off-policy mismatch ($\pi_{roll} \neq \pi_\theta$), leading to approximation errors between the surrogate and the true objective. We demonstrate that classical trust region bounds on this error scale as $O(T^2)$ with sequence length $T$, rendering them vacuous for long-horizon tasks. To address this, we derive two new bounds: a Pinsker-Marginal bound scaling as $O(T^{1.5})$ and a Mixed bound scaling as $O(T)$. Crucially, both bounds depend on $D_{KL}^{max}$, the maximum token-level KL divergence across the sequence. As a sequence-level term, the divergence cannot be controlled by previous token-independent methods like PPO clipping. We propose Trust Region Masking (TRM), which masks entire sequences that violate the trust region. TRM reduces the vacuity of traditional bounds, offering a principled framework to mitigate training collapse in long-horizon LLM-RL.
Lay Summary
- Problem: In modern LLM-RL, off-policy mismatch is not a minor implementation bug but an inevitable systemic conflict between generation and training software. Over long responses, these minor discrepancies snowball, rendering traditional safety bounds completely useless and causing the learning process to catastrophically collapse. - Solution: We establish that stable learning depends strictly on the maximum divergence at any single point in a response—a bottleneck that standard token-level clipping cannot control. We develop Trust Region Masking (TRM), a sequence-level filter that instantly throws away any generated response where a single token violates safety thresholds. - Impact: TRM completely eliminates training collapse, providing the first mathematically sound framework to keep LLM-RL training stable over extended horizons.
Successful Page Load