arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2602.04417cs.LGcs.AI

EMA策略梯度:通过EMA锚点和Top-k KL驯服LLM的强化学习

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

University of Toronto(多伦多大学)

Lunjun Zhang, Jimmy Ba

AI总结:

EMA-PG通过EMA锚点和Top-k KL估计器提升LLM强化学习性能,显著提高数学推理和代理任务表现

摘要:

强化学习(RL)使大型语言模型(LLMs)能够获得日益复杂的推理和代理行为。在本工作中,我们提出两种简单的技术来改进LLMs的策略梯度算法。首先,我们将RL中的固定锚定策略替换为指数移动平均(EMA),类似于深度Q学习中的目标网络。其次,我们引入Top-k KL估计器,这允许在精确KL和采样KL之间进行灵活的插值。我们推导了使用EMA锚点的稳定性条件;此外,我们证明了我们的Top-k KL估计器在任何k值下都能产生无偏的KL值和无偏的梯度,同时带来精确KL的优势。当与GRPO结合时,这两种技术(EMA-PG)显著提升了性能。在数学推理方面,它使R1蒸馏的Qwen-1.5B在OlympiadBench上达到53.9%,比GRPO的50.8%有所提升。在代理RL领域,使用Qwen-3B基础模型时,EMA-PG在7个QA与搜索引擎的数据库上平均提高了33.3%,包括在HotpotQA上从29.7%提升到44.1%,在2WikiMultiHopQA上从27.4%提升到40.1%。总体而言,我们展示了EMA-PG是一种简单、原则性且强大的方法,用于扩展LLMs的强化学习。代码:https://github.com/LunjunZhang/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

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