# Approximation of Log-Partition Function in Policy Mirror Descent Induces Implicit Regularization for LLM Post-Training

arXiv: https://arxiv.org/abs/2602.05933 · HTML: https://mlsanity.com/paper/2602.05933/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 35 of its 42 references that are in the map.
- First posted on arXiv 5 Feb 2026.

## Builds on (most cited first, 30 of 35)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [Gemini 2.5: Pushing the Frontier with Advanced Reasoning, Multimodality, Long Context, and Next Generation Agentic Capabilities](https://mlsanity.com/paper/2507.06261/) (2025)
- [Constitutional AI: Harmlessness from AI Feedback](https://mlsanity.com/paper/2212.08073/) (2022)
- [DAPO: An Open-Source LLM Reinforcement Learning System at Scale](https://mlsanity.com/paper/2503.14476/) (2025)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [Kimi k1.5: Scaling Reinforcement Learning with LLMs](https://mlsanity.com/paper/2501.12599/) (2025)
- [Tulu 3: Pushing Frontiers in Open Language Model Post-Training](https://mlsanity.com/paper/2411.15124/) (2024)
- [Group Sequence Policy Optimization](https://mlsanity.com/paper/2507.18071/) (2025)
- [Kimi K2: Open Agentic Intelligence](https://mlsanity.com/paper/2507.20534/) (2025)
- [A Theory of Regularized Markov Decision Processes](https://mlsanity.com/paper/1901.11275/) (2019)
- [MiniMax-M1: Scaling Test-Time Compute Efficiently with Lightning Attention](https://mlsanity.com/paper/2506.13585/) (2025)
- [AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning](https://mlsanity.com/paper/2505.24298/) (2025)
- [Bridging the Gap Between Value and Policy Based Reinforcement Learning](https://mlsanity.com/paper/1702.08892/) (2017)
- [Beyond Reverse KL: Generalizing Direct Preference Optimization with Diverse Divergence Constraints](https://mlsanity.com/paper/2309.16240/) (2023)
- [Policy Mirror Descent for Reinforcement Learning: Linear Convergence, New Sampling Complexity, and Generalized Problem Classes](https://mlsanity.com/paper/2102.00135/) (2021)
- [On the Convergence Rates of Policy Gradient Methods](https://mlsanity.com/paper/2201.07443/) (2022)
- [Mirror Descent Policy Optimization](https://mlsanity.com/paper/2005.09814/) (2020)
- [Policy Mirror Descent for Regularized Reinforcement Learning: A Generalized Framework with Linear Convergence](https://mlsanity.com/paper/2105.11066/) (2021)
- [Correcting the Mythos of KL-Regularization: Direct Alignment without Overoptimization via Chi-Squared Preference Optimization](https://mlsanity.com/paper/2407.13399/) (2024)
- [Asynchronous RLHF: Faster and More Efficient Off-Policy RL for Language Models](https://mlsanity.com/paper/2410.18252/) (2024)
- [REBEL: Reinforcement Learning via Regressing Relative Rewards](https://mlsanity.com/paper/2404.16767/) (2024)
- [Offline Regularised Reinforcement Learning for Large Language Models Alignment](https://mlsanity.com/paper/2405.19107/) (2024)

## Abstract (arXiv)

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--chi2 regularizer. This additional chi2 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.