Skip to main content
Back to timeline
arXivSource publication:

A fixed KL reference can collapse the odds of optimal answers: resetting it to the latest policy restores global convergence for PPO and GRPO variants

Synopsis

Policy optimization in language-model post-training starts from a capable base model, yet REINFORCE, the basis of PPO and GRPO, can drive the probability of optimal responses toward zero even as reward improves, and convergence can take exponentially many steps even with exact gradients and KL regularization to the base model. Periodically resetting the KL reference to the most recent policy, formalized as Approximate Policy Mirror Descent (APMD), lets each subproblem be solved locally under strong regularization while successive resets accumulate into global optimality. For softmax linear policies, this yields the first global convergence and coverage-preserving guarantees for PPO and GRPO variants with reference resets and for Kimi-style squared-loss algorithms.

Source-provided article image: A Theory of Policy Optimization in Post-Training: Coverage Collapse and Reference Resets
arXiv

Interpretation

Policy optimization in post-training can raise reward while losing the responses the base model was already good at, and this coverage collapse makes the gradient signal vanish so optimization stalls on a reward plateau. Best-of-N sampling and supervised finetuning on reward-filtered samples already had sample-complexity guarantees governed by the base model's coverage, while policy optimization lacked an analogous guarantee; existing convergence analyses either assume a uniform initial policy or cannot rule out poor performance even with a high-quality base model. In a tabular softmax, single-prompt bandit construction with exact gradients, the base model has constant coverage of the optimal response, yet for any learning rate and any KL regularization strength, if REINFORCE is near-optimal from some iteration onward, some intermediate iteration has exponentially small probability on the optimal response, and expected reward rises strictly throughout that phase.

Periodically replacing the reference policy that anchors regularization with the most recent policy lets the learner make only local improvements under strong regularization, and successive resets accumulate into global optimality. With a fixed reference, regularization must be weak enough to permit substantial improvement, and weak regularization cannot prevent coverage collapse; reference resets allow strong regularization and continued improvement at the same time. The intervention is formalized as Approximate Policy Mirror Descent (APMD): each round sets the reference to the current policy, then an inner loop approximately solves a KL-regularized subproblem; provided each subproblem stays within a small KL divergence of the exact mirror descent target, the final policy retains the sublinear mirror descent rate, and on most prompts every iterate generates optimal responses at least half as often as the base model.

For softmax linear policies, which model learning over frozen pretrained features, PPO and GRPO-style variants with reference resets and Kimi-style squared-loss algorithms are all shown to be coverage-preserving. These algorithms previously lacked end-to-end global convergence guarantees, and the separate roles of regularization, clipping and objective choice had not been isolated. PPO-NoClip completes its inner optimization under strong regularization with controlled importance weights; PPO-Lite with clipping ties the required batch size to preset clipping thresholds rather than to the magnitude of policy ratios during training; KimiPO with a mean baseline, KimiPO-Norm with a normalizer baseline, and Rebel-OnPolicy, which cancels the baseline using paired responses, each satisfy the coverage-preserving condition under the corresponding regularization and sample-size choices.

Perspective

The demonstrated scope is the one-shot reasoning bandit setting with softmax linear policy classes and linear rewards, and the coverage-preserving sample and gradient-step bounds are given within that setting; the sample-complexity result for PPO-Lite requires a finite policy class and each inner loop meeting a prescribed optimization tolerance. The authors explicitly list multi-step reasoning, agentic multi-round interaction and richer policy parameterizations as future work.

The coverage-collapse lower bound uses tabular parameters and exact gradients, and the authors expect the phenomenon to persist under complex optimization landscapes and stochastic noise but do not prove it there; the clipping analysis for PPO-Lite does not cover a specific inner optimization procedure when clipping is active, because the objective is non-smooth at the clipping thresholds; KimiPO needs stronger regularization to control the bias of its mean baseline, which adds a dependence on target accuracy relative to PPO-NoClip; response coverage is only one notion of base model quality, and whether post-training can go beyond amplifying existing behaviors remains open.