Benadering van de Log-Partitiefunctie in Policy Mirror Descent Induceert Impliciete Regularisatie voor Post-Training van Taalmodellen
Approximation of Log-Partition Function in Policy Mirror Descent Induces Implicit Regularization for LLM Post-Training
February 5, 2026
Auteurs: Zhenghao Xu, Qin Lu, Changlong Yu, Tuo Zhao
cs.AI
Samenvatting
Policy Mirror Descent (PMD) biedt een principieel raamwerk voor reinforcement learning (RL) door iteratief KL-geregulariseerde beleidsverbeteringssubproblemen op te lossen. Hoewel deze aanpak is overgenomen bij het trainen van geavanceerde LLM's zoals Kimi K1.5/K2, vereisen de ideale gesloten-vorm PMD-updates een betrouwbare schatting van de partitiefunctie, een aanzienlijke uitdaging bij het werken met beperkte rollouts in de uitgestrekte actieruimten van LLM's. Wij onderzoeken een praktisch algoritme, genaamd PMD-mean, dat de log-partitieterm benadert met de gemiddelde beloning onder het bemonsteringsbeleid en regressie uitvoert in de log-beleidsruimte. Specifiek karakteriseren we de populatie-oplossing van PMD-mean en tonen aan dat het impliciet spiegelafstammingssubproblemen optimaliseert met een adaptieve gemengde KL-χ²-regularisator. Deze extra χ²-regularisatie beperkt grote kansveranderingen, wat conservatievere updates produceert wanneer de verwachte beloningen laag zijn en de robuustheid tegen schattingsfouten met eindige steekproeven verbetert. Experimenten met wiskundige redeneertaken tonen aan dat PMD-mean superieure prestaties bereikt met verbeterde stabiliteit en tijdsefficiëntie. Deze bevindingen verdiepen ons begrip van PMD-mean en belichten wegen naar principiële verbeteringen in RL-algoritmen voor LLM's. Code is beschikbaar op 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.