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

TL;DR

提出EMA锚点与Top-k KL,提升LLMs强化学习性能,达成53.9%数学推理准确率。

cs.LG 🔴 高级 2026-02-04 67 次浏览
Lunjun Zhang Jimmy Ba
强化学习 大规模语言模型 策略梯度 KL估计 模型稳定性

核心发现

方法论

本文提出结合指数移动平均(EMA)锚点与Top-k KL估计器的策略梯度算法,改进LLMs的RL训练。EMA类似深度Q学习中的目标网络,用于平滑策略参数,提升训练稳定性。Top-k KL在保持偏差无误的同时,灵活调节精确KL与采样KL的权衡,通过对logits的top-k索引进行部分精确计算,有效降低内存消耗。算法在GRPO基础上实现,结合token级别的KL估计与EMA锚点,优化数学推理和代理任务中的表现。

关键结果

  • 在数学推理任务中,R1蒸馏Qwen-1.5B模型在OlympiadBench达成53.9%的准确率,相较GRPO的50.8%显著提升。
  • 在Q&A搜索引擎任务中,Qwen-3B模型通过EMA-PG平均提升33.3%,HotpotQA由29.7%跃升至44.1%,2WikiMultiHopQA由27.4%提升至40.1%。
  • Top-k KL估计器在保持偏差无误的同时,显著降低梯度误差,增强样本效率,尤其在长序列推理和大规模词汇空间中表现优越。

研究意义

该研究突破了RL在大规模语言模型中的应用瓶颈,通过引入EMA锚点与Top-k KL,显著提升训练稳定性与样本效率,推动LLMs在复杂推理和代理任务中的性能极限。其理论分析与实证验证,为未来RL优化提供了新范式,有助于模型在多样化应用场景中实现更高水平的自主学习与推理能力。

技术贡献

技术创新在于引入EMA目标网络以增强训练稳定性,结合Top-k KL估计器实现偏差无误的高效KL计算,解决传统采样估计的偏差与高方差问题。提出的理论稳定性条件为大规模RL训练提供数学保障,且算法兼容多种策略梯度方法,具有广泛适用性。此方法在推理和代理任务中均展现出优越性能,推动RL在LLMs中的应用边界。

新颖性

首次在大规模LLMs中系统引入EMA锚点与Top-k KL估计器,结合token级别的KL优化策略,突破了序列级别KL的限制。创新点在于偏差无误的高效KL估计与动态锚点机制的结合,为RL训练中的稳定性与样本效率提供新思路,区别于现有仅采用采样或序列级别KL的方案。

局限性

  • 方法在极端长序列或极大词汇空间中仍面临内存与计算瓶颈,尤其在多模态任务中适应性待验证。
  • EMA锚点参数η的选择对训练稳定性影响显著,缺乏自适应调节机制,可能在不同任务中表现不一致。
  • 当前实验主要集中在文本推理与问答任务,泛化到其他复杂任务(如多模态、多任务学习)仍需进一步验证。

未来方向

未来将探索自适应EMA参数调节机制,结合多模态数据优化KL估计器,扩展算法在多任务、多模态场景中的应用。同时,结合更高效的内存管理策略,推动算法在超大模型中的实用化,提升模型自主推理与交互能力。

AI 总览摘要

随着大规模语言模型(LLMs)在复杂推理和自主行为方面的突破,强化学习(RL)成为提升模型性能的重要工具。然而,传统RL方法在训练稳定性和样本效率方面仍面临挑战。本文提出EMA策略梯度(EMA-PG),结合指数移动平均(EMA)锚点与Top-k KL估计器,有效缓解训练中的不稳定性和偏差问题。

EMA锚点通过平滑目标网络参数,类似深度Q学习中的目标网络机制,提升训练的稳定性和收敛速度。Top-k KL估计器则在保持偏差无误的同时,动态调节KL的计算范围,显著降低内存消耗和梯度误差。这两项技术结合,使得在数学推理和代理任务中的RL性能得到大幅提升。

实验证明,R1蒸馏Qwen-1.5B模型在OlympiadBench达到53.9%的准确率,超越传统GRPO的50.8%。在Q&A搜索任务中,Qwen-3B模型通过EMA-PG平均提升33.3%,HotpotQA由29.7%跃升至44.1%,2WikiMultiHopQA由27.4%提升至40.1%。这些结果验证了方法在长序列推理和大词汇空间中的优越性。

该研究不仅提供了理论分析的稳定性条件,还为RL在LLMs中的应用提供了新范式。未来,结合自适应参数调节和多模态扩展,有望推动模型在更复杂场景中的自主学习和推理能力,开启大规模RL训练的新篇章。

深度分析

研究背景

近年来,LLMs在自然语言处理中的表现持续提升,尤其在预训练和微调阶段取得显著成果。RL在模型对齐、推理和自主行为中扮演关键角色,代表性方法如RLHF、GRPO等,推动模型在复杂任务中的表现。然而,RL训练中的不稳定性、样本效率不足和高内存需求仍是瓶颈。深度Q学习中的目标网络和EMA机制已被证明能缓解训练波动,但在LLMs中的应用尚未充分探索。KL正则化作为策略优化的重要组成部分,其估计方式直接影响训练效果。传统的序列级别KL在长序列任务中表现有限,而token级别KL逐渐成为主流。本文基于此背景,提出结合EMA锚点与Top-k KL的创新方法,旨在解决现有方法的不足,推动RL在大规模语言模型中的广泛应用。

