---
title: PMD-mean in KL-Regularized RL
url: https://www.emergentmind.com/topics/pmd-mean
type: topic
---

# PMD-mean in KL-Regularized RL

The term **PMD-mean** typically refers to an approximation method for mirror-descent updates in KL-regularized policy optimization, particularly in large-action-space reinforcement learning (RL) for language models. The core methodology of PMD-mean is to replace the log-partition normalization (which is impractical to estimate in vast action spaces) with the sample mean reward under the sampling policy, thus inducing a specific regression objective in log-policy space. This construction leads to robust policy improvement dynamics, especially under limited or noisy reward signals [2602.05933].

## 1. Policy Mirror Descent and the Partition Function Challenge

Policy Mirror Descent (PMD) considers the regularized objective, for each context $x$:
\[
\pi_{t+1}(\cdot|x) = \arg\max_{\pi \in \Delta(\mathcal{Y})} \Big\{ \mathbb{E}_{y \sim \pi}[\, r(x, y) \,] - \tau\, \mathrm{KL}(\pi \| \pi_t) \Big\}
\]
Its exact solution is the Boltzmann update:
\[
\pi_{t+1}(y|x) = \frac{\pi_t(y|x) \exp(r(x, y)/\tau)}{Z_t(x)}, \quad Z_t(x) = \mathbb{E}_{y \sim \pi_t}[\, e^{r(x, y)/\tau} \,]
\]
However, computing $Z_t(x)$ is infeasible in high-dimensional output spaces due to the expectation over exponentially many actions.

## 2. The PMD-mean Algorithm: Formulation and Implementation

PMD-mean circumvents the partition function calculation by using the mean reward as a surrogate. Specifically:
- Compute the baseline $b(x) = \mathbb{E}_{y \sim \pi_t}[r(x,y)]$.
- Use the "advantage" $\Delta(x, y) = r(x, y) - b(x)$ as the regression target.
The regression objective is:
\[
\mathcal{L}_{\text{mean}}(\pi) = \mathbb{E}_{x \sim \mathcal{D}} \mathbb{E}_{y \sim \pi_t(\cdot|x)} \left[ \frac{1}{2}\left( \log \frac{\pi(y|x)}{\pi_t(y|x)} - \frac{\Delta(x, y)}{\tau} \right)^2 \right]
\]
This selects $\log$-ratios in policy space to match (in the squared sense) the scaled advantage, subject to the normalization constraint $\mathbb{E}_{y \sim \pi_t}[e^{u(y)}] = 1$ for $u(y) = \log\left( \frac{\pi(y)}{\pi_t(y)} \right)$.

The closed-form population solution involves the Lambert-$W$ function:
\[
\pi_{t+1}(y) = \pi_t(y)\, \exp\left( \frac{\Delta_y}{\tau} - W\left( \frac{\lambda}{\tau^2} e^{\Delta_y/\tau} \right) \right)
\]
with $\lambda$ chosen to ensure normalization. In practice, models are trained using stochastic gradient-based regression on the log-policy ratios [2602.05933].

The update is summarized in the table below:

| Step            | Operation                                                        | Note                         |
|-----------------|------------------------------------------------------------------|------------------------------|
| Baseline        | $b(x) = \mathbb{E}_{y \sim \pi_t}[r(x,y)]$                      | Sample mean over rollouts    |
| Advantage       | $\Delta(x, y) = r(x, y) - b(x)$                                 | Scalar difference            |
| Loss            | $\frac{1}{2}\left( \log \frac{\pi(y|x)}{\pi_t(y|x)} - \Delta/\tau \right)^2$ | Regression in $\log$-space    |
| Constraint      | $\mathbb{E}_{\pi_t}[e^{u(y)}] = 1$                              | Normalization                |

## 3. Theoretical Implications: Implicit Mixed Regularization

The statistical solution to this regression coincides with the optimizer of a mirror-descent subproblem using a mixed KL–$\chi^2$ regularizer:
\[
\max_\pi \left\{ \mathbb{E}_\pi[r] - \tau\, \mathrm{KL}(\pi \|\pi_t) - \frac{\lambda}{2\tau} \chi^2(\pi \|\pi_t) \right\}
\]
where $\chi^2(\pi \|\pi_t) = \mathbb{E}_{y \sim \pi_t}\left[ \left(\frac{\pi(y)}{\pi_t(y)} - 1\right)^2 \right ]$.

This regularizer dynamically penalizes large changes in the policy distribution, especially when the average reward is low. In the binary reward setting, as the expected reward $p$ becomes small (weak models or hard tasks), the effective $\chi^2$ regularization coefficient $\lambda/\tau$ remains $O(1)$, producing conservative policy updates and improving robustness against estimation error in the normalization constant [2602.05933].

## 4. PMD-mean Versus Exact PMD and Practical Behavior

When compared with the partition-function ("PMD-part") update, PMD-mean is inherently less aggressive in its policy ratio updates, both on positive and negative actions—particularly early in training when rewards are sparse. The ideal Boltzmann update in PMD contracted policies faster but was far more susceptible to instability and policy collapse under finite-sample error in $Z_t(x)$. In contrast, PMD-mean attains slower, more controlled contraction, yielding better stability and robustness in practical settings, as confirmed by both theoretical sample-complexity bounds and direct experiments [2602.05933].

## 5. Empirical Results and Stability Properties

Experiments on large language models (Qwen2.5-7B, Qwen3-30B-A3B-Base) for binary-reward math reasoning (DAPO-Math-17k, AIME’24/’25) confirm the practical strengths of PMD-mean:
- Achieves superior accuracy compared to PPO-style methods and even strong on-policy gradients, with overall improvements in pass@1 scores on math benchmarks.
- Maintains training stability (avoids collapse) even for small $\tau$ or large global batch sizes (high staleness), unlike PMD-part.
- Enables efficient batching, tolerates larger rollout staleness without explicit importance weights, resulting in $4$--$5\times$ time speedup for token generation.
- Policy ratio statistics match theoretical envelopes, confirming conservative update dynamics predicted by the mixed regularization effect [2602.05933].

## 6. Practical Recommendations and Implementation Guidance

Key recommendations for applying PMD-mean in large-model post-training include:
- Use small $\tau$ (e.g., $0.005$–$0.02$ for 7B, $0.1$ for 30B MoE), leveraging the adaptive $\chi^2$ penalty for stability.
- Choose rollout budgets proportional to the inverse mean reward in early training ($K \sim O(1/p)$), with $K = 16$ or $32$ generally sufficient.
- Allow large global batch sizes for efficiency; PMD-mean tolerates these without requiring importance corrections.
- Apply explicit log-policy ratio clipping to enforce theoretical ratio bounds, ensuring stability.
- Monitor policy ratio histograms during training to verify that updates remain within predicted conservative limits.

## 7. Broader Significance and Comparison to Related Methodologies

PMD-mean exemplifies a principled, sample-efficient approach to RL-based policy improvement for LLMs and other high-action-space generative models. Its primary advance lies in enabling stable, sample-efficient, and automatically adaptive KL-regularized policy optimization without the need for precise normalization over the action space. This stands in contrast to alternative approaches that either require expensive importance sampling, rely on variance-prone partition estimates, or lack the self-correcting implicit regularization property of PMD-mean [2602.05933].

The success of PMD-mean in LLM post-training suggests applicability to other structured decision-making settings, especially those exhibiting sparse rewards and prohibitive action spaces, highlighting a robust paradigm for scalable and stable RL in complex generative modeling tasks.

Source: https://www.emergentmind.com/topics/pmd-mean