Theory says REINFORCE can collapse coverage in LM post-training
Original title: A Theory of Policy Optimization in Post-Training: Coverage Collapse and Reference Resets
Resetting the KL reference in Approximate Policy Mirror Descent is shown to restore global convergence for stylized PPO and GRPO.
A theory paper argues that REINFORCE, the basis of PPO and GRPO, can drive the probability of optimal responses toward zero even while reward improves, a failure called coverage collapse. The stall can require exponentially many steps despite good initial coverage, and the authors say it persists with exact gradients and KL regularization to the base model. They formalize repeated KL-reference resets as Approximate Policy Mirror Descent and claim stronger regularization plus resets can reach global optimality. For softmax linear policies over pretrained features, they give coverage-preserving guarantees for stylized PPO and GRPO with explicit KL regularization and resets, and for squared-loss variants including those used to train the Kimi models.
Unconfirmed: The global guarantees are proved for stylized algorithms and softmax linear policies, not full production training runs.
- Coverage collapse can occur even as measured reward improves.
- The claimed failure persists under KL regularization to the base model.
- APMD approximately solves a sequence of KL-regularized RL subproblems.
- Guarantees cover stylized PPO, GRPO, and some squared-loss algorithms.
From the original article
1 min read at arxiv.orgPolicy 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.
By Audrey Huang · arxiv.org
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.