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

• URL: http://arxiv.org/abs/2602.05933v1
• Date: Thu, 05 Feb 2026 17:44:28 GMT
• Title: Approximation of Log-Partition Function in Policy Mirror Descent Induces Implicit Regularization for LLM Post-Training
• Authors: Zhenghao Xu, Qin Lu, Changlong Yu, Tuo Zhao,
• Abstract summary: Policy mirror descent (PMD) provides a principled framework for reinforcement learning (RL)<n>We investigate a practical algorithm, termed PMD-mean, that approximates the log-partition term with the mean reward under the sampling policy.<n>Experiments on math reasoning tasks show that PMD-mean achieves superior performance with improved stability and time efficiency.
• Score: 33.61029387987583
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Rethinking the Trust Region in LLM Reinforcement Learning [72.25890308541334]
  Proximal Policy Optimization (PPO) serves as the de facto standard algorithm for Large Language Models (LLMs)<n>We propose Divergence Proximal Policy Optimization (DPPO), which substitutes clipping with a more principled constraint.<n>DPPO achieves superior training and efficiency compared to existing methods, offering a more robust foundation for RL-based fine-tuning.
  arXiv  Detail & Related papers  (2026-02-04T18:59:04Z)
• A Step Back: Prefix Importance Ratio Stabilizes Policy Optimization [58.116300485427764]
  Reinforcement learning post-training can elicit reasoning behaviors in large language models.<n> token-level correction often leads to unstable training dynamics when the degree of off-policyness is large.<n>We propose a simple yet effective objective, Minimum Prefix Ratio (MinPRO)
  arXiv  Detail & Related papers  (2026-01-30T08:47:19Z)
• Stabilizing Policy Gradients for Sample-Efficient Reinforcement Learning in LLM Reasoning [77.92320830700797]
  Reinforcement Learning has played a central role in enabling reasoning capabilities of Large Language Models.<n>We propose a tractable computational framework that tracks and leverages curvature information during policy updates.<n>The algorithm, Curvature-Aware Policy Optimization (CAPO), identifies samples that contribute to unstable updates and masks them out.
  arXiv  Detail & Related papers  (2025-10-01T12:29:32Z)
• wd1: Weighted Policy Optimization for Reasoning in Diffusion Language Models [15.638885149395657]
  Intractability of dLLMs likelihood function requires approximating the current, old, and reference policy likelihoods at each policy optimization step.<n>We introduce $mathttwd1$, a novel policy optimization approach that reformulates the objective as a weighted likelihood.<n>Experiments on widely used reasoning benchmarks demonstrate that $mathttwd1$, without supervised fine-tuning (SFT) or any supervised data, outperforms existing RL methods for dLLMs.
  arXiv  Detail & Related papers  (2025-07-07T21:27:25Z)
• StaQ it! Growing neural networks for Policy Mirror Descent [4.672862669694739]
  In Reinforcement Learning (RL), regularization has emerged as a popular tool both in theory and practice.<n>We propose and analyze PMD-like algorithms that only keep the last $M$ Q-functions in memory.<n>We show that for finite and large enough $M$, a convergent algorithm can be derived, introducing no error in the policy update.
  arXiv  Detail & Related papers  (2025-06-16T18:00:01Z)
• Accelerating RL for LLM Reasoning with Optimal Advantage Regression [52.0792918455501]
  We propose a novel two-stage policy optimization framework that directly approximates the optimal advantage function.<n>$A$*-PO achieves competitive performance across a wide range of mathematical reasoning benchmarks.<n>It reduces training time by up to 2$times$ and peak memory usage by over 30% compared to PPO, GRPO, and REBEL.
  arXiv  Detail & Related papers  (2025-05-27T03:58:50Z)
• Trajectory Bellman Residual Minimization: A Simple Value-Based Method for LLM Reasoning [55.33984461046492]
  Policy-based methods currently dominate reinforcement learning pipelines for large language model (LLM) reasoning.<n>We introduce Trajectory Bellman Residual Minimization (TBRM), an algorithm that naturally adapts this idea to LLMs.<n>We prove convergence to the near-optimal KL-regularized policy from arbitrary off-policy via an improved change-of-trajectory-measure analysis.
  arXiv  Detail & Related papers  (2025-05-21T09:41:53Z)
• Provably Mitigating Overoptimization in RLHF: Your SFT Loss is Implicitly an Adversarial Regularizer [52.09480867526656]
  We identify the source of misalignment as a form of distributional shift and uncertainty in learning human preferences.<n>To mitigate overoptimization, we first propose a theoretical algorithm that chooses the best policy for an adversarially chosen reward model.<n>Using the equivalence between reward models and the corresponding optimal policy, the algorithm features a simple objective that combines a preference optimization loss and a supervised learning loss.
  arXiv  Detail & Related papers  (2024-05-26T05:38:50Z)
• Mirror Descent Policy Optimization [41.46894905097985]
  We propose an efficient RL algorithm, called em mirror descent policy optimization (MDPO) MDPO iteratively updates the policy by em approximately solving a trust-region problem. We highlight the connections between on-policy MDPO and two popular trust-region RL algorithms: TRPO and PPO, and show that explicitly enforcing the trust-region constraint is in fact em not a necessity for high performance gains in TRPO.
  arXiv  Detail & Related papers  (2020-05-20T01:30:43Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.