策略镜像下降中对数配分函数近似为LLM后训练引入隐式正则化

PMD-MEAN通过均值奖励近似隐式引入KL-χ²混合正则化，提升LLM RL训练稳定性

前置知识

策略镜像下降 (Policy Mirror Descent, PMD)

PMD是强化学习中一种经典的策略优化框架，其核心思想是在每一步迭代中求解一个KL正则化的策略改进子问题。具体来说，在全局步 $t$，PMD通过求解 $\pi_{t+1}(\cdot|x) = \arg\max_{\pi} \mathbb{E}_{y\sim\pi}[r(x,y)] - \tau \cdot \text{KL}(\pi \| \pi_t)$ 来更新策略，其中 $\tau > 0$ 是正则化参数。该子问题有闭式解，本质上是对旧策略按奖励信号进行指数重加权并用配分函数归一化。PMD为TRPO、PPO等实用算法提供了统一的理论框架。

本文的核心贡献正是分析PMD的一个实用近似变体PMD-MEAN的理论性质，因此理解PMD的基本框架是读懂全文的前提。

配分函数 (Partition Function)

在PMD的理想闭式解中，配分函数 $Z_t(x) = \mathbb{E}_{y\sim\pi_t}[e^{r(x,y)/\tau}]$ 是一个归一化常数，确保更新后的策略 $\pi_{t+1}$ 仍然是合法的概率分布。在LLM场景下，动作空间（可能的生成序列）极其庞大，精确计算这个期望值在计算上是不可行的，必须从有限的采样中估计，这引入了显著的估计误差。

本文的核心动机正是配分函数难以估计的问题——PMD-MEAN正是通过用均值奖励替代对数配分函数来规避这一困难，理解配分函数的作用才能理解这一近似为何有意义。

Lambert-W函数

Lambert-W函数 $W(\cdot)$ 是函数 $f(w) = w \cdot e^w$ 的反函数，即满足 $W(z)e^{W(z)} = z$。它在组合数学、物理学和概率论中有广泛应用。本文的关键理论结果之一就是证明PMD-MEAN的人口解可以通过Lambert-W函数写出精确闭式表达，这与标准PMD的Boltzmann分布形式形成鲜明对比。

Lambert-W闭式解是本文定理3.1的核心结果，理解这个函数的性质（单调性、渐近行为）对于理解PMD-MEAN为何比PMD-PART更保守至关重要。

KL散度与χ²散度

KL散度 $\text{KL}(p\|q) = \mathbb{E}_{y\sim p}[\log \frac{p(y)}{q(y)}]$ 和χ²散度 $\chi^2(p\|q) = \mathbb{E}_{y\sim q}[(\frac{p(y)}{q(y)}-1)^2]$ 都是衡量两个概率分布差异的度量。对于全支撑分布，总有 $\text{KL}(p\|q) \leq \chi^2(p\|q)$，即χ²散度对概率比的大幅变化施加更强的惩罚。χ²散度在策略优化中可以防止策略更新过于激进。

本文证明PMD-MEAN隐式地优化了KL-χ²混合正则化目标，χ²项是理解该算法为何更稳定的关键机制。

GRPO (Group Relative Policy Optimization)

GRPO是Shao等人2024年提出的LLM强化学习算法，它通过为每个prompt生成多个响应（组），用组内相对奖励作为基线来估计优势函数，从而避免了训练额外的价值网络（critic）。GRPO已成为DeepSeek等模型训练中的标准方法，但在大规模MoE模型训练中存在稳定性问题。

GRPO是本文实验的主要基线方法，理解GRPO的工作原理才能理解PMD-MEAN相比现有方法的改进幅度和原因。

Off-policy与On-policy学习

On-policy方法要求用于更新策略的数据来自当前策略本身的采样，这在LLM场景下意味着每次策略更新前都需要重新生成响应，计算开销巨大。Off-policy方法则允许使用旧策略生成的数据进行更新，可以通过增大全局批大小来摊销推理成本，但会引入采样策略与当前策略之间的分布偏移（staleness），需要通过重要性采样等技术进行校正。

