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

• URL: http://arxiv.org/abs/2602.04417v1
• Date: Wed, 04 Feb 2026 10:50:17 GMT
• Title: EMA Policy Gradient: Taming Reinforcement Learning for LLMs with EMA Anchor and Top-k KL
• Abstract summary: We propose two techniques to improve gradient policy algorithms for Large Language Models (LLMs)<n>First, we replace the fixed anchor policy during RL with an Exponential Moving Average (EMA)<n>Second, we introduce Top-k KL estimator, which allows for flexible between exact KL and sampled KL.
• Score: 19.463033284177087
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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

Related papers

• fg-expo: Frontier-guided exploration-prioritized policy optimization via adaptive kl and gaussian curriculum [11.537163059885687]
  We propose FG-ExPO, short for Frontier-Guided Exploration-Prioritized Policy Optimization.<n> Accuracy-Conditioned KL Scaling (AKL) adjusts the KL penalty strength through a smooth nonlinear function of batch average accuracy.<n>We conduct evaluations on DeepSeek-R1-Distill-Qwen-1.5B and Qwen3-8B-Base across six mainstream mathematical reasoning benchmarks.
  arXiv  Detail & Related papers  (2026-05-12T01:48:48Z)
• TRACE: Distilling Where It Matters via Token-Routed Self On-Policy Alignment [20.277178104190536]
  On-policy self-distillation (self-OPD) densifies reinforcement learning with verifiable rewards (RLVR) by letting a policy teach itself under privileged context.<n>We propose Token-Routed Alignment for Critical rEasoning (TRACE), which distills only on annotator-marked critical spans.<n>Our analysis explains TRACE through two effects: forward KL provides non-vanishing lift to teacher-supported tokens that the student under-allocates, while span masking and decay keep cumulative privileged-gradient exposure finite.
  arXiv  Detail & Related papers  (2026-05-11T08:45:03Z)
• expo: Exploration-prioritized policy optimization via adaptive kl regulation and gaussian curriculum sampling [11.537163059885687]
  We propose Exploration-Prioritized Policy Optimization (EXPO) with two lightweight plug-in modules.<n>We conduct experiments on DeepSeek-R1-Distill-Qwen-1.5B and Qwen3-8B-Base over six mathematical reasoning benchmarks.<n>It obtains an absolute gain of 13.34 on AIME 2025 pass@32, rising from 63.33 percent to 76.67 percent, and achieves an average pass@32 improvement of 2.66 on the 8B model.
  arXiv  Detail & Related papers  (2026-05-11T03:19:04Z)
• FIPO: Eliciting Deep Reasoning with Future-KL Influenced Policy Optimization [84.58281577727566]
  We present Future-KL Influenced Policy Optimization (FIPO), a reinforcement learning algorithm designed to overcome reasoning bottlenecks in large language models.<n>FIPO addresses this by incorporating discounted future-KL divergence into the policy update, creating a dense advantage formulation that re-weights tokens based on their influence on subsequent trajectory behavior.<n> evaluated on Qwen2.5-32B, FIPO extends the average chain-of-thought length from roughly 4,000 to over 10,000 tokens and increases AIME 2024 Pass@1 accuracy from 50.0% to a peak of 58.0%.
  arXiv  Detail & Related papers  (2026-03-20T10:24:50Z)
• Dr. MAS: Stable Reinforcement Learning for Multi-Agent LLM Systems [20.971694319263353]
  We propose Dr. MAS, a simple and stable RL training recipe for multi-agent LLM systems.<n>Dr. MAS uses agent-wise remedy: normalizing advantages per agent using each agent's own reward statistics.<n>Dr. MAS provides an end-to-end RL training framework for multi-agent LLM systems.
  arXiv  Detail & Related papers  (2026-02-09T16:13:39Z)
• Achieving Logarithmic Regret in KL-Regularized Zero-Sum Markov Games [53.447182734351]
  We develop and analyze algorithms that provably achieve improved sample efficiency under Reverse Kullback-Leibler (KL) regularization.<n>We study both two-player zero-sum Matrix games and Markov games: for Matrix games, we propose OMG, an algorithm based on best response sampling with optimistic bonuses, and extend this idea to Markov games through the algorithm SOMG.<n>Both algorithms achieve a logarithmic regret in $T$ that scales inversely with the KL regularization strength $beta$ in addition to the standard $widetildemathcalO(sqrtT)
  arXiv  Detail & Related papers  (2025-10-15T01:00:54Z)
