Approximation de la Fonction Log-Partition dans la Descente de Miroir de Politique Induit une Régularisation Implicite pour le Post-Entraînement des LLM

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

Le « Policy Mirror Descent » (PMD) offre un cadre théorique rigoureux pour l'apprentissage par renforcement (RL) en résolvant itérativement des sous-problèmes d'amélioration de politique régularisés par la divergence KL. Bien que cette approche ait été adoptée pour entraîner des modèles de langage avancés tels que Kimi K1.5/K2, les mises à jour PMD idéales sous forme close nécessitent une estimation fiable de la fonction de partition, un défi majeur lorsqu'on travaille avec un nombre limité d'épisodes dans les vastes espaces d'actions des LLMs. Nous étudions un algorithme pratique, nommé PMD-moyenne, qui approxime le terme du log-partition par la récompense moyenne sous la politique d'échantillonnage et effectue une régression dans l'espace log-politique. Plus précisément, nous caractérisons la solution de population de PMD-moyenne et démontrons qu'elle optimise implicitement les sous-problèmes de descente miroir avec un régulariseur mixte adaptatif KL–χ². Cette régularisation χ² supplémentaire contraint les changements de probabilité importants, produisant des mises à jour plus conservatrices lorsque les récompenses attendues sont faibles et améliorant la robustesse aux erreurs d'estimation sur échantillon fini. Les expériences sur des tâches de raisonnement mathématique montrent que PMD-moyenne atteint des performances supérieures avec une stabilité et une efficacité temporelle améliorées. Ces résultats approfondissent notre compréhension de PMD-moyenne et éclairent des voies vers des améliorations principlées des algorithmes de RL pour les LLMs. Le code est disponible à l'adresse 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.