PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zhenghao Xu, Qin Lu, Changlong Yu, Tuo Zhao

arXiv:2602.05933v1cs.LG

TL;DR

PMD’s ideal updates are difficult to implement in LLM post-training because partition functions and finite-rollout targets are hard to estimate reliably. The paper analyzes PMD-MEAN, which replaces the log-partition term with mean reward, and shows that its induced adaptive χ2 regularization moderates policy changes and improves stability under limited rollouts.

  • Problem

    Reliable partition-function estimation and ideal-target fitting remain challenging for KL-regularized PMD in LLMs’ large action spaces and finite-rollout settings.

  • Method

    PMD-MEAN fits a log-policy-space regression target that approximates the log-partition term using the sampling policy’s mean reward.

  • Results

    PMD-MEAN is mathematically equivalent to mirror descent with an adaptive mixed KL–χ2 regularizer and is less sensitive to finite-sample rollout errors than PMD-PART.

  • Takeaways & Limitations

    The induced χ2 term constrains large probability-ratio changes, providing a mechanism for more stable policy updates in data-constrained LLM post-training.

  • Takeaways & Limitations

    PMD-MEAN is accurate as τ approaches infinity, but for small τ its mean-reward approximation can diverge from the log-partition function and ideal KL-subproblem solution.

Abstract

from arXiv · show

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.

1 Introduction

The paper studies PMD-MEAN, an off-policy regression approach that approximates the log-partition term with mean reward, and characterizes its optimization objective. It shows that the method induces adaptive mixed KL–χ2 regularization, improving stability under limited rollouts.

  • Motivation: RL post-training algorithms can be formalized as regularized policy improvement, while PMD provides KL-regularized subproblems with partition-normalized closed-form updates.In practice, estimating the partition function and fitting ideal targets from finite rollouts is challenging in large action spaces.
  • Motivation: Large or asynchronous generation batches create sampling-policy staleness, producing a training/inference mismatch and instability.Practitioners commonly mitigate this mismatch through importance weighting with clipping.
  • PMD-MEAN: PMD-MEAN replaces the exact partition-normalized target with a mean-reward approximation under the sampling policy and fits the result through off-policy regression.The approximation can diverge substantially from partition-normalized updates under the smaller regularization typical in practice.
  • Theory: PMD-MEAN has a closed-form Lambert-W solution equivalent to mirror descent with an adaptive mixed KL–χ2 regularizer.The χ2 weight depends dynamically on the current policy’s mean reward.
  • Regularization and stability: The induced χ2 term penalizes probability changes more strongly, controls probability ratios even with minimal nominal KL, and reduces sensitivity to finite-sample target errors.The effect is strongest when mean rewards are low, moderating early convergence and reducing overfitting to misestimated targets relative to PMD-PART.
  • Convergence and experiments: The paper develops inexact-PMD convergence analysis and reports improved efficiency, stability, and performance over standard GRPO on math reasoning tasks.The convergence analysis distinguishes PMD-MEAN from PMD-PART and characterizes their separations.

2 Preliminaries

The paper formulates LLM post-training as a contextual bandit and introduces PMD-MEAN as a practical alternative to partition-function-based policy mirror descent updates. The approximation replaces the log-partition function with a mean-reward baseline, enabling efficient regression from sampled rollouts.

  • Contextual bandit formulation: LLM post-training is modeled as a contextual bandit with prompts as states, generated responses as actions, and rewards bounded in [0, 1].A policy maps each prompt to a distribution over responses and induces expected reward under the prompt distribution.
  • Policy mirror descent: KL-regularized PMD updates each policy by maximizing expected reward minus τ times the KL divergence from the previous policy.The regularization parameter τ > 0 controls the strength of the update constraint.
  • Policy mirror descent: The ideal PMD solution exponentially reweights the previous policy by reward and uses a partition function for normalization.This closed-form update resembles a Boltzmann distribution but requires evaluating the partition function across a large action space.
  • Fitting the target by regression: Finite-rollout estimation of the partition function can become highly unstable in large action spaces and with small τ, causing pathological update behavior.Direct log-policy regression recovers the ideal update only under realizability and precise knowledge of Z_t(x).
  • PMD-MEAN: Log-partition approximation: PMD-MEAN approximates the log-partition function with the average reward and instead regresses using the advantage Δ(x, y) = r(x, y) − E_y′∼π_t(·|x)[r(x, y′)].The mean reward is efficiently estimated per prompt by Monte Carlo averaging over sampled responses; the method has been adopted for Kimi K1.5 and K2.

