A Theory of Policy Optimization in Post-Training: Coverage Collapse and Reference Resets

Paper: 2611.14467
Authors: Audrey Huang, Nan Jiang, Sham M. Kakade, Sadhika Malladi, Akshay Krishnamurthy, Dylan J. Foster (Kempner Institute at Harvard University, UIUC, UCSD, Microsoft Research NYC)
Categories: cs.LG, stat.ML

The Gap

Reinforcement Learning (RL) post-training has driven major breakthroughs in frontier reasoning models (e.g., OpenAI o-series, DeepSeek-R1, Kimi k1.5). In post-training, an already capable pretrained base model $\pi_{\text{base}}$ is adapted using reward feedback. The core promise is that the base model already possesses high response coverage—it can generate optimal answers with non-trivial probability—and RL should amplify this capability.

However, despite massive engineering investments in algorithms like PPO (Proximal Policy Optimization) and GRPO (Group Relative Policy Optimization), the field suffers from a lack of foundational theory:

1. The Coverage Puzzle: In practice, policy optimization frequently encounters prolonged training plateaus and catastrophic capability collapse. Why do algorithms that optimize expected reward stall, even when the base model starts with favorable coverage of correct answers?
2. The Inefficacy of Static KL Regularization: Practitioners routinely add a Kullback-Leibler (KL) penalty against the base model $\pi_{\text{base}}$: $\mathcal{J}_\eta(\pi_\theta) := \mathcal{J}(\pi_\theta) - \frac{1}{\eta} D_{\text{KL}}(\pi_\theta \parallel \pi_{\text{base}})$ If $\eta$ is large (weak regularization), the policy quickly drifts and destabilizes. If $\eta$ is small (strong regularization), the policy is pinned to the base model and fails to reach optimality.
3. The Disconnect in Algorithmic Recipes: Why do leading open recipes disagree on fundamental choices—such as clipping versus KL regularization, policy gradient versus squared-loss objectives, or reference resets?

                    THE COVERAGE COLLAPSE PHENOMENON

  Pretrained Base Model π_base: Has good response coverage (C_base ≤ 6)
  
  [Standard REINFORCE / PPO / GRPO with fixed base reference π_base]:
  Iterate t=0 ──> Iterate t=100 ──> Iterate t=1,000 ...
  Expected Reward J(π_t) rises monotonically: 0.2 ──> 0.5 ──> 0.8
  
  BUT Probability of Optimal Response π_t(Y^*_x | x) crashes!
  Coverage Metric: C_{π_t} explodes from 6 to exp(Ω(1 / √ε))!
  
  Result: Optimal solutions become exponentially rare to sample.
          Gradient signals vanish into an exponential plateau.
          Optimization stalls completely despite good initial coverage!

The Increment

One sentence: The authors uncover coverage collapse—proving that standard policy gradients drive optimal response probabilities to near zero even under exact gradients and static KL penalties—and prove that periodically resetting the KL reference to the active policy (formalized as Approximate Policy Mirror Descent / APMD) guarantees coverage preservation and polynomial convergence, theoretically validating the design of models like Kimi k1.5.

                  APMD: MITIGATING COVERAGE COLLAPSE VIA REFERENCE RESETS

  Outer Loop t = 0, 1, 2, ... , T:
  
  1. REFERENCE RESET:
     Anchor KL penalty to CURRENT policy π_t (NOT static π_base!):
     Target: π^*_{η,t}(y|x) ∝ π_t(y|x) · exp(η · r(x, y))
  
  2. INNER SOLVE:
     Solve KL-regularized subproblem: D_KL(π_{t+1} || π^*_{η,t}) ≤ ε_in
     ├── Enables STRONG local regularization (large 1/η) -> Subproblem is easy to solve!
     └── Resets allow GLOBAL progression -> Sublinear mirror descent convergence!
  
  Guaranteed Outcome (Theorem 3.1):
  For all iterates t: π_t(Y^*_x | x) ≥ 0.5 · π_base(Y^*_x | x)  [Coverage is Preserved!]
  Final Suboptimality: J^* - J(π_T) ≤ log(2 C_base) / (η T)

Core Framework & Mechanisms

1. Mathematical Formalization of Response Coverage

Let prompt $x \sim \mu$, and $Y^*_x := \{ y \in \mathcal{Y} : r(x, y) = \max_{y'} r(x, y') \}$ be the set of optimal responses. The response coverage of a policy $\pi$ is: $\mathcal{C}_\pi := \mathbb{E}_{x \sim \mu} \left[ \frac{1}{\pi(Y^*_x | x)} \right]$ Small $\mathcal{C}_\pi$ means fewer samples are required to observe an optimal response (e.g., Best-of-$N$ requires $N = \widetilde{\mathcal{O}}(\mathcal{C}_{\text{base}} / \varepsilon)$ samples).

2. Theorem 2.1: The Lower Bound on Coverage Collapse

The authors prove that standard policy gradients inherently destroy coverage:

