Approximation of Log-Partition Function in Policy Mirror Descent Induces Implicit Regularization for LLM Post-Training

Policy mirror descent (PMD) provides a principled framework for reinforcement learning (RL) by iteratively solving KL-regularized policy improvement subproblems. While this approach has been adopted in training advanced LLMs such as Kimi K1.5/K2, the ideal closed-form PMD updates require reliable partition function estimation, a significant challenge when working with limited rollouts in the vast action spaces of LLMs. We investigate a practical algorithm, termed PMD-mean, that approximates the log-partition term with the mean reward under the sampling policy and performs regression in log-policy space. Specifically, we characterize the population solution of PMD-mean and demonstrate that it implicitly optimizes mirror descent subproblems with an adaptive mixed KL--$\chi^2$ regularizer. This additional $\chi^2$ regularization constrains large probability changes, producing more conservative updates when expected rewards are low and enhancing robustness against finite-sample estimation errors. Experiments on math reasoning tasks show that PMD-mean improves stability and remains competitive in the stale-rollout settings while enabling higher throughput through larger rollout batches. These findings deepen our understanding of PMD-mean and illuminate pathways toward principled improvements in RL algorithms for LLMs. Code is available at https://github.com/horizon-llm/OpenKimi.

Large language models can improve their reasoning ability by practicing on problems and learning from whether their answers are correct. However, generating many answers during training is expensive, so practical systems often reuse answers produced by an older version of the model. This can make training faster, but it can also make learning unstable. This paper studies a simple training method inspired by recent large-scale reasoning systems. We explain mathematically why this method tends to make more cautious updates, especially when the model is still solving only a small fraction of problems correctly. This caution helps prevent the model from changing too aggressively based on limited or outdated training samples. We also test the method on math-reasoning tasks and find that it remains stable and competitive while allowing more efficient training with reused model outputs.