3 Implicit Regularization of PMD-MEAN

PMD-MEAN can diverge from the ideal KL-based PMD update because its average-reward approximation is inaccurate for small or moderate τ. Its population solution instead corresponds to an adaptive mixed KL–χ2-regularized improvement problem that yields more conservative policy changes.

  • Approximation gap: The average-reward approximation is accurate as τ →∞ but can differ substantially from the log-partition function for small or moderate τ.Consequently, PMD-MEAN may converge to a different subproblem solution than the ideal KL update even with infinite samples.
  • Population solution: Theorem 3.1 characterizes PMD-MEAN’s unique population minimizer using the mean-baseline advantage ∆y = r(y) − Eπt[r(y)].The theorem assumes πt(y) > 0 for every action and defines the solution over the probability simplex.
  • Population solution: PMD-MEAN heterogeneously normalizes action probabilities through the principal Lambert-W function, unlike the action-independent log-partition normalization in the KL solution.Because W is monotone, larger-advantage actions are suppressed relative to the KL solution, making PMD-MEAN less aggressive than PMD-PART.
  • Binary-reward behavior: When the average reward p is small, PMD-MEAN produces a substantially larger gap from PMD-PART in positive-action ratios and clearer separation for negative actions.Both positive-action ratios approach 1/p from below as τ →0, while the shared upper bound 1/p is approached at distribution-dependent exponential rates O(e^p/τ) versus O(e^1/τ).
  • Implicit regularization: PMD-MEAN is equivalent to mirror descent with an additional χ2 penalty, which directly suppresses large policy-ratio spikes and is stronger than KL regularization under full support.The equivalence follows from comparing the KKT conditions of the two optimization problems.
  • Implicit regularization: For binary rewards, the effective χ2 strength λ/τ remains O(1) as τ →0, so PMD-MEAN continues regularizing updates even when nominal KL regularization is small.This adaptive regularization explains why the method constrains policy changes in low-temperature regimes.

4 Implications on Convergence

The convergence analysis shows a tradeoff: PMD-PART has a faster ideal improvement rate, while PMD-MEAN is more robust to finite-rollout estimation errors when pass rates are low. A refined analysis further shows that PMD-MEAN recovers the ideal target policy as rollout size grows and shares PMD-PART’s asymptotic error scaling.

  • Ideal convergence rates: PMD-PART converges faster ideally than PMD-MEAN when the pass rate p_t is small and τ is small, although both approach one-step convergence as τ → 0.The comparison concerns the large-rollout regime, where estimation error is negligible.
  • Finite-sample robustness: For small τ, PMD-PART’s log-ratio factors are worse than PMD-MEAN’s, with the gap becoming significant when p_t is small.These larger factors amplify the effect of inexact updates in the convergence bound.
  • Target estimation: PMD-MEAN’s estimated target is systematically larger than the ideal target for positive actions but preserves more conservative shrinkage for negative actions at rate exp(−p_t/τ).The positive-action mismatch produces more aggressive improvement, whereas negative actions remain closer to the ideal target.
  • Finite-sample robustness: PMD-MEAN has smaller target-estimation error than PMD-PART when both the pass rate p_t and rollout size n are small.This explains why PMD-PART’s faster ideal rate can become unstable in practice with limited rollouts.
  • Refined PMD-MEAN analysis: As n → ∞, the empirical PMD-MEAN minimizer recovers the ideal target policy because the constraint E_πt[es_π] = 1 pulls log-ratios back from exact advantage fitting.The refined analysis therefore removes the apparent error floor from the basic target-estimation bound.

5 Experiments

