Aproximação da Função Log-Partição no Descenso de Política com Espelho Induz Regularização Implícita no Pós-Treinamento de LLM
Approximation of Log-Partition Function in Policy Mirror Descent Induces Implicit Regularization for LLM Post-Training
February 5, 2026
Autores: Zhenghao Xu, Qin Lu, Changlong Yu, Tuo Zhao
cs.AI
Resumo
O Descenso de Políticas por Espelho (PMD) fornece uma estrutura fundamentada para a aprendizagem por reforço (RL) ao resolver iterativamente subproblemas de melhoria de políticas regularizados por KL. Embora esta abordagem tenha sido adotada no treino de LLMs avançados, como o Kimi K1.5/K2, as atualizações ideais de PMD de forma fechada requerem uma estimativa confiável da função de partição, um desafio significativo ao trabalhar com trajectórias limitadas nos vastos espaços de ação dos LLMs. Investigamos um algoritmo prático, denominado PMD-média, que aproxima o termo do log-partição com a recompensa média sob a política de amostragem e executa regressão no espaço log-política. Especificamente, caracterizamos a solução populacional do PMD-média e demonstramos que este otimiza implicitamente subproblemas de descenso por espelho com um regularizador misto adaptativo KL--χ^2. Esta regularização χ^2 adicional restringe grandes alterações de probabilidade, produzindo atualizações mais conservadoras quando as recompensas esperadas são baixas e aumentando a robustez contra erros de estimativa de amostra finita. Experiências em tarefas de raciocínio matemático mostram que o PMD-média alcança um desempenho superior com melhor estabilidade e eficiência temporal. Estes resultados aprofundam a nossa compreensão do PMD-média e iluminam caminhos para melhorias fundamentadas em algoritmos de RL para LLMs. O código está disponível em 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.