Trust Region Masking for Long-Horizon LLM Reinforcement Learning

Trust region methods [Kakade & Langford, 2002, Schulman et al., 2015] provide a principled framework for policy optimization by constructing a surrogate objective L(π θ ) that can be computed from samples of a rollout policy π roll . The key theoretical result is a monotonic improvement guarantee: if the surrogate improves and the policy stays within a trust region, the true objective J(π θ ) is guaranteed to improve. This guarantee relies on bounding the approximation error |J(π θ )-J(π roll ) -L(π θ )|, which depends on the divergence between π roll and π θ .

Off-Policy Mismatch in Modern LLM-RL. Recent work has shown that off-policy mismatch (π roll ̸ = π θ ) is unavoidable in modern LLM-RL systems [Liu et al., 2025, Yao et al., 2025]. Several factors contribute to this mismatch:

1. Implementation divergence: Different numerical implementations for inference (vLLM [Kwon et al., 2023], SGLang [Zheng et al., 2024]) versus training (Megatron-LM [Shoeybi et al., 2019], PyTorch FSDP [Zhao et al., 2023]) produce different logits from identical weights.

In mixture-of-experts models [Shazeer et al., 2017, DeepSeek-AI, 2024], small numerical differences can trigger different expert selections, causing discrete jumps in token probabilities.

3. Distributed staleness: Asynchronous training pipelines [Espeholt et al., 2018, Nair et al., 2015] create lag between rollout generation and gradient computation, so training occurs with updated weights π θ while rollouts were generated with stale weights π roll .

The Long-Horizon Problem.

Given that π roll ̸ = π θ is unavoidable, approximation error becomes a central concern. Classical error bounds [Kakade & Langford, 2002, Schulman et al., 2015] 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 tok,max KL = 10 -4 ), the classical bound yields an error of ≈ 1677, far exceeding any plausible improvement. This means existing theory provides no guarantee that optimization actually improves performance.

Contributions. We make the following contributions:

1. Tighter error bounds: We derive two new bounds on the approximation error (Section 3):

2. Failure of token-level methods: Both bounds depend on D tok,max KL -the maximum tokenlevel divergence across all positions in the sequence. This is inherently a sequence-level quantity, and therefore cannot be controlled by token-independent methods like PPO clipping or token masking (Section 4).

We propose Trust Region Masking, which excludes entire sequences from the gradient update if any token violates the trust region. This ensures D tok,max KL ≤ δ for all accepted sequences, providing the first non-vacuous monotonic improvement guarantees for long-horizon LLM-RL (Section 5).

We consider autoregressive language generation where a policy π θ generates a response y = (y 1 , . . . , y T ) to a prompt x. Each token y t is drawn from a vocabulary V according to:

where y <t = (y 1 , . . . , y t-1 ) denotes the tokens generated before position t. The full trajectory distribution factorizes as:

We define the context at step t as c = (x, y <t ), and the context visitation distribution under policy π as:

This is the probability of reaching context (x, y <t ) when following policy π.

Given a reward function R(x, y) ∈ [0, 1], the objective is:

The fundamental challenge is that we can only sample from a rollout policy π roll , which may differ from the training policy π θ we wish to optimize. This off-policy setting requires importance sampling or surrogate objectives.

Following Kakade & Langford [2002] and Schulman et al. [2015], we define the surrogate objective:

where A = R(x, y) -b is the trajectory advantage (with baseline b), and

is the importance ratio at token t.

The key property of this surrogate is that it matches the true gradient at the reference point (see [Kakade & Langford, 2002]):

This makes L a valid local approximation to J, but the approximation degrades as π θ moves away from π roll .

We use the following divergence measures throughout:

Following TRPO [Schulman et al., 2015], we use D KL (π roll ∥π θ )-the KL from the rollout policy to the training policy. This is the natural choice because: (1) it matches the TRPO constraint formulation, and (2) it can be computed exactly from stored rollout logits.

Definition 2.2 (Maximum and sequence-level divergences).

