用均值近似对数归一化项,让大模型强化学习更稳定高效。
Approximation of Log-Partition Function in Policy Mirror Descent Induces Implicit Regularization for LLM Post-Training
- 用采样策略的平均奖励近似对数归一化项,简化计算。
- 实验显示在数学推理任务上性能更好,收敛更快且更稳定。
- 适合研究大模型强化学习优化与鲁棒性改进的开发者。
策略镜面下降(PMD)通过迭代求解带KL正则的策略改进子问题,为强化学习提供严谨框架。尽管该方法已用于训练如Kimi K1.5/K2等先进大模型,但理想闭式更新需可靠的归一化函数估计,在大模型动作空间和有限回放条件下极具挑战。本文提出实用算法PMD-mean,用采样策略下的平均奖励近似对数归一化项,并在对数策略空间中进行回归。我们刻画了PMD-mean的总体解,证明其隐式优化了带有自适应混合KL-χ²正则的镜面下降子问题。额外的χ²正则限制了概率的大幅变化,在期望奖励较低时产生更保守的更新,提升了对有限样本估计误差的鲁棒性。数学推理任务上的实验表明,PMD-mean在性能、稳定性与时间效率方面均表现更优。这些发现深化了对PMD-mean的理解,为大模型强化学习算法的原理性改进提供了新路径。代码见https://github.com/horizon-rl/OpenKimi。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。