• FlowRL: Matching Reward Distributions for LLM Reasoning [69.88820066093798]
  We propose FlowRL: matching the full reward distribution via flow balancing instead of maximizing rewards in large language model (LLM) reinforcement learning (RL)<n>We transform scalar rewards into a normalized target distribution using a learnable partition function, and then minimize the reverse KL divergence between the policy and the target distribution.
  arXiv  Detail & Related papers  (2025-09-18T17:56:36Z)
• KL-Regularised Q-Learning: A Token-level Action-Value perspective on Online RLHF [1.8665975431697432]
  We develop a new action-value RL method for the LM-RLHF setting, KL-regularised Q-Learning (KLQ)<n>We show that our method is equivalent to a version of PPO in a certain specific sense, despite its very different motivation.<n>We demonstrate that KLQ performs on-par with PPO at optimising the LM-RLHF objective, and achieves a consistently higher win-rate against PPO on LLM-as-a-judge evaluations.
  arXiv  Detail & Related papers  (2025-08-23T11:50:54Z)
• GRPO-CARE: Consistency-Aware Reinforcement Learning for Multimodal Reasoning [53.894789613838654]
  We introduce SEED-Bench-R1, a benchmark with complex real-world videos requiring balanced perception and reasoning.<n>Using SEED-Bench-R1, we find that standard GRPO, while improving answer accuracy, often reduces logical coherence between reasoning steps and answers, with only a 57.9% consistency rate.<n>We propose GRPO-CARE, a consistency-aware RL framework optimizing both answer correctness and reasoning coherence without explicit supervision.
  arXiv  Detail & Related papers  (2025-06-19T08:49:13Z)
• On the Design of KL-Regularized Policy Gradient Algorithms for LLM Reasoning [59.11784194183928]
  Policy gradient algorithms have been successfully applied to enhance the reasoning capabilities of large language models (LLMs)<n>Regularized Policy Gradient (RPG) view shows that the widely-used $k_3$ penalty is exactly the unnormalized KL.<n>RPG-REINFORCE with RPG-Style Clip improves accuracy by up to $+6$ absolute percentage points over DAPO.
  arXiv  Detail & Related papers  (2025-05-23T06:01:21Z)
• $Q\sharp$: Provably Optimal Distributional RL for LLM Post-Training [63.602824642605775]
  $Qsharp$ is a value-based algorithm for KL-regularized RL that guides the reference policy using the optimal regularized $Q$ function.<n>Our results highlight $Qsharp$ as an effective approach for post-training LLMs, offering both improved performance and theoretical guarantees.
  arXiv  Detail & Related papers  (2025-02-27T21:43:00Z)
• Iterative Nash Policy Optimization: Aligning LLMs with General Preferences via No-Regret Learning [55.65738319966385]
  We propose a novel online algorithm, iterative Nash policy optimization (INPO)<n>Unlike previous methods, INPO bypasses the need for estimating the expected win rate for individual responses.<n>With an LLaMA-3-8B-based SFT model, INPO achieves a 42.6% length-controlled win rate on AlpacaEval 2.0 and a 37.8% win rate on Arena-Hard.
  arXiv  Detail & Related papers  (2024-06-30T08:00:34Z)
• WARP: On the Benefits of Weight Averaged Rewarded Policies [66.95013068137115]
  We introduce a novel alignment strategy named Weight Averaged Rewarded Policies (WARP) WARP merges policies in the weight space at three distinct stages. Experiments with GEMMA policies validate that WARP improves their quality and alignment, outperforming other open-source LLMs.
  arXiv  Detail & Related papers  (2024-06-24T16:24:34Z)
• Generalized Munchausen Reinforcement Learning using Tsallis KL Divergence [22.400759435696102]
  We investigate a generalized KL divergence, called the Tsallis KL divergence, which use the $q$-logarithm in the definition. We characterize the types of policies learned under the Tsallis KL, and motivate when $q >1$ could be beneficial. We show that this generalized MVI($q$) obtains significant improvements over the standard MVI($q = 1$) across 35 Atari games.
  arXiv  Detail & Related papers  (2023-01-27T00:31:51Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.