The chain rule (11) is an equality, which will be crucial for our improved bounds.

Key property: Pinsker’s inequality [Pinsker, 1964]. Total variation is bounded by KL divergence:

Importantly, Pinsker holds for both KL directions (since TV is symmetric), so bounds derived using D KL (π roll ∥π θ ) are equally valid as those using D KL (π θ ∥π roll ). We use the former throughout as it aligns with TRPO and is computable in practice.

We develop tighter error bounds for the surrogate objective. Define the approximation error:

This measures how well L approximates the true improvement. If we can bound |Error|, we can guarantee that improving L also improves J.

Performance Difference Identity. The error decomposes as a sum over timesteps. [Kakade & Langford, 2002]:

The error arises from evaluating g t under the wrong context distribution:

Key lemmas. Two ingredients bound this error (proofs in Section B):

The KL chain rule (11) gives

For T = 4096 and D tok,max KL = 10 -4 , this gives |Error| ≤ 1677-vacuous.

New Bound 1: Pinsker-Marginal. Our key insight is to apply Pinsker’s inequality [Pinsker, 1964, Cover & Thomas, 2006] to the marginal KL, not the per-step TV. . Apply Pinsker (12) to the marginal KL:

The advantage bound via Pinsker gives ∥g t ∥ ∞ ≤ 2 D tok,max KL /2. Combining:

For T = 4096: |Error| ≤ 35.0, which is 48× tighter than TRPO.

New Bound 2: Mixed. An alternative bounds the TV shift uniformly using the sequence-level KL.

Theorem

The Mixed bound is tighter when the sequence-level KL is small relative to the token maximum (sparse KL). For D seq KL = 0.01: |Error| ≤ 8.2, which is 200× tighter than TRPO.

Adaptive bound and monotonic improvement. The best bound is:

Define the minorizer M(π θ ) := L(π θ ) -|Error| bound . If M(π θ ) > 0, then J(π θ ) > J(π roll ) (monotonic improvement).

Scaling Value (T = 4096) ))-the worstcase token-level KL divergence across the entire sequence. This is a sequence-level quantity: no bound using only D seq KL exists, since divergence can concentrate at rare contexts (Section C). Following TRPO, we use D KL (π roll ∥π θ ), which can be computed exactly from stored rollout logits.

Our bounds depend critically on D tok,max KL -a property of the entire sequence. We now show that token-level methods cannot control this quantity.

PPO [Schulman et al., 2017] uses ratio clipping:

Gradient leakage. The min operator creates asymmetric behavior:

When ρ t ≫ 1 (e.g., from MoE routing flip) and A t < 0 (noisy reward), the gradient is unbounded. Under systemic mismatch, this injects large erroneous gradients.

A natural fix is to mask tokens with excessive divergence:

where

The theoretical problem. This does not satisfy our bounds. If token k has large divergence, masking it changes the gradient we compute, but D tok,max KL is unchanged -the divergence still exists in the sequence. The error bound remains vacuous.

Proposition 4.1. Token masking changes the optimization target but does not reduce D tok,max KL . Therefore, the monotonic improvement guarantee does not apply. The only solution. If any token violates the trust region, we must exclude the entire sequence from the gradient computation.

Since our bounds require controlling D tok,max KL -a sequence-level quantity-we must exclude entire sequences that violate the trust region from gradient computation. Define a binary mask M (x, y) = I[(x, y) ∈ Trust Region] and the masked surrogate:

The gradient is estimated as

t |c), dividing by total batch size N (not the accepted count).

Why sequence masking works. Masked sequences (M = 0) contribute zero to the gradientwe simply choose not to learn from them, which is a valid reweighting. Accepted sequences (M = 1) satisfy D tok,max KL ≤ δ by construction, so our error bounds apply. In contrast, token masking keeps the sequence but removes tokens, changing the gradient target to something that doesn’t correspond to any valid objective.

Exact computation of D KL . Following TRPO, we use D KL (π roll ∥π θ ), which can be computed exactly:

