PMD-mean in KL-Regularized RL

• PMD-mean is an approach that approximates mirror-descent updates by replacing the log-partition with the sample mean reward, addressing challenges in large-action-space settings.
• The algorithm computes a baseline from expected rewards and uses the resultant advantage in a regression objective within log-policy space, leading to conservative yet robust policy updates.
• Empirical results demonstrate that PMD-mean improves training stability and efficiency in language models, yielding effective policy enhancements even under sparse or noisy reward conditions.

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 LLMs. 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 (Xu et al., 5 Feb 2026).

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-$$\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\}$$0 function: $$\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\}$$1 with $$\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\}$$2 chosen to ensure normalization. In practice, models are trained using stochastic gradient-based regression on the log-policy ratios (Xu et al., 5 Feb 2026).

The update is summarized in the table below:

Step	Operation	Note
Baseline	$$\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\}$$3	Sample mean over rollouts
Advantage	$$\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\}$$4	Scalar difference
Loss	$$\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\}$$5	Regression in $$\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\}$$6-space
Constraint	$$\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\}$$7	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–$$\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\}$$8 regularizer: $$\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\}$$9 where $$\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} \,]$$0.

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 $$\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} \,]$$1 becomes small (weak models or hard tasks), the effective $$\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} \,]$$2 regularization coefficient $$\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} \,]$$3 remains $$\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} \,]$$4, producing conservative policy updates and improving robustness against estimation error in the normalization constant (Xu et al., 5 Feb 2026).

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 $$\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} \,]$$5. 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 (Xu et al., 5 Feb 2026).

5. Empirical Results and Stability Properties

Experiments on LLMs (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 $$\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} \,]$$6 or large global batch sizes (high staleness), unlike PMD-part.
• Enables efficient batching, tolerates larger rollout staleness without explicit importance weights, resulting in $$\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} \,]$$7--$$\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} \,]$$8 time speedup for token generation.
• Policy ratio statistics match theoretical envelopes, confirming conservative update dynamics predicted by the mixed regularization effect (Xu et al., 5 Feb 2026).

6. Practical Recommendations and Implementation Guidance

Key recommendations for applying PMD-mean in large-model post-training include:

• Use small $$\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} \,]$$9 (e.g., $Z_t(x)$0–$Z_t(x)$1 for 7B, $Z_t(x)$2 for 30B MoE), leveraging the adaptive $Z_t(x)$3 penalty for stability.
• Choose rollout budgets proportional to the inverse mean reward in early training ($Z_t(x)$4), with $Z_t(x)$5 or $Z_t(x)$6 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 (Xu et al., 5 Feb 2026).

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.