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

The evolution of large language models (LLMs) has been marked by pretraining on massive corpora, followed by fine-tuning and alignment techniques such as RLHF. Despite these advances, RL training remains challenging due to instability, high variance, and memory constraints, especially when optimizing complex objectives like reasoning or agentic behaviors. Prior work like PPO, GRPO, and WARP has made progress, but issues like biased KL estimation and inefficient sampling persist. Deep Q-learning’s target network and EMA mechanisms have demonstrated success in stabilizing training in reinforcement learning, inspiring their adaptation to NLP. Additionally, the importance of accurate KL regularization for model alignment and robustness has driven research into better estimators, including sequence-level and token-level approaches. However, existing methods often trade bias for variance or require prohibitive memory. This paper builds on these foundations, proposing a unified framework that leverages EMA smoothing and adaptive KL estimation to overcome these limitations, aiming for more stable, sample-efficient RL training for large-scale language models.

The core challenge in RL for LLMs lies in balancing training stability, sample efficiency, and computational feasibility. Sequence-level KL regularization, while straightforward, hampers long-sequence reasoning due to its bias and high memory cost. Token-level KL offers finer control but introduces bias and high variance in estimators like K1, K2, and K3. Moreover, existing estimators either do not provide unbiased gradients or are memory-intensive, limiting scalability. The instability caused by fluctuating targets and biased regularization leads to slow convergence and suboptimal performance. Addressing these issues requires a method that ensures unbiased, low-variance KL estimation, stable training dynamics, and manageable memory footprint, especially for reasoning tasks involving long sequences and large vocabularies.

The paper introduces three key innovations: 1) EMA anchor network, which smooths policy parameters over training steps, stabilizing updates; 2) Top-k KL estimator, which computes exact KL on top-k logits and estimates the tail via sampling, maintaining unbiasedness with lower variance; 3) a rigorous stability analysis deriving conditions on learning rate, EMA decay, and Fisher information for convergence. These innovations enable bias-free, memory-efficient KL regularization, improving training stability and sample efficiency. Unlike prior methods relying solely on sampling or sequence-level KL, this approach adaptively interpolates between exact and sampled estimates, providing theoretical guarantees and practical robustness. The combined framework significantly enhances RL performance on reasoning and agentic tasks.

• �� Implement an EMA target network θema, updating via exponential smoothing to serve as a stable reference.
• �� Design Top-k KL estimators that compute exact KL on the top-k logits, with the remaining probability mass estimated through importance sampling.
• �� Derive stability conditions for EMA parameters, involving Fisher information and learning rate, ensuring convergence.
• �� Incorporate token-level KL regularization into policy gradients, replacing sequence-level KL.
• �� Use theoretical analysis to confirm unbiasedness of estimators and stability of training dynamics.
• �� Validate on reasoning benchmarks (OlympiadBench) and search-based QA datasets, comparing against baseline algorithms like GRPO.
• �� Conduct ablation studies on k, EMA decay η, and estimator variants to optimize performance.
• �� Demonstrate improvements in accuracy, success rate, and training stability across tasks.

The experimental setup involves training Qwen-1.5B and Qwen-3B models on RL datasets with 40k question-answer pairs, focusing on math reasoning and agentic search tasks. Evaluation metrics include accuracy on OlympiadBench, success rates on HotpotQA and 2WikiMultiHopQA, and success rate improvements over baseline GRPO. The experiments compare EMA-PG with standard RL algorithms, analyzing the impact of EMA decay η, Top-k parameter k, and different KL estimators. Extensive ablation studies assess the stability conditions, estimator bias, and sample efficiency. The models are trained with hyperparameters tuned for stability, and results are statistically validated to confirm significant improvements. The setup emphasizes real-world reasoning benchmarks and multi-task scenarios to demonstrate robustness.

Empirical results show that EMA-PG achieves 53.9% accuracy on OlympiadBench, outperforming GRPO’s 50.8%. In multi-hop QA tasks, HotpotQA success rate increases from 29.7% to 44.1%, and 2WikiMultiHopQA from 27.4% to 40.1%. The success rate improvements are consistent across datasets, with an average boost of 33.3% in search-based agentic tasks. The Top-k KL estimators demonstrate lower gradient variance and unbiasedness, especially with moderate k values, balancing bias and variance effectively. The stability analysis confirms that proper tuning of EMA decay η and learning rate α ensures convergence, validating the theoretical conditions derived. Ablation on k shows optimal performance at intermediate values, confirming the estimator’s flexibility and efficiency.

The proposed EMA-PG framework is suitable for training large-scale language models in tasks requiring reasoning, dialogue, and autonomous decision-making. It can be integrated into existing RL pipelines for model alignment, fine-tuning, and multi-task learning, especially where long sequences and large vocabularies are involved. The method enables more stable training with fewer samples, reducing computational costs and improving robustness. Long-term, this approach can facilitate development of autonomous agents capable of complex reasoning, multi-modal understanding, and real-time interaction, impacting industries like virtual assistants, education, and AI research.

Despite its advantages, the method still faces challenges with extremely long sequences or very large vocabularies, where memory and computation remain bottlenecks. The choice of EMA decay η and k requires careful tuning, lacking adaptive mechanisms. The experiments focus mainly on NLP tasks, and its effectiveness in multi-modal or real-world deployment scenarios needs further validation. Additionally, reliance on logits for exact KL may limit applicability in black-box or privacy-sensitive settings. Future work should address these limitations by developing adaptive algorithms, more memory-efficient estimators, and broader application testing.