Theorem 2.1 (Coverage Collapse in REINFORCE): Fix $\varepsilon \in (0, 1/64]$. There exists a bandit instance with base model $\mathcal{C}_{\text{base}} \le 6$ such that for any learning rate $\alpha \le 1/5$ and any static regularization parameter $\eta \in (0, \infty]$, REINFORCE produces an intermediate iterate $t < T$ where: $\mathcal{C}_{\pi_t} = \exp\left( \Theta\left(\frac{1}{\sqrt{\varepsilon}}\right) \right), \quad \text{requiring } T = \Omega\left( \frac{\mathcal{C}_{\pi_t}}{\alpha \sqrt{\varepsilon}} \right) \text{ steps to converge.}$

Because the RL objective is non-concave for softmax policies, the Łojasiewicz gradient inequality: $J^* - J(\pi_t) \le c \cdot \| \pi^* / \pi_t \|_\infty \cdot \| \nabla \mathcal{J}(\pi_t) \|$ is throttled by $\| \pi^* / \pi_t \|_\infty$. When optimal actions are prematurely deprioritized, $\| \nabla \mathcal{J}(\pi_t) \| \to 0$ while the suboptimality gap remains massive, freezing the policy on an exponential plateau.

3. Approximate Policy Mirror Descent (APMD)

To prevent collapse, APMD updates the policy via sequential KL-regularized proximal targets: $\pi^*_{t,\eta} := \arg\max_{\pi} \mathbb{E}_\pi [r(x, y)] - \frac{1}{\eta} D_{\text{KL}}(\pi \parallel \pi_t) \implies \pi^*_{t,\eta}(y|x) \propto \pi_t(y|x) e^{\eta r(x, y)}$

Key Theoretical Guarantees:

• Coverage Preservation: If each inner subproblem achieves error $D_{\text{KL}}(\pi_{t+1} \parallel \pi^*_{t,\eta}) \le \varepsilon_{\text{in}} := \frac{1}{64 \mathcal{C}_{\text{base}} T^2 \log(eT)}$, then with high probability: $\pi_t(Y^*_x | x) \ge \frac{1}{2} \pi_{\text{base}}(Y^*_x | x), \quad \forall 0 \le t \le T$
• Global Sublinear Convergence: The suboptimality gap shrinks monotonically as: $J^* - J(\pi_T) \le \frac{\log(2 \mathcal{C}_{\text{base}})}{\eta T}$

                           TAXONOMY OF POST-TRAINING ALGORITHMS

  Algorithm Family         Objective / Loss              Reference Scheme   Coverage Preserving?
  ─────────────────────────────────────────────────────────────────────────────────────────────
  Naïve REINFORCE          Policy gradient on J(π)       None               NO (Collapses exp.)
  Standard PPO / GRPO      Clipped ratio to π_base       Static π_base      NO (Plateaus)
  APMD-PPO / APMD-GRPO     KL-regularized proximal       Reset to π_t       YES (Theorem 3.1)
  Squared-Loss (Kimi k1.5) Implicit mirror descent       Implicit Reset     YES (Theorem 4.1)

Empirical Validation & Algorithmic Implications

1. Validating Coverage Dynamics

In controlled experiments comparing REINFORCE against APMD on instances where $\mathcal{C}_{\text{base}} = 6$:

• In REINFORCE with static base KL, coverage $\mathcal{C}_{\pi_t}$ explodes from $10^1$ to over $10^4$ by step 1,500, stalling reward improvement in a prolonged plateau lasting over 8,000 gradient steps.
• In APMD (with reference resets every $K=100$ steps), coverage $\mathcal{C}_{\pi_t}$ remains bounded below $10^1$ throughout training, smoothly ascending to $J^* = 1.0$ within 1,200 steps (an $8\times$ optimization speedup).

2. Theoretical Grounding for Kimi k1.5 Squared-Loss RL

A major contribution of Section 4 is establishing formal equivalence between APMD and the squared-loss policy optimization recently adopted in frontier reasoning models like Moonshot AI’s Kimi k1.5: $\mathcal{L}_{\text{square}}(\theta) := \mathbb{E}_{x, y \sim \pi_t} \left[ \left( \log \pi_\theta(y|x) - \log \pi_t(y|x) - \eta A^{\pi_t}(x, y) \right)^2 \right]$ The authors prove that minimizing this regression loss implicitly executes an APMD step with reference reset to $\pi_t$. Consequently, Kimi’s squared-loss formulation is fundamentally coverage-preserving, explaining why it avoids the policy collapse that plagues standard PPO.

Practical Implications

1. Stop Anchoring KL Regularization to $\pi_{\text{base}}$: Static base-model KL penalties do not prevent policy collapse. Modern RL post-training infrastructure should anchor KL regularization to the current checkpoint $\pi_t$ (or update the reference periodically in an outer loop).
2. Aggressive Regularization Becomes Viable: With reference resets, engineers can safely crank up the regularization penalty $\frac{1}{\eta}$ without fearing under-optimization, making each inner gradient step far more stable and variance-reduced.
3. Formal Justification for Regression-Based RL: The mathematical convergence guarantees confirm that squared-loss objectives (regressing log-probabilities onto advantage targets) are not mere empirical heuristics, but theoretically grounded mirror descent implementations that preserve base model coverage.