本文的PMD-MEAN本质上是一种off-policy回归方法，它的优势正是在off-policy（staleness=16）设置下依然保持稳定性，这与传统需要复杂重要性采样校正的方法形成对比。

研究动机

在LLM后训练中使用强化学习时，策略镜像下降（PMD）提供了优雅的理论框架，其KL正则化子问题有闭式解，形式为对旧策略按奖励进行指数重加权并用配分函数 $Z_t(x) = \mathbb{E}_{y\sim\pi_t}[e^{r(x,y)/\tau}]$ 归一化。然而在实践中，LLM的动作空间极其庞大（所有可能的生成序列），从有限的rollout样本中可靠地估计这个配分函数是一个重大挑战。当正则化参数 $\tau$ 较小（实践中常见）时，配分函数的估计对样本噪声极为敏感，导致策略更新可能严重偏离理想目标。PMD-PART（直接拟配分归一化的目标）在rollout样本量 $n$ 较小、奖励通过率 $p_t$ 较低（训练早期）时，目标估计误差可以大到 $O(\frac{e^{2/\tau}p_t}{n})$ 的量级，这会导致训练崩溃。此外，现代高效RL实现倾向于使用大生成批次或异步rollout来避免长尾生成的计算瓶颈，这不可避免地引入了采样策略与被更新策略之间的staleness，传统方法通过重要性采样裁剪等启发式技术来缓解，但大大增加了实现和理论分析的复杂度。

本文的目标是本文旨在深入理解一种已被Kimi K1.5/K2等先进LLM采用的实用PMD变体——PMD-MEAN的理论基础。具体而言，PMD-MEAN用采样策略下的平均奖励近似对数配分函数，并在对数策略空间中进行回归。本文的核心目标是回答一个根本性问题：PMD-MEAN究竟优化的是什么目标？这种近似在算法层面有什么具体后果？通过严格的数学分析，揭示PMD-MEAN的隐式正则化机制，并解释其在实践中表现出的优越稳定性和效率。

与已有工作不同的是，现有文献对PMD-MEAN的理论理解存在显著空白。虽然该算法已被部署在Kimi K1.5/K2等前沿模型中并取得了优异的实验效果，但其数学基础一直缺乏严格的理论刻画。标准PMD理论假设可以精确计算配分函数或使用足够多的样本来近似，但PMD-MEAN的均值近似在小 $\tau$ 时会系统性地偏离理想目标，这意味着它实际上优化的是一个不同于标准KL正则化PMD的目标。本文的独特切入角度是从回归视角出发，推导PMD-MEAN人口解的精确闭式表达，将其与Lambert-W函数和自适应KL-χ²混合正则化建立等价关系，从而为这一实用算法提供坚实的理论支撑。这种分析方法不仅解释了PMD-MEAN的成功，还揭示了均值近似如何隐式地引入了有价值的正则化效果。

核心方法

本文从回归视角分析PMD-MEAN的整体思路是：首先将PMD的策略更新问题转化为对数策略空间中的回归问题，然后推导PMD-MEAN（均值奖励近似）的人口最小化解的精确闭式表达，最后证明这个解等价于求解一个带有自适应KL-χ²混合正则化的镜像下降子问题。直觉上，标准PMD-PART的目标是拟合 $\frac{r(x,y)}{\tau} - \log Z_t(x)$，其中配分函数项 $\log Z_t(x)$ 与具体动作 $y$ 无关，是一个全局归一化常数。而PMD-MEAN则用平均奖励 $\mathbb{E}_{\pi_t}[r(x,y)]$ 替代 $\log Z_t(x)$，使目标变为优势函数 $\Delta(x,y)/\tau$。这个看似简单的替换会导致解的结构发生根本性变化——从Boltzmann分布变为涉及Lambert-W函数的形式，且归一化常数变得与具体动作相关，高优势动作的概率增长被系统性地抑制。

本文的核心创新在于揭示PMD-MEAN的均值近似并非仅仅是配分函数的一个粗糙近似，而是隐式地求解了一个结构完全不同的正则化子问题。具体而言，PMD-MEAN的人口解 $\pi_{t+1}(y) = \pi_t(y) \exp\left(\frac{\Delta_y}{\tau} - \frac{\lambda}{\tau} W\left(\frac{\Delta_y}{\tau} e^{\Delta_y/\tau}\right)\right)$ 中的Lambert-W项本质上起到了χ²正则化的效果。这与标准PMD-PART的根本区别在于：PMD-PART只使用KL散度作为正则化器，而PMD-MEAN等价于使用自适应权重的KL-χ²混合正则化，其中χ²权重 $\lambda/\tau$ 根据当前策略下的奖励分布自适应调节。这种额外的χ²正则化直接抑制概率比的大幅变化，效果在平均奖励较低时（即训练早期）尤为显著，从而产生更保守的更新，增强了对有限样本估计误差的鲁棒性。

方法步骤详情

PMD-MEAN的理论分析包含以下关键步骤。第一步，建立回归框架：将PMD策略更新形式化为对数策略空间中的回归问题，定义 $s_\pi(y) = \log \frac{\pi(y)}{\pi_t(y)}$，目标是用 $s_\pi$ 拟合理想目标 $s^\star$。PMD-PART的目标为 $s^\star_{\text{part}}(y) = \frac{r(y)}{\tau} - \log Z_t$，PMD-MEAN的目标为 $s^\star_{\text{mean}}(y) = \frac{\Delta(y)}{\tau}$，其中 $\Delta(y) = r(y) - \mathbb{E}_{\pi_t}[r]$ 是均值基线优势。第二步，推导闭式解（定理3.1）：对PMD-MEAN的平方回归损失求人口最小化，利用KKT条件证明唯一最小化解满足 $\pi_{t+1}(y) = \pi_t(y) \exp(\frac{\Delta_y}{\tau} - \frac{\lambda}{\tau} W(\frac{\Delta_y}{\tau} e^{\Delta_y/\tau}))$，其中 $W$ 是Lambert-W主分支，$\lambda$ 是归一化常数。第三步，建立等价正则化问题（命题3.2）：证明上述解等价于求解 $\max_\pi \mathbb{E}_\pi[r] - \tau \text{KL}(\pi\|\pi_t) - \frac{\lambda}{2\tau} \chi^2(\pi\|\pi_t)$。第四步，收敛分析：在不精确PMD框架下，利用引理4.5将ERM误差分解为目标估计误差 $\Delta^2$ 和统计复杂度项，然后在二元奖励模型下分别实例化PMD-MEAN和PMD-PART的理想收敛速率 $\eta_t$、对数比界 $(B, B_+)$ 和目标估计误差 $\Delta^2$。

技术新颖性

本文在技术层面有多个新颖之处。首先，这是首次为PMD-MEAN（即Kimi风格PMD）给出精确的数学刻画，Lambert-W函数闭式解的发现揭示了均值近似带来的结构变化远比表面上看起来更深刻。其次，将PMD-MEAN与自适应KL-χ²混合正则化建立等价关系是一个重要的概念性贡献，它将一个看似ad hoc的近似技术与有坚实理论基础的正则化框架联系起来。第三，本文的收敛分析能够精确刻画PMD-MEAN和PMD-PART在有限样本下的分离：PMD-PART有更快的理想收敛速率（接近一步收敛），但对有限样本估计误差极其敏感；PMD-MEAN虽然理想速率稍慢，但目标估计误差显著更小，尤其在 $p_t$ 小（训练早期）和 $n$ 小的条件下。第四，命题4.11对二元奖励下两种方法的目标估计误差给出了精细的分布依赖界，PMD-PART的误差在小 $n$ 和小 $p_t$ 时可高达 $O(\frac{e^{2/\tau}p_t}{n})$，而PMD-MEAN仅为 $O(\frac{p_t(1-p_t)^2}{\tau^2} + \frac{p_t(1-p_t)}{\tau^2 n})$，差距可达数个数量级。

**Figure 2**: PMD-MEAN和PMD-PART在二元奖励下的对数概率比率对比

**Figure 5**: PMD-MEAN和PMD-PART的最小对数策略比率

实验结果

本文的实验在DAPO-Math-17k数据集上进行，使用Qwen2.5-7B和Qwen3-30B-A3B-Base作为基础模型，在AIME 2024和AIME 2025上评估数学推理能力。核心发现包括：第一，PMD-MEAN显著优于GRPO基线——在7B模型上，$\tau=0.005$ 在AIME24上获得+2.6%绝对提升，在AIME25上获得+9.0%绝对提升；在30B MoE模型上，$\tau=0.1$ 在AIME24上获得+14.6%绝对提升，在AIME25上获得+8.1%绝对提升。第二，效率优势明显：相比on-policy梯度（staleness=1），off-policy PMD-MEAN在取得相当性能的同时，通过利用更大的全局rollout批大小（512 prompts vs 32）摊销推理成本，实现了4.6倍的加速（整体每token仅需0.0126ms vs 0.0569ms）。第三，稳定性是关键优势：训练曲线显示PMD-MEAN在训练过程中保持稳定，而PMD-PART即使使用更大的 $\tau$ 也高度不稳定甚至可能崩溃。第四，PMD-MEAN的策略比率分析验证了理论预测——负动作的概率衰减比PMD-PART更温和，且随着训练推进和准确率提高而逐渐增强。第五，PMD-MEAN在7B模型上超越了GSPO（一种针对MoE稳定性优化的高级GRPO变体），在30B模型上取得相当性能（平均44.01 vs 43.96）。

**Table 1**: 总体评估分数（Avg@32），Staleness表示使用同一旧策略rollout的mini步数

**Table 2**: On-policy梯度和PMD-MEAN的效率对比，时间为每token毫秒数

**Figure 3**: PMD-MEAN和PMD-PART在有限rollout下的目标估计误差对比

**Figure 4**: PMD-MEAN和PMD-PART的训练曲线（Qwen2.5-7B，DAPO-Math-17k）

查看结构化数据

任务	指标	本文	基线	提升
AIME 2024 数学推理 (Qwen2.5-7B)	Avg@32	PMD-MEAN (τ=0.005): 19.69	GRPO: 17.08, On-policy: 18.65	相比GRPO +2.61绝对提升(+15.3%), 相比On-policy +1.04
AIME 2025 数学推理 (Qwen2.5-7B)	Avg@32	PMD-MEAN (τ=0.005): 19.48	GRPO: 10.52, On-policy: 18.33	相比GRPO +8.96绝对提升(+85.2%), 相比On-policy +1.15
AIME 2024 数学推理 (Qwen3-30B-A3B-Base)	Avg@32	PMD-MEAN (τ=0.1): 50.83	GRPO: 36.56	相比GRPO +14.27绝对提升(+39.0%)
AIME 2025 数学推理 (Qwen3-30B-A3B-Base)	Avg@32	PMD-MEAN (τ=0.1): 37.19	GRPO: 27.92	相比GRPO +9.27绝对提升(+33.2%)
训练效率对比 (Qwen2.5-7B)	ms/token	PMD-MEAN: 0.0126	On-policy: 0.0569	4.5倍加速，generation成本从0.0512降至0.0062 ms/token
AIME 2024 数学推理 vs GSPO (7B)	Avg@32	PMD-MEAN: 19.69	GSPO: 15.52	相比GSPO +4.17绝对提升(+26.9%)

局限与改进