核心问题

当前RL在LLMs中的应用面临多重挑战,包括训练不稳定、样本效率低、内存消耗大以及KL估计偏差。传统策略梯度方法如REINFORCE和PPO在大模型上易出现梯度震荡,影响收敛速度。序列级别KL在长推理链中难以提供有效的梯度信号,而token级别KL虽更细粒度,但存在偏差和高方差问题。此外,模型训练中引入的目标网络(如EMA)虽能稳定训练,但缺乏理论指导,参数选择困难。解决这些问题,需在保持偏差无误的同时,降低内存和计算成本,提升训练稳定性和样本效率,成为当前研究的核心难题。

核心创新

本研究的创新点在于:一是引入EMA锚点机制,通过平滑目标网络参数,增强训练稳定性,减少震荡;二是提出Top-k KL估计器,结合部分精确计算与采样,偏差无误且高效,适应大词汇空间;三是分析并建立了EMA训练的稳定性条件,为理论提供保障。这些创新区别于传统仅用采样估计或序列级别KL的方案,提供了更稳健、更高效的训练策略,显著提升模型在推理和代理任务中的性能。

方法详解

  • �� 采用策略梯度框架,结合奖励函数和KL正则化,优化模型参数。
  • �� 引入EMA目标网络,利用指数移动平均平滑参数,作为锚点,提升训练稳定性。
  • �� 设计Top-k KL估计器,在每个时间步只对logits的前k个进行精确计算,剩余部分用采样补偿,确保偏差无误。
  • �� 通过分析训练动态,推导EMA参数的稳定性条件,确保训练收敛。
  • �� 在数学推理和问答任务中,结合具体数据集(如OlympiadBench、HotpotQA)进行实证验证。
  • �� 对比传统采样KL和序列级别KL,验证新方法在样本效率和性能上的优势。

实验设计

实验采用Math Reasoning和Agentic RL两大场景,使用Qwen-1.5B和Qwen-3B模型,基于40k问答对的RL数据集进行训练。评估指标包括准确率、成功率等,比较GRPO、WARP等基线。重点验证EMA锚点对训练稳定性和性能提升作用,以及Top-k KL在不同k值下的效果。通过 ablation 研究,分析参数η、k值对模型性能的影响。多任务、多数据集的实验确保方法的普适性和鲁棒性。

结果分析

在数学推理任务中,R1蒸馏Qwen-1.5B模型在OlympiadBench达成53.9%的准确率,超越GRPO的50.8%。在HotpotQA中由29.7%提升至44.1%,在2WikiMultiHopQA由27.4%跃升至40.1%。Qwen-3B模型在搜索问答中,平均提升33.3%。Top-k KL估计器在偏差无误的基础上,显著降低梯度误差,增强样本效率,尤其在长序列和大词汇空间中表现优越。这些结果验证了算法在复杂推理和多任务场景中的优越性。

应用场景

该方法适用于需要高效推理和自主学习的LLMs,如智能问答、自动推理、对话系统。结合强化学习,可以在模型微调、对齐、任务适应中实现更快收敛和更优性能。未来还可扩展到多模态任务,提升模型在视觉、语音等多模态场景中的自主推理能力。

局限与展望

当前算法在极长序列或超大词汇空间中仍存在内存和计算瓶颈,模型参数选择对训练效果影响大,缺乏自适应机制。实验主要集中在文本推理和问答,泛化到多模态、多任务场景仍需验证。未来需优化内存管理和参数调节策略,以实现更广泛的应用。

通俗解读 非专业人士也能看懂

想象你在一家工厂里,工人们每天都要完成各种任务。有些任务很复杂,需要很多步骤,而且每次都可能出错。为了保证工厂正常运转,管理者会用一种特别的方法:他们会用一份“平滑的计划”来指导工人,这份计划会不断更新,确保工人们不会偏离目标太远。这就像论文中的EMA锚点,用来让模型在训练时更稳定。

另外,工厂里有很多原材料(像词汇一样),但大部分原材料都集中在少数几种上。工厂只关注这些“重要的几种材料”,而对其他的少量材料只做简单的估算。这就是Top-k KL的思想,只关注最重要的部分,既节省资源,又保证效果。

通过这些方法,工厂的生产效率大大提高,出错率降低,生产速度加快。类似的,论文中的算法让大模型在学习复杂任务时更稳定、更高效,能更快地掌握知识和推理能力。

简单解释 像给14岁少年讲一样

想象你在学校里学新东西,老师会给你一些练习题,帮助你理解。可是,有时候题目太多,你不知道该专注哪一部分。这个论文就像是老师用一种聪明的方法帮你安排学习:他会用一份“特别的笔记”来提醒你重点内容,而且会不断更新这个笔记,让你不会迷失方向。

还有,老师知道你大脑里最重要的知识点集中在几块区域,所以他只让你多练这几块,其他的部分只简单提一下,这样既节省时间,又能学得更好。这就像Top-k KL,只关注最重要的内容。

用这些方法,你学习得更快,考试也能拿更高分。论文里的技术也是一样,让大模型在学习复杂任务时变得更稳定、更聪明,能更快理解和推理,就像你变成了学习高手!

原文摘要

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

cs.LG cs.AI