用EMA和Top-kKL改进强化学习,让大模型推理更稳定高效
EMA Policy Gradient: Taming Reinforcement Learning for LLMs with EMA Anchor and Top-k KL
- 用EMA动态更新策略锚点,提升训练稳定性
- 在数学推理上达53.9%准确率,搜索问答任务平均提升33.3%
- 适合想提升大模型强化学习性能的研究者与工程师
强化学习使大语言模型逐步具备复杂推理与智能体行为。本文提出两种简单有效的方法改进策略梯度算法:一是用指数移动平均(EMA)替代固定锚点策略,类似深度Q学习中的目标网络;二是引入Top-k KL估计器,可在精确KL与采样KL间灵活插值。我们推导了使用EMA锚点的稳定性条件,并证明所提方法在任意k下均能保持无偏的KL值与梯度,同时兼具精确KL优势。结合GRPO后,该方法(EMA-PG)性能显著提升:在数学推理任务中,经微调的Qwen-1.5B模型在OlympiadBench上达到53.9%,高于GRPO的50.8%;在7个需搜索引擎的问答任务中,基于Qwen-3B的EMA-PG平均提升33.3%,包括HotpotQA从29.7%升至44.1%,2WikiMultiHopQA从27.4%升至40.1%。结果表明,EMA-PG是一种简单、原理清晰且强大的大模型强化学习扩展方案。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。