where π roll logits are stored during rollout and π θ logits are computed during training. In code: D_kl = F.kl_div(log_softmax(logits_θ), softmax(logits_roll), reduction=‘batchmean’).

Masking criterion. We use the max-based criterion M (x, y) = I[max t D KL (c t ) ≤ δ], which directly bounds D tok,max KL and is length-invariant: the threshold δ does not grow with T . In practice, one may also add an average-based criterion 1 T t D KL (c t ) ≤ δ avg to tolerate occasional outliers while maintaining robustness.

Sample-based approximation. When storing full logits is infeasible, we can use sample-based divergence estimates. The choice depends on the criterion: (1) For the max criterion, use f (ρ) = | log ρ|, which is symmetric-both ρ ≫ 1 and ρ ≪ 1 produce large values, detecting divergence in either direction. (2) For the average criterion, use f (ρ) = ρ -1 -log ρ (the k 3 estimator), which is both unbiased (E[f (ρ)] = D KL ) and non-negative (preventing cancellation). See Section D for detailed analysis. Note that each context has only one sample, so sample-based guarantees remain approximate.

Algorithm 1 Trust Region Masking (TRM)

; stored π roll logits 1: for each (x, y) ∈ D do 2:

Compute π θ logits via forward pass 3:

Set M i = I[max t D KL (c t ) ≤ δ] 5: end for 6: Compute masked gradient and update θ Theorem 5.1 (TRM Guarantee). TRM with exact KL computation and threshold δ provides:

1. Valid gradient: ∇L masked is an unbiased estimate of the masked surrogate gradient.

2. Bounded divergence: For accepted sequences, max t D KL (c t ) ≤ δ (exactly verifiable).

3. Length-invariant threshold: δ does not grow with sequence length T .

For T = 4096 and δ = 10 -4 , the error bounds are 35.0 (PM) and 8.2 (Mixed with D seq KL = 0.01)non-vacuous, unlike the classical bound of 1677.

Limitations. Our analysis assumes bounded rewards R ∈ [0, 1]. Extending to unbounded rewards requires additional assumptions on reward tail behavior. The sample-based divergence estimation is conservative; tighter bounds may be possible with distributional information.

Practical considerations. The max criterion may mask many sequences under high MoE noise. The average criterion provides a practical relaxation, trading theoretical tightness for sample efficiency. Monitoring the mask rate provides a diagnostic for training health.

Connection to other methods. Our analysis applies to any method using the surrogate objective (5), including REINFORCE [Williams, 1992] and PPO variants. The key insight is that trust region constraints must be enforced at the sequence level, not the token level.

Future work. Promising directions include: (1) soft masking with importance weights, (2) adaptive threshold selection, (3) integration with KL penalties in the reward, and (4) extensions to multi-turn agentic tasks.

Off-policy mismatch is unavoidable in modern LLM-RL due to implementation divergence, MoE routing discontinuities, and distributed staleness. Classical trust region bounds scale as O(T 2 ), becoming vacuous for long-horizon tasks. We derived two new bounds-Pinsker-Marginal at O(T 3/2 ) and Mixed at O(T )-that provide significantly tighter guarantees. Crucially, these bounds depend on the maximum token-level KL divergence across the sequence, which cannot be controlled by token-level methods like PPO clipping. We proposed Trust Region Masking (TRM), which excludes entire sequences violating the trust region from gradient computation, providing the first nonvacuous monotonic improvement guarantees for long-horizon LLM-RL.

Since R ∈ [0, 1], we have |A t | ≤ 1, so: Proof. The joint distribution factors as P π (x, y <t ) = P (x)

3.2 (Mixed bound). |Error| ≤ 2T • D tok,max KL • D seq KL Proof. Since the marginal KL is a partial sum of the sequence KL, D KL (d π roll t ∥d π θ t ) ≤ D seq KL for all t. t ∥ TV ≤ D seq KL /2 uniformly. Thus: