EMA Policy Gradient: Taming Reinforcement Learning for LLMs with EMA Anchor and Top-k KL

📄 arXiv: 2602.04417v1 📥 PDF

作者: Lunjun Zhang, Jimmy Ba

分类: cs.LG, cs.AI

发布日期: 2026-02-04

🔗 代码/项目: GITHUB

💡 一句话要点

EMA-PG：通过EMA锚定和Top-k KL提升LLM强化学习的稳定性和性能

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 强化学习 大型语言模型 策略梯度 EMA锚定 KL散度估计

📋 核心要点

  1. 现有LLM强化学习方法在训练过程中存在不稳定性，难以充分利用数据。
  2. 采用EMA锚定策略和Top-k KL估计器，提高策略梯度算法的稳定性和数据利用率。
  3. 实验表明，EMA-PG在数学推理和agentic RL任务上均显著优于现有方法。

📝 摘要（中文）

本文提出了两种简单技术，以改进大型语言模型（LLM）的策略梯度强化学习算法。首先，类似于深度Q学习中的目标网络，我们将强化学习期间的固定锚定策略替换为指数移动平均（EMA）。其次，我们引入了Top-k KL估计器，它允许在精确KL散度和采样KL散度之间进行灵活的插值。我们推导了使用EMA锚定的稳定性条件；此外，我们证明了我们的Top-k KL估计器在任何k下都能产生无偏的KL值和无偏的梯度，同时带来精确KL的好处。当与GRPO结合使用时，这两种技术（EMA-PG）可以显著提高性能。在数学推理方面，它使R1-distilled Qwen-1.5B在OlympiadBench上达到53.9％，而GRPO为50.8％。在agentic RL领域，使用Qwen-3B作为基础模型，EMA-PG在7个搜索引擎问答数据集中平均提高了GRPO 33.3％，包括HotpotQA上的29.7％→44.1％，2WikiMultiHopQA上的27.4％→40.1％。总而言之，我们表明EMA-PG是一种简单、有原则且功能强大的方法，可以扩展LLM的强化学习。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）强化学习训练过程中的不稳定性和数据利用率不足的问题。传统的策略梯度方法在更新策略时，通常使用一个固定的锚定策略来计算KL散度，这可能导致策略更新过于激进或保守，从而影响训练的稳定性和收敛速度。此外，采样KL散度估计可能引入偏差，而精确KL散度计算成本高昂。

核心思路：论文的核心思路是引入EMA（指数移动平均）锚定策略和Top-k KL估计器，以提高策略梯度算法的稳定性和数据利用率。EMA锚定策略通过对历史策略进行加权平均，可以更平滑地更新锚定策略，从而减少策略更新的方差。Top-k KL估计器则通过选择KL散度最大的前k个样本进行计算，可以在保证无偏性的同时，降低计算成本。

技术框架：EMA-PG算法的整体框架与标准的策略梯度算法类似，主要包括以下几个步骤：1. 使用当前策略生成样本；2. 计算每个样本的奖励；3. 使用EMA更新锚定策略；4. 使用Top-k KL估计器计算KL散度；5. 使用策略梯度算法更新策略。其中，EMA锚定策略和Top-k KL估计器是EMA-PG算法的核心组成部分。

关键创新：论文最重要的技术创新点在于提出了EMA锚定策略和Top-k KL估计器。EMA锚定策略通过对历史策略进行加权平均，可以更平滑地更新锚定策略，从而减少策略更新的方差。Top-k KL估计器则通过选择KL散度最大的前k个样本进行计算，可以在保证无偏性的同时，降低计算成本。与现有方法的本质区别在于，EMA-PG算法可以更稳定、更高效地训练LLM。

关键设计：EMA锚定策略的关键参数是EMA的衰减系数，该系数控制了历史策略对当前锚定策略的影响程度。Top-k KL估计器的关键参数是k，该参数控制了用于计算KL散度的样本数量。论文中，作者通过实验验证了不同参数设置对算法性能的影响，并给出了推荐的参数设置。

📊 实验亮点

实验结果表明，EMA-PG算法在数学推理和agentic RL任务上均显著优于现有方法。在OlympiadBench上，使用R1-distilled Qwen-1.5B模型，EMA-PG算法达到了53.9％的准确率，而GRPO算法为50.8％。在agentic RL领域，使用Qwen-3B模型，EMA-PG算法在7个搜索引擎问答数据集中平均提高了GRPO 33.3％，包括HotpotQA上的29.7％→44.1％，2WikiMultiHopQA上的27.4％→40.1％。

🎯 应用场景

EMA-PG算法可应用于各种需要利用强化学习来训练LLM的场景，例如对话系统、文本生成、智能体决策等。该算法能够提高LLM的性能和稳定性，使其能够更好地完成各种复杂的任务。该研究的实际价值在于降低了LLM强化学习的训练难度，加速了LLM在实际应用中的部署。未来，EMA-PG算法可以进一步扩展到其他类型的模型和任务中。

📄 摘要（原文）

Reinforcement Learning (RL) has enabled Large Language Models (LLMs) to acquire increasingly complex reasoning and agentic behaviors. In this work, we propose two simple techniques to improve policy gradient algorithms for LLMs. First, we replace the fixed anchor policy during RL with an Exponential Moving Average (EMA), similar to a target network in deep Q-learning. Second, we introduce Top-k KL estimator, which allows for flexible interpolation between exact KL and sampled KL. We derive the stability conditions for using EMA anchor; moreover, we show that our Top-k KL estimator yields both unbiased KL values and unbiased gradients at any k, while bringing the benefits of exact KL. When combined with GRPO, the two techniques (EMA-PG) lead to a significant performance boost. On math reasoning, it allows R1-distilled Qwen-1.5B to reach 53.9% on OlympiadBench compared to 50.8% by GRPO. On agentic RL domains, with Qwen-3B base, EMA-PG improves GRPO by an average of 33.3% across 7 datasets of Q&A with search engines, including 29.7% $\rightarrow$ 44.1% on HotpotQA, 27.4% $\rightarrow$ 40.1% on 2WikiMultiHopQA. Overall, we show that EMA-PG is a simple, principled, and powerful approach to scaling RL for LLMs. Code: https://github.com/LunjunZhang/ema-pg