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

TL;DR

This paper introduces PMD-mean, an approximation method for policy mirror descent in reinforcement learning for large language models, which enhances stability and efficiency through implicit regularization.

Contribution

It proposes a novel approximation of the log-partition function in PMD, revealing its implicit regularization effect and demonstrating improved performance in math reasoning tasks.

Findings

PMD-mean achieves superior performance on math reasoning tasks.

It improves stability and time efficiency in RL for LLMs.

The method introduces an adaptive mixed KL–χ² regularizer.

Abstract

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 achieves superior performance with improved stability and time efficiency. 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-rl/OpenKimi.