Approximation der Log-Partitionsfunktion in Policy Mirror Descent induziert implizite Regularisierung für das Post-Training von LLMs

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

February 5, 2026

papers.authors: Zhenghao Xu, Qin Lu, Changlong Yu, Tuo Zhao

cs.AI

papers.abstract

Policy Mirror Descent (PMD) bietet einen prinzipiellen Rahmen für bestärkendes Lernen (Reinforcement Learning, RL), indem es KL-regularisierte Policy-Improvement-Teilprobleme iterativ löst. Obwohl dieser Ansatz beim Training fortgeschrittener LLMs wie Kimi K1.5/K2 Anwendung findet, erfordern die idealen geschlossenen PMD-Updates eine zuverlässige Schätzung der Partitionfunktion – eine erhebliche Herausforderung bei begrenzten Rollouts in den riesigen Aktionsräumen von LLMs. Wir untersuchen einen praktischen Algorithmus, genannt PMD-Mean, der den Log-Partition-Term mit dem mittleren Reward unter der Sampling-Policy approximiert und Regression im Log-Policy-Raum durchführt. Konkret charakterisieren wir die Populationslösung von PMD-Mean und zeigen, dass diese implizit Mirror-Descent-Teilprobleme mit einem adaptiven gemischten KL-χ²-Regularisierer optimiert. Diese zusätzliche χ²-Regularisierung beschränkt große Wahrscheinlichkeitsänderungen, erzeugt konservativere Updates bei niedrigen erwarteten Rewards und erhöht die Robustheit gegenüber Schätzfehlern mit begrenzten Stichproben. Experimente zu mathematischen Reasoning-Aufgaben zeigen, dass PMD-Mean eine überlegene Leistung mit verbesserter Stabilität und ZeitEffizienz erzielt. Diese Ergebnisse vertiefen unser Verständnis von PMD-Mean und weisen Wege zu prinzipiellen Verbesserungen von RL-Algorithmen für LLMs auf. Code ist verfügbar unter https://github.com/horizon-rl/OpenKimi.

English

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--χ^2 regularizer. This additional χ^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.