策略镜像下降中的对数配分函数近似引发 LLM 后训练中的隐式正则化
机器学习
2026-02-06 v1
摘要
策略镜像下降(PMD)为强化学习(RL)提供了一个原则框架，通过迭代求解 KL 正则化策略改进子问题。虽然该方法已被采用于训练高级 LLM 如 Kimi K1.5/K2，但理想的闭式 PMD 更新需要可靠的配分函数估计，而在 LLM 庞大的动作空间中进行有限抽样时，这是一大挑战。我们研究了一种实用算法，即 PMD-mean，通过对采样策略下的平均奖励来近似对数配分项，并在对数策略空间中进行回归。具体而言，我们刻画了 PMD-mean 的 population 解，表明它隐式地对镜像下降子问题进行优化，采用自适应混合 KL-- 正则化器。这种额外的 正则化约束了大概率变动，使得在预期奖励较低时产生更保守的更新，从而增强了对有限样本估计误差的鲁棒性。在数学推理任务上的实验表明，PMD-mean 在性能、稳定性和时间效率方面均优于基线方法。这些发现深化了对 PMD-mean 的理解，并阐明了针对 LLM 强化学习算法的原则性改进之路。代码已公开于 https://github.com/horizon-rl/OpenKimi。
引用
@article{arxiv.2602.05933,
title = {Approximation of Log-Partition Function in Policy Mirror Descent Induces Implicit Regularization for LLM Post-Training},
author = {Zhenghao Xu and Qin Lu and Changlong Yu and Tuo Zhao},
journal= {arXiv preprint arXiv:2602.05933},
year = {2026}
}