mlsanityData to 7 Oct 2026

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

arXiv 2602.05933 ↗First posted 5 Feb 2026cs.LGcited by 0

Not on the map. Why?

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.

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

Builds on

35 of 42 references.

References · 35

  1. 2022Training language models to follow instructions with human feedbackcited by 15,712
  2. 2017Proximal Policy Optimization Algorithmscited by 14,441
  3. 2024DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Modelscited by 8,851
  4. 2025Qwen3 Technical Reportcited by 8,195
  5. 2023Direct Preference Optimization: Your Language Model is Secretly a Reward Modelcited by 7,916
  6. 2025DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learningcited by 6,146
  7. 2024Qwen2.5 Technical Reportcited by 4,713
  8. 2025Gemini 2.5: Pushing the Frontier with Advanced Reasoning, Multimodality, Long Context, and Next Generation Agentic Capabilitiescited by 3,655
  9. 2022Constitutional AI: Harmlessness from AI Feedbackcited by 2,950
  10. 2025DAPO: An Open-Source LLM Reinforcement Learning System at Scalecited by 2,678
  11. 2015Trust Region Policy Optimizationcited by 2,503
  12. 2019Fine-Tuning Language Models from Human Preferencescited by 2,124
23 more
  1. 2024HybridFlow: A Flexible and Efficient RLHF Frameworkcited by 1,980
  2. 2025Kimi k1.5: Scaling Reinforcement Learning with LLMscited by 1,193
  3. 2024Tulu 3: Pushing Frontiers in Open Language Model Post-Trainingcited by 988
  4. 2025Group Sequence Policy Optimizationcited by 788
  5. 2025Kimi K2: Open Agentic Intelligencecited by 523
  6. 2019A Theory of Regularized Markov Decision Processescited by 279
  7. 2025MiniMax-M1: Scaling Test-Time Compute Efficiently with Lightning Attentioncited by 248
  8. 2025AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoningcited by 234
  9. 2017Bridging the Gap Between Value and Policy Based Reinforcement Learningcited by 178
  10. 2023Beyond Reverse KL: Generalizing Direct Preference Optimization with Diverse Divergence Constraintscited by 174
  11. 2021Policy Mirror Descent for Reinforcement Learning: Linear Convergence, New Sampling Complexity, and Generalized Problem Classescited by 149
  12. 2022On the Convergence Rates of Policy Gradient Methodscited by 115
  13. 2020Mirror Descent Policy Optimizationcited by 102
  14. 2021Policy Mirror Descent for Regularized Reinforcement Learning: A Generalized Framework with Linear Convergencecited by 83
  15. 2024Correcting the Mythos of KL-Regularization: Direct Alignment without Overoptimization via Chi-Squared Preference Optimizationcited by 78
  16. 2024REBEL: Reinforcement Learning via Regressing Relative Rewardscited by 66
  17. 2024Asynchronous RLHF: Faster and More Efficient Off-Policy RL for Language Modelscited by 65
  18. 2024Offline Regularised Reinforcement Learning for Large Language Models Alignmentcited by 48
  19. 2022Linear Convergence of Natural Policy Gradient Methods with Log-Linear Policiescited by 44
  20. 2023Foundations of Reinforcement Learning and Interactive Decision Makingcited by 33
  21. 2023A Novel Framework for Policy Mirror Descent with General Parameterization and Linear Convergencecited by 29
  22. 2025Trajectory Balance with Asynchrony: Decoupling Exploration and Learning for Fast, Scalable LLM Post-Trainingcited by 27
  23. 2023Sample Complexity of Neural Policy Mirror Descent for Policy Optimization on Low-Dimensional Manifoldscited by 1

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2024Operator World Models for Reinforcement Learning
  2. 2025Reinforcement Learning with Discrete Diffusion Policies for Combinatorial Action Spaces
  3. 2026Beyond State-Wise Mirror Descent: Offline Policy Optimization with Parametric Policies
  4. 2025One-Step Flow Policy Mirror Descent
  5. 2025Stabilizing Reinforcement Learning with LLMs: Formulation and Practices
  6. 2026Bellman Policy Optimization
  7. 2026DDO-RM: Distribution-Level Policy Improvement after Reward Learning
  8. 2025EPO: Entropy-regularized Policy Optimization for LLM Agents Reinforcement Learning