Experiments on math reasoning RL show that PMD-MEAN improves performance over GRPO, matches or exceeds relevant alternatives, and offers faster, more stable training. Its observed policy-ratio behavior also supports the theoretical analysis of weaker policy decreases than PMD-PART.

  • Main results: PMD-MEAN significantly outperforms GRPO, gaining +2.6% AIME24 and +9.0% AIME25 on 7B, and +14.6% AIME24 and +8.1% AIME25 on 30B MoE.The gains use τ = 0.005 for the 7B model and τ = 0.1 for the 30B MoE model.
  • Efficiency: Off-policy PMD-MEAN achieves comparable performance to on-policy gradient with staleness 1 while providing a 4.6× speedup through larger rollout batches.The larger global batch amortizes inference cost, while actor update cost remains comparable.
  • Stability: PMD-MEAN remains stable during training, whereas PMD-PART is highly unstable and can collapse even with a much larger τ.This comparison is shown in the training curves for Qwen2.5-7B on DAPO-Math-17k.
  • Policy ratios: The policy decrease in PMD-MEAN is weaker than in PMD-PART, then strengthens as training proceeds and accuracy improves.The trend is measured using minimum log policy ratios as an approximation of the next policy update.
  • Beyond standard GRPO: Against GSPO, PMD-MEAN outperforms on Qwen2.5-7B and achieves comparable performance on the Qwen3-30B-A3B-Base MoE model.GSPO is a GRPO variant using sequence-level importance sampling with clipping and geometric mean normalization for MoE stability.

6 Related Work

Related work situates PMD-mean within RL-based LLM post-training, critic-free policy optimization, off-policy correction, and policy mirror descent. Prior methods use value networks, pairwise relative rewards, or group-based approximations to address the log-partition term in ideal KL updates.

  • RL for Post-Training of LLMs: RLHF/RLAIF and RLVR are predominant paradigms for LLM post-training, with particular efficacy in mathematical reasoning, coding, and logical tasks.These approaches have inspired large-scale RL methodologies and architectural designs for increasingly complex agentic capabilities.
  • RL for Post-Training of LLMs: TRPO and PPO are foundational policy-gradient methods, but critic models create estimation bias and computational overhead in LLM post-training.GRPO and RLOO address these limitations by eliminating the critic-model requirement.
  • RL for Post-Training of LLMs: PPO and GRPO use token-level importance sampling with clipping, while GSPO and CISPO use sequence-level or detached-clipping importance sampling to improve stability in MoE training.These methods depend on sampling distributions closely matching the current policy and therefore require off-policy correction mechanisms.
  • Policy Mirror Descent: Policy mirror descent has a classical role in reinforcement-learning policy optimization, but much of its iteration- and sample-complexity theory focuses on tabular or function-approximation settings.The supplied passage notes extensive theoretical analyses while indicating that these settings differ from the paper’s focus.
  • PMD-PART: Prior approximations to ideal KL solutions include value-network estimation, partition-free pairwise relative rewards, and group-based estimates using other loss terms.PMD-PART uses a value network, Gao et al. fit pairwise relative rewards, and Bartoldson et al. use a stop-gradient group average.

7 Conclusion … A.2 Hyperparameters

The paper concludes that PMD-MEAN has a mathematically characterized adaptive mixed KL–χ2 regularization mechanism, while identifying advanced sampling corrections as future work. The appendices document the datasets, base models, prompt format, and hyperparameter summary used in the experiments.

  • 7 Conclusion: PMD-MEAN’s population update is characterized exactly through the Lambert-W function and is equivalent to mirror descent with an adaptive mixed KL–χ2 regularizer.This framework explains a concrete mechanism underlying the algorithm’s stability.
  • 7 Conclusion: Oversampling strategies and importance sampling corrections for training/inference engine mismatches are identified as potential directions for future research.The analysis focuses deliberately on the principled form of PMD-MEAN to enable clearer theoretical understanding.
  • A Experimental Details: Experiments use the deduplicated DAPO-Math-17k dataset with Qwen2.5-7B and Qwen3-30B-A3B-Base as base models.These are the stated training dataset and base-model configurations.
  • A.1 Prompt Template: The prompt template follows Yu et al. (2025) and processes each question in the specified problem_statement format.The appendix states that this template is used for all questions.
  • A.1 Prompt Template: The Chain-of-Thought prompt instructs models to solve each math problem step by step and place the final response on its own line after Answer:.The required final line has the form Answer: $Answer, without quotes.
  • A.2 Hyperparameters: The appendix summarizes the key hyperparameters for the 7B dense-model and 30B MoE-model experiments in Table 4.The supplied passage identifies the table but does not provide its parameter values.

A.3 Implementation Details · B Missing Proofs in Section 3

