A recent paper published on arXiv, "A Theory of Policy Optimization in Post-Training: Coverage Collapse and Reference Resets," details a critical limitation in how language models learn after their initial training. The authors, including researchers from the Kempner Institute and UIUC, identify a problem they term "coverage collapse," which affects algorithms like REINFORCE, a foundational component of widely used methods such as Proximal Policy Optimization (PPO) and Group Relative Policy Optimization (GRPO).
Coverage collapse occurs when the probability of a language model producing optimal responses falls to near zero, even as the overall reward signal improves. This can dramatically impede the optimization process, potentially requiring an exponentially larger number of steps to converge to a solution, despite the model initially having a good "coverage" of optimal responses. The researchers note that this issue persists even with exact gradients and when Kullback-Leibler (KL) regularization is applied to the base model.
The paper defines "response coverage" as a measure of how readily a base model generates optimal responses. Algorithms that are "coverage-preserving" aim to amplify capabilities already present in the base model through sampling. The research suggests that REINFORCE, and by extension algorithms based on it, can inadvertently drive down the likelihood of producing these desired responses, creating a bottleneck in learning.
To address coverage collapse, the researchers propose "reference resets." This involves repeatedly resetting the KL reference to the most recent policy during training. They formalize this intervention within Approximate Policy Mirror Descent (APMD), a meta-algorithm designed to solve a sequence of KL-regularized reinforcement learning subproblems. APMD with reference resets allows for stronger regularization, simplifying each subproblem, while still enabling convergence to global optimality.
The study provides the first global convergence guarantees for coverage-preserving algorithms, specifically for softmax linear policies over pretrained features. These include variants of PPO and GRPO that incorporate explicit KL regularization and reference resets. The authors also note that squared-loss algorithms, including those used to train the Kimi models, implicitly implement these mechanisms.
Policy optimization is a central element in language model post-training, which encompasses various techniques like Reinforcement Learning from Human Feedback (RLHF), Supervised Fine-Tuning (SFT), and Direct Preference Optimization (DPO). Post-training aims to refine a language model's behavior, enabling it to follow instructions, solve problems, and produce preferred outputs effectively.
PPO and GRPO are prominent policy optimization methods. PPO, introduced by OpenAI in 2017, uses generalized advantage estimation and an actor-critic interplay to stabilize updates. GRPO, often used in reasoning-oriented tasks, eliminates the need for a separate value function by comparing multiple generated outputs and using their average reward as a baseline. The findings from this new theoretical work clarify how elements such as regularization, clipping, and the choice of objective function influence the sampling and optimization costs associated with post-training.