Aproximación de la Función Log-Partición en el Descenso de Espejo de Política Induce una Regularización Implícita para el Post-Entrenamiento de Modelos de Lenguaje Grandes
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
Resumen
El descenso de políticas por espejo (PMD) proporciona un marco fundamentado para el aprendizaje por refuerzo (RL) mediante la resolución iterativa de subproblemas de mejora de políticas regularizados por la divergencia KL. Si bien este enfoque ha sido adoptado en el entrenamiento de modelos de lenguaje avanzados como Kimi K1.5/K2, las actualizaciones ideales de forma cerrada de PMD requieren una estimación fiable de la función de partición, un desafío significativo al trabajar con trayectorias limitadas en los vastos espacios de acción de los LLM. Investigamos un algoritmo práctico, denominado PMD-media, que aproxima el término del log-partition con la recompensa media bajo la política de muestreo y realiza regresión en el espacio logarítmico de políticas. Específicamente, caracterizamos la solución poblacional de PMD-media y demostramos que optimiza implícitamente subproblemas de descenso por espejo con un regularizador mixto adaptativo KL-χ². Esta regularización χ² adicional restringe los grandes cambios de probabilidad, produciendo actualizaciones más conservadoras cuando las recompensas esperadas son bajas y mejorando la robustez frente a errores de estimación con muestras finitas. Los experimentos en tareas de razonamiento matemático muestran que PMD-media logra un rendimiento superior con una estabilidad y eficiencia temporal mejoradas. Estos hallazgos profundizan nuestra comprensión de PMD-media e iluminan vías hacia mejoras fundamentadas en los algoritmos de RL para LLMs. El código está disponible en 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.