Policy optimization struggles and fixes for language model post-training
A Theory of Policy Optimization in Post-Training: Coverage Collapse and Reference Resets
Machine Learning
Summary
Training language models to improve output through policy optimization can run into a problem called coverage collapse, where chances of generating good answers drop even when rewards improve. The authors found this happens with popular methods like REINFORCE and PPO, and it makes learning very slow. They propose a way called Approximate Policy Mirror Descent that resets a reference point during training, which avoids this collapse and helps models learn better responses. This work clarifies how adjusting training methods and regularization can improve the efficiency of language model post-training.
What this means in practice
- •For language model developers: Improve post-training methods to avoid coverage collapse and achieve more reliable learning of optimal responses.
- •For reinforcement learning engineers: Design algorithms with reference resets to ensure convergence when optimizing policies over pretrained model features.
A theory result. No direct application yet.
Authors
Audrey Huang, Nan Jiang, Sham M. Kakade, Sadhika Malladi, Akshay Krishnamurthy, Dylan J. Foster
Abstract
Policy optimization is central to language model post-training, yet principled guidance remains limited. We revisit its theoretical foundations to identify algorithms that learn near-optimal policies using total samples and gradient steps governed by the base model's response coverage, a measure of how readily it generates optimal responses. We call such algorithms coverage-preserving. This guarantee captures the goal of amplifying capabilities accessible through sampling from the base model. (1) Coverage collapse. REINFORCE, the basis for PPO and GRPO, can drive the probability of optimal responses to near zero, even as reward improves. This coverage collapse stalls optimization, and convergence can require exponentially many steps despite good initial coverage. The failure arises even with exact gradients and persists under KL regularization to the base model. (2) Reference resets. We show that repeatedly resetting the KL reference to the most recent policy mitigates coverage collapse. We formalize this intervention in Approximate Policy Mirror Descent (APMD), a meta-algorithm that approximately solves a sequence of KL-regularized RL subproblems. APMD enables stronger regularization, making each subproblem easier to solve, while reference resets allow convergence to global optimality. (3) Coverage-preserving algorithms. For softmax linear policies over pretrained features, we use APMD to establish the first global convergence guarantees showing that widely used families of post-training algorithms are coverage-preserving. These include stylized PPO and GRPO variants with explicit KL regularization and reference resets, and squared-loss algorithms that implement them implicitly, including those used to train the Kimi models. These results clarify how regularization, clipping, and objective choice affect the sampling and optimization costs of post-training.