The implementation builds on verl5, disables explicit reference-model KL penalties for GRPO and GSPO, and compares several policy-optimization losses. It samples K responses from the old policy, assigns correctness rewards in {−1, +1}, and applies specified normalization and tuning conventions.

  • A.3 Implementation Details: The implementation is based on verl5 and follows the DAPO recipe by disabling explicit KL penalties for GRPO and GSPO.This KL setting is with respect to the base reference model.
  • A.3 Implementation Details: For each prompt, the method samples K responses from the old policy and assigns each response a correctness reward in {−1, +1}.The trainable policy is denoted πθ, while the sampling policy is πt.
  • A.3 Implementation Details: The implementation defines token probability ratios and geometric normalized sequence probability ratios for policy updates.These ratios are introduced before defining the advantage estimators and optimization losses.
  • A.3 Implementation Details: GRPO uses an explicitly defined loss, and when global and mini-batch sizes match, it reduces to an on-policy gradient with the GRPO advantage estimator.The reduction is stated for equal global batch and mini-batch sizes.
  • A.3 Implementation Details: On-policy gradient experiments use the length-normalized RLOO loss because its empirical performance is similar to using the leave-one-out advantage.The stated empirical comparison motivates the selected loss.
  • A.3 Implementation Details: GSPO is implemented with ϵlow = 3 × 10−4 and ϵhigh = 4 × 10−4, while PMD-MEAN and PMD-PART use their specified losses.The clipping values are given as suggested settings.
  • A.3 Implementation Details: The factor τ limits gradient-norm variation when tuning τ, and length normalization matches the aggregation mode used by other methods.Both conventions are described as implementation choices in the loss.

B.1 Proof of Theorem 3.1 … C Missing Proofs in Section 4.1

The appendices establish the PMD-MEAN optimizer through KKT conditions, derive binary-reward policy-ratio asymptotics, identify its mixed-subproblem equivalence, and provide a Bernstein self-bounding variance lemma for Section 4.1.

  • B.1 Proof of Theorem 3.1: Under πt(y) > 0 for all y, stationary points of Lmean satisfy the required conditions, an assumption holding for LLMs without top-p/top-k constraints.The proof uses positivity of the sampling policy over all actions.
  • B.1 Proof of Theorem 3.1: The KKT conditions determine a unique λ ≥ 0, and the Lagrangian Hessian in u is positive definite, making the stationary point the minimizer of Lmean.Monotonicity in λ and positive definiteness establish uniqueness and optimality.
  • B.1 Proof of Theorem 3.1: For binary rewards with p = Eπt[r(y)], the small-temperature analysis yields λ = τp(1 −p)(1 + o(1)).The result follows from the Lambert-W asymptotics as τ → 0.
  • B.2 Policy Ratio: Proposition B.1 characterizes PMD-MEAN and PMD-PART policy ratios for binary rewards, with the PMD-MEAN claims established using Lambert-W expansions.The proposition assumes r(y) ∈{0, 1} and p ∈ (0, 1).
  • B.2 Policy Ratio: For PMD-PART, the policy-ratio update is explicit and is analyzed by expanding (1 + u)^−1 = 1 −u + O(u2).The expansion uses u = (1 −p)e−1/τ.
  • B.3 Proof of Proposition 3.2: The PMD-MEAN population solution satisfies exactly the same KKT conditions as mixed subproblem (12), so it also solves that subproblem with the same λ.The proof rewrites rewards using the centered quantity ∆y := r(y) −Eπt[r].
  • C Missing Proofs in Section 4.1: Lemma C.1 applies Bernstein’s inequality to i.i.d. bounded variables with nonnegative mean and self-bounding variance, giving a high-probability concentration bound.Its assumptions include E[Zi] = µ ≥0, |Zi| ≤R, and E[Zi^2] ≤vµ, with probability at least 1 −δ.

C.1 Proof of Lemma 4.5 … C.5 Proof of Proposition 4.11

