Approssimazione della Funzione Log-Partizione nel Discesa a Specchio della Politica Induce una Regolarizzazione Implicita per il Post-Addestramento di LLM
Approximation of Log-Partition Function in Policy Mirror Descent Induces Implicit Regularization for LLM Post-Training
February 5, 2026
Autori: Zhenghao Xu, Qin Lu, Changlong Yu, Tuo Zhao
cs.AI
Abstract
Il Policy Mirror Descent (PMD) fornisce un quadro metodologico rigoroso per l'apprendimento per rinforzo (RL) risolvendo iterativamente sottoproblemi di miglioramento della politica regolarizzati con la divergenza KL. Sebbene questo approccio sia stato adottato nell'addestramento di LLM avanzati come Kimi K1.5/K2, gli aggiornamenti ideali in forma chiusa del PMD richiedono una stima affidabile della funzione di partizione, una sfida significativa quando si lavora con rollout limitati nei vasti spazi d'azione degli LLM. Investigiamo un algoritmo pratico, denominato PMD-mean, che approssima il termine del log-partizione con la ricompensa media sotto la politica di campionamento ed esegue una regressione nello spazio del log-policy. Nello specifico, caratterizziamo la soluzione di popolazione di PMD-mean e dimostriamo che esso ottimizza implicitamente i sottoproblemi del mirror descent con un regolarizzatore misto adattativo KL-χ². Questa regolarizzazione χ² aggiuntiva vincola ampi cambiamenti di probabilità, producendo aggiornamenti più conservativi quando le ricompense attese sono basse e migliorando la robustezza contro gli errori di stima su campioni finiti. Esperimenti su compiti di ragionamento matematico mostrano che PMD-mean raggiunge prestazioni superiori con una stabilità e un'efficienza temporale migliorate. Questi risultati approfondiscono la nostra comprensione di PMD-mean e illuminano percorsi verso miglioramenti metodologici negli algoritmi di RL per LLM. Il codice è disponibile all'indirizzo 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.