Part 3: Trust Region Optimization via Sequence Masking

Authors: Yingru Li, Jiacai Liu

Citation

@online{liu-li-2025-rl-collapse,
  title = {When Speed Kills Stability: Demystifying {RL} Collapse from the Training-Inference Mismatch},
  author = {Liu, Jiacai and Li, Yingru and Fu, Yuqian and Wang, Jiawei and Liu, Qian and Shen, Yu},
  year = {2025},
  month = sep,
  url = {https://richardli.xyz/rl-collapse}
}

Recap: The Trust Region Framework

In Part 1, we established the theoretical foundation for trust region optimization. Here we recap the key results that motivate the methods in this part.

The Surrogate Objective and Its Limitations

When optimizing a policy $\pi$ using samples from a behavior policy $\mu$, we cannot directly optimize the true objective $J(\pi)$. Instead, we optimize a surrogate objective:

where $d_\mu$ is the state visitation distribution under $\mu$, and $A_\mu$ is the advantage function.

The surrogate is a first-order approximation that satisfies:

• $L_\mu(\mu) = J(\mu)$ (equal values at $\pi = \mu$)
• $\nabla L_\mu(\pi)\big|_{\pi=\mu} = \nabla J(\pi)\big|_{\pi=\mu}$ (equal gradients at $\pi = \mu$)

However, the approximation degrades as $\pi$ moves away from $\mu$.

The TRPO Lower Bound

The Performance Difference Lemma quantifies the gap between the surrogate and true objectives:

The key difference from the surrogate is that the true improvement uses the state distribution $d_\pi$ (under the new policy), while the surrogate uses $d_\mu$ (under the old policy).

Using the Simulation Lemma, which bounds how state distributions diverge over time:

we derive the TRPO lower bound:

where:

• $C$ is a constant depending on the maximum advantage $\max_{s,a}|A_\mu(s,a)|$
• $T$ is the horizon (sequence length)
• $D_{TV}^{\max} = \max_s D_{TV}(\pi(\cdot|s) | \mu(\cdot|s))$ is the maximum per-token TV distance

The Trust Region Requirement

For surrogate optimization to guarantee improvement in the true objective, the policy must stay within a trust region:

The critical insight is that the valid trust region size must shrink with sequence length:

This $T^2$ dependence arises because state distribution errors accumulate linearly over $T$ steps, and the total error (summed over all steps) scales quadratically.

Soft vs. Hard Trust Regions

The TRPO framework suggests two approaches to enforce trust regions:

Soft trust regions use clipped importance sampling: $\min(\rho_t, 1+\epsilon)$. This is computationally efficient but retains potentially problematic samples.

This part develops hard trust region methods—Seq-MIS and Geo-RS—that completely exclude samples outside the trusted region. We show when and why hard rejection outperforms soft clipping.

1. OOD High-Weight Samples: Why Rejection Outperforms Clipping

1.1 The Problem: Clipping Retains Problematic Samples

In Part 2, we derived the Seq-TIS estimator:

The implicit assumption was that all samples $y \sim \mu$ are valid learning signals—samples with high weights $\rho(y) = \pi(y)/\mu(y)$ simply require variance control via clipping.

However, this assumption fails in practice. Consider a sample with $\rho(y) = 10,000$. This means:

Such extreme ratios typically arise when $\mu(y)$ is near the numerical precision floor (e.g., $\mu(y) \approx 10^{-9}$). These are Out-of-Distribution (OOD) samples—samples generated by policies outside the trust region. They occur due to:

1. Numerical precision artifacts in probability computation
2. Distribution shift between $\pi$ and $\mu$ beyond the valid IS regime

The Clipping Problem: When we apply Seq-TIS, we compute $\min(10000, C) \cdot f(y) = C \cdot f(y)$. The sample is still included in the gradient update with weight $C$. If $f(y)$ is malformed (due to the OOD nature of $y$), we introduce a systematic error into every gradient step.

1.2 The Solution: Hard Trust Region via Rejection (Seq-MIS)

Instead of soft clipping, we enforce a Hard Trust Region: samples outside the trusted region are rejected entirely.

Mathematical Interpretation: The trust region is defined as:

Only samples within this region contribute to the gradient. Samples with $\rho(y) \gt C$ are treated as unreliable and excluded.

Connection to TRPO: This implements the trust region concept from TRPO theory (Part 1), where the trust region constraint $D_{TV}(\pi | \mu) \le \delta$ must be enforced to guarantee improvement. Seq-MIS enforces this constraint via rejection (hard trust region) rather than penalization (soft trust region), ensuring that gradient updates only use samples from the valid trust region.

1.3 Bias-Variance Analysis of Seq-MIS

Bias: The Seq-MIS estimator is biased. By importance sampling identity:

The true gradient is $g = \mathbb{E}_\pi[f]$. Therefore:

The bias is the negative of the contribution from rejected samples (weighted by $\pi$).

Variance: Since $\mathbb{I}(\rho \le C) \cdot \rho \le C$ when the indicator is 1, and 0 otherwise:

1.4 When to Use Each Estimator

The choice depends on the reliability of high-weight samples. When the behavior policy $\mu$ is well-calibrated and mismatch is controlled, Seq-TIS extracts more information. When OOD samples are prevalent, Seq-MIS provides a Hard Trust Region that prevents gradient corruption.

2. Length-Dependent Rejection Bias: The Failure of Sequence-Level IS for Long Horizons

2.1 The Problem: Exponential Growth of the Importance Ratio

For autoregressive generation, the sequence-level importance ratio is a product of per-token ratios:

Even when the per-token mismatch is small, this product grows (or shrinks) exponentially with sequence length $T$.

Formal Analysis: Let $\bar{\rho} = \mathbb{E}[\rho_t]$ denote the expected per-token ratio. If we assume (for simplicity) that the $\rho_t$ are independent with $\bar{\rho} \gt 1$, then:

Similarly, for the log-ratio:

If each $\log \rho_t$ has mean $\delta \gt 0$ (i.e., $\pi$ assigns slightly higher probability than $\mu$ on average), then:

Numerical Example: Consider $\bar{\rho} = 1.001$ (0.1% per-token drift):

2.2 The Consequence: Systematic Length Bias

This exponential scaling creates a length-dependent acceptance probability. For any fixed threshold $C$, there exists a critical length $T^*$ beyond which almost all samples are rejected:

For $C = 5$ and $\bar{\rho} = 1.001$: $T^* = \frac{\log 5}{\log 1.001} \approx 1609$ tokens.

The Problem for Reasoning Models: Chain-of-Thought (CoT) models and agents often generate sequences of 2,000-10,000+ tokens. With standard Seq-TIS or Seq-MIS:

1. Short responses (< 1000 tokens) are almost always accepted
2. Long reasoning chains (> 2000 tokens) are almost always rejected or heavily clipped
3. The model receives systematically biased feedback favoring short outputs

This is not a variance problem—it is a structural bias against long-horizon reasoning, independent of the quality of individual reasoning steps.

3. Geometric Rejection Sampling: A Per-Token Trust Region

3.1 From Extensive to Intensive Metrics

The fundamental problem with sequence-level IS is that $\rho(y) = \prod_t \rho_t$ is an extensive quantity—it scales with sequence length. We need an intensive quantity that measures the average per-token divergence, independent of length.

This is the geometric mean of the per-token ratios. It is length-invariant: if every $\rho_t = r$, then $\rho_{\text{geo}} = r$ regardless of $T$.

3.2 Mathematical Foundation: Connection to Per-Token KL Divergence

The geometric mean has a natural interpretation in terms of KL divergence. Taking the logarithm:

This is the sample average of the per-token log-ratios along trajectory $y$.

Connection to KL Divergence: Let $s_t = (x, y_{\lt t})$ denote the state (context) at step $t$. Recall that:

• Forward KL: $D_{KL}(\pi \| \mu) = \mathbb{E}_{y_t \sim \pi}\left[\log \frac{\pi(y_t|s_t)}{\mu(y_t|s_t)}\right]$
• Reverse KL: $D_{KL}(\mu \| \pi) = \mathbb{E}_{y_t \sim \mu}\left[\log \frac{\mu(y_t|s_t)}{\pi(y_t|s_t)}\right] = -\mathbb{E}_{y_t \sim \mu}\left[\log \frac{\pi(y_t|s_t)}{\mu(y_t|s_t)}\right]$

Since samples are drawn from $\mu$ (the behavior policy), each term $\log \rho_t = \log \frac{\pi(y_t|s_t)}{\mu(y_t|s_t)}$ is a single-sample estimate of the negative reverse KL:

Therefore, $\log \rho_{\text{geo}}(y)$ can be interpreted as a trajectory-averaged sample of the negative reverse KL:

Interpretation: $\log \rho_{\text{geo}}(y)$ measures the average per-step log-likelihood ratio along the specific trajectory $y$. Unlike the sequence-level ratio, this quantity:

1. Does not grow with $T$ (it’s an average, not a sum)
2. Can be positive or negative ($\log \rho_{\text{geo}} \gt 0$ when $\pi$ assigns higher probability, $\lt 0$ when $\mu$ assigns higher probability)
3. Detects both directions of drift ($\rho_{\text{geo}} \ll 1$: policy forgetting; $\rho_{\text{geo}} \gg 1$: policy collapse)

3.3 The Two-Sided Hard Trust Region (Geo-RS)

With the geometric ratio, we can define a Per-Token Trust Region that is independent of sequence length:

Why Two-Sided? The trust region enforces constraints in both directions:

Typical Values: $C_{\text{low}} = 0.5$, $C_{\text{high}} = 2.0$ (or equivalently, $|\log \rho_{\text{geo}}| \le \log 2 \approx 0.69$).

3.4 Bias-Variance Analysis of Geo-RS

Bias: Geo-RS is biased because it rejects samples:

Unlike importance sampling, Geo-RS does not reweight by $\rho$—it only filters. This introduces bias but eliminates the variance explosion from high weights.

Variance: Since there is no importance weight multiplication, the variance is bounded by:

Length Invariance: The acceptance criterion $C_{\text{low}} \le \rho_{\text{geo}} \le C_{\text{high}}$ is length-independent. A 100-token sequence and a 10,000-token sequence are judged by the same per-token divergence threshold.

3.5 Combining Geo-RS with Seq-TIS

In practice, we may want to combine both mechanisms:

1. Geo-RS filter: Reject samples with extreme per-token divergence (length-invariant safety)
2. Seq-TIS weighting: Apply importance sampling with clipping for accepted samples (bias correction)

This estimator:

1. First checks if the sample is within the per-token trust region (Geo-RS)
2. Then applies clipped importance weighting for accepted samples (Seq-TIS)

When to use which component:

4. Summary: Hierarchy of Estimators and Selection Guidelines

We have developed a hierarchy of estimators, each addressing specific failure modes:

For long-horizon reasoning tasks, Geometric Rejection Sampling (Geo-RS) provides a principled, length-invariant Hard Trust Region that prevents the systematic length bias inherent in standard importance sampling estimators.