The appendix proves the main finite-sample and convergence claims through concentration, policy-mismatch bounds, and explicit binary-reward analyses. It also establishes the leave-one-out concentration and PMD-PART error control used in Proposition 4.11.

  • C.1 Proof of Lemma 4.5: Lemma 4.5 defines a residual and clean empirical loss, then applies concentration and a union bound over policies to prove the target bound.The proof uses E[Xi(π)] = Lt(π), sets v = R = 1/(2M^2), and concludes (20) using M ≤ 2B.
  • C.2 Proof of Theorem 4.6: Theorem 4.6 bounds policy-value differences by total variation and combines this with an exponential-moment estimate and Lemma 4.5 to obtain (22).Because r ∈ [0, 1], |J(p) − J(q)| ≤ TV(p, q); the prefactor is at most e^B+/2 up to constants.
  • C.3 Proof of Proposition 4.7: For zero rewards, Proposition 4.7 derives the contraction of total probability mass on the negative set from the PMD update formula.The proof begins from equation (10) with r(y) = 0.
  • C.4 Proof of Proposition 4.8: Proposition 4.8 computes the total mass on r = 1 and identifies its recursion with equation (21) under condition (24).The resulting expression is t+1 = 1−pt pte1/τ +(1−pt).
  • C.5 Proof of Proposition 4.11: Lemma C.2 establishes a simultaneous leave-one-out mean concentration bound for i.i.d. Bernoulli variables using Bernstein’s inequality and a union bound over indices and signs.The variables have mean p, with p−i defined as the average excluding observation i.
  • C.5 Proof of Proposition 4.11: For PMD-PART, Proposition 4.11 expresses each log-ratio mismatch through the mean value theorem and bounds it using the leave-one-out deviation and the uniform inequality |∆i| ≤ 1/τ.The proof writes ∆i = log(1 + apt) − log(1 + ap−i) and combines the resulting inequalities after squaring and averaging.
  • C.5 Proof of Proposition 4.11: The remaining Proposition 4.11 argument bounds the two-value reward contribution using τ→0 asymptotics and combines the terms to prove (25).As τ → 0 with fixed pt ∈ (0, 1), τw+ = (1 − pt) + o(1) and τw− = o(1).

D Refined Analysis for PMD-MEAN · D.1 Proof of Lemma D.1

The refined analysis connects PMD-MEAN’s population squared loss to its ideal log-policy target and population objective. Lemma D.1 underpins a refined ERM analysis that removes the target-estimation error floor.

  • D Refined Analysis for PMD-MEAN: The refinement begins by relating the population squared loss L_t in (15) to the ideal target s⋆(y) = log π_t(y) with the population PMD-MEAN objective.This establishes the analytical connection used throughout the refined PMD-MEAN analysis.
  • D Refined Analysis for PMD-MEAN: Lemma D.1 fixes a global step t and centers rewards through Δ_y := r(y) − E_y′∼π_t[r(y′)].The centered reward difference is then scaled by τ in the proof notation.
  • D Refined Analysis for PMD-MEAN: The lemma defines the PMD-MEAN population objective as the bandit specialization of objective (4).This objective is the population counterpart connected to the loss analysis.
  • D.1 Proof of Lemma D.1: In the proof of Lemma D.1, the centered reward is abbreviated as g(y) := Δ_y/τ before analyzing any π ∈ Π.This notation initiates the proof’s algebraic argument.
  • D Refined Analysis for PMD-MEAN: Using Lemma D.1, the analysis refines PMD-MEAN’s ERM guarantee and eliminates the error floor in target-estimation error.The result is stated as Lemma D.2 under Assumptions 4.1 to 4.4, binary rewards, and p_t := E_y∼π_t[r(y)].
  • D Refined Analysis for PMD-MEAN: Under the stated assumptions, Lemma D.2 provides a high-probability refined ERM result for any δ ∈ (0, 1), using ε_n(p_t, δ) from Proposition 4.11.The guarantee holds with probability at least 1 − δ.

D.2 Proof of Lemma D.2

The proof decomposes the PMD-MEAN empirical targets, controls the leave-one-out deviation on a high-probability event, and bounds population and empirical excess risks. It combines concentration, a union bound, and Lemma D.1 to establish the result.

  • The proof begins by recalling the PMD-MEAN leave-one-out target and decomposing the empirical targets.
  • On the event from Proposition 4.11, the leave-one-out deviation is controlled by ε_n(p_t, δ/2).
  • The proof bounds the target-error term using the score bound |s_π(y)| ≤ B and binary rewards r(y_i) ∈ {0, 1}.
  • For population excess risk, it bounds Z_i(π), applies Lemma D.1, and then invokes Lemma C.1 with v = 2M_τ^2.
  • A union bound with δ′ = δ/(2|Π|) yields the population-risk bound simultaneously for all π ∈ Π with probability at least 1 − δ/2.
  • The proof finally bounds empirical excess risk using Assumptions 4.1 and 4.2 together with inequality (45), then combines this with (46).
Loading 2602.05933v1…