本文存在若干值得注意的局限性。首先，理论分析基于简化假设：将LLM后训练形式化为contextual bandit问题，忽略了自回归生成中的token级序列决策结构，这使得分析更清晰但可能遗漏了序列依赖性带来的复杂性。其次，收敛分析依赖于可实现性假设（$\pi^\star_{t+1} \in \Pi$，即理想目标在策略类中）和有限策略类假设，这些在神经网络参数化的深度策略类中难以严格满足。第三，实验仅在数学推理任务上验证，使用二元正确性奖励，未探索更复杂的奖励信号（如人类偏好、代码执行反馈）或多步RL场景。第四，PMD-MEAN的均值近似引入了一个不可消除的系统性偏差：正动作的目标被系统性地高估（命题4.11），虽然引理4.12表明这个偏差在大 $n$ 极限下可以通过归一化约束消除，但在有限样本下它仍然影响训练动态。第五，作者自己也承认，本文聚焦于PMD-MEAN的基本形式以保持理论分析的清晰性，未考虑过采样策略和重要性采样校正等实际工程技巧，这些可能进一步提升性能但会增加算法复杂度。

独立分析的弱点

从独立分析的角度来看，PMD-MEAN存在几个值得改进的弱点。第一，均值近似的系统性偏差问题：对于正动作，估计目标 $\hat{s}^\star_{-i}$ 系统性地大于理想目标 $s^\star$，导致对这些动作的改进过于激进（论文公式25附近讨论），这在奖励通过率极低的训练初期可能导致对少数正确响应的过度拟合。改进方向可以是引入偏差校正项或使用leave-one-out估计来直接减少这一偏差。第二，$\tau$ 的选择是全局统一的，但论文自身分析（命题4.7-4.8和引理4.12）表明最优的 $\tau$ 应该与每个prompt的通过率 $p_t$ 自适应相关——$p_t$ 小时需要更强的正则化，$p_t$ 大时可以更激进。作者在4.2.4节末尾也提出了这一点作为未来方向。第三，二元奖励假设限制了理论的适用范围，真实场景中的奖励往往是连续值或多维的，Lambert-W解在这种情况下可能有更复杂的行为。第四，论文未讨论PMD-MEAN对超参数（特别是采样温度和group size）的敏感性，这些在实践中对性能有重要影响。

未来方向

本文为多个有前景的研究方向奠定了基础。作者明确提出的方向包括：自适应正则化方案，即根据每个prompt的通过率 $p_t$ 动态调整 $\tau$，引理4.12的分析表明这可以获得更好的误差-收敛速率权衡。将PMD-MEAN与过采样策略和重要性采样校正相结合，以处理训练/推理引擎不匹配的问题。基于本文的理论框架，未来可以探索：将KL-χ²混合正则化的自适应机制推广到其他RL算法（如PPO、GRPO），设计出具有理论保证的自适应正则化调度器；将分析扩展到多步RL和token级决策场景，研究序列依赖性如何影响隐式正则化效果；在更多样化的任务和奖励类型上验证PMD-MEAN的有效性，包括代码生成、工具使用和长链推理任务；探索PMD-MEAN与DAPO等高级优化技巧的结合，进一步提升大规模模型训练的效率和稳定性。

复现评估

本文在可复现性方面提供了较好的支持。代码已开源在 https://github.com/horizon-rl/OpenKimi，基于verl框架实现，这是当前LLM RL训练的主流开源框架之一。实验设置明确：DAPO-Math-17k数据集公开可用，基础模型Qwen2.5-7B和Qwen3-30B-A3B-Base均可获取。训练超参数（全局批大小512、组大小16、采样温度1、最大响应长度8192或20480、mini-batch大小32、学习率 $1 \times 10^{-6}$）都有明确记录。然而，复现所需的算力门槛较高：30B MoE模型训练300步、7B模型训练495步，且每个评估问题采样32个解，需要相当规模的GPU集群。此外，论文提到的Figure 6（在4.2.3节引用）在正文中未完整展示，可能需要查看附录才能获得完整的实验细节。整体而言，有开源代码和明确的实验设置，复现难度主要在算力层面。

策略镜像下降中对数配分函数近似为LLM后训练引入隐式正则化 Approximation of Log-Partition Function in Policy Mirror Descent Induces Implicit Regularization for LLM Post-Training