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

Motivation

• Known: RL은 LLM의 reasoning 및 agentic 능력 향상에 핵심적인 역할을 하며, GRPO와 같은 policy gradient 알고리즘이 널리 쓰인다. KL regularization은 trust-region 유지, reward hacking 방지, base model 능력 보존을 위해 사용되며, K1/K2/K3 같은 token-level sampled KL estimator가 흔히 쓰인다.
• Gap: 기존 anchor policy는 학습 내내 고정되어 있어 신선한 supervision을 제공하지 못하고, prior work(TR-DPO, WARP)의 EMA anchor는 sequence-level KL만 다뤄 reasoning/agentic 설정에 부적합하다. 또한 기존 K1/K2/K3 KL estimator는 KL 값과 gradient를 동시에 unbiased하게 추정하지 못하고, exact KL은 계산량이 vocabulary 크기에 비례해 메모리 부담이 크다.
• Why: RL을 통한 LLM의 reasoning 및 agentic 능력 확장이 중요한 연구 방향이 된 상황에서, 알고리즘의 근본 구성요소인 anchor policy와 KL estimator를 개선하면 다양한 policy gradient 알고리즘(REINFORCE, PPO, GRPO, CISPO 등)에 폭넓게 적용 가능한 원리적이고 실용적인 성능 향상을 가져올 수 있기 때문이다.
• Approach: deep Q-learning의 target network처럼 EMA로 anchor policy를 갱신하는 EMA anchor와, top-k logit 인덱스에서는 exact KL을, 나머지에서는 sampled KL을 결합해 임의의 k에서 unbiased KL 값과 gradient를 모두 보장하는 Top-k KL estimator를 제안하고, 이를 GRPO에 결합한 EMA-PG를 제시한다.

Achievement

Figure 5. EMA-PG not only learns faster and but reaches higher asymptotic performance on agentic tasks of Q&A with searc

1. EMA anchor의 안정성 조건 도출: 고정 anchor 대신 EMA 기반 anchor policy를 사용할 때의 stability condition을 이론적으로 도출했다.
2. Unbiased Top-k KL estimator 설계: TopkReverseKL, TopkForwardKL 등 임의의 k에서 KL 값과 gradient 모두 unbiased인 estimator(K3++, K4, K5 포함)를 구성하고 기존 K1/K2/K3의 한계를 규명했다.
3. 수학 reasoning 성능 향상: R1-distilled Qwen-1.5B 기준 OlympiadBench에서 GRPO의 50.8% 대비 EMA-PG는 53.9%를 달성했다.
4. Agentic search RL 성능 향상: Qwen-3B base로 검색엔진 기반 7개 Q&A 데이터셋에서 평균 33.3% 성능 향상을 보였으며, HotpotQA는 29.7%→44.1%, 2WikiMultiHopQA는 27.4%→40.1%, Bamboogle은 12.8%→33.1%로 향상되었다.
5. 학습 속도 및 최종 성능 동시 개선: EMA-PG가 더 빠르게 학습할 뿐 아니라 더 높은 asymptotic performance에 도달함을 보였다.

How

Figure 2. Top-k KL computes exact KL on the top-k indices, and adds a masked sampled KL to keep the estimator unbiased.

• 기존 RL 목적함수의 KL 항을 고정 πref 대신 EMA로 업데이트되는 θema (θema ← η·θema + (1−η)·θ) 기반 policy로 대체
• token-level KL을 sequence-level KL 대신 사용하는 것이 reasoning/agentic 세팅에 더 적합함을 이론적·실증적으로 분석
• K1, K2, K3 등 기존 sampled KL estimator가 KL 값과 gradient의 unbiasedness를 동시에 만족하지 못함을 규명하고, 이를 만족하는 K3++, K4, TopkReverseKL, K5, TopkForwardKL 등 새로운 estimator를 유도
• TopkReverseKL/TopkForwardKL 알고리즘: sampling policy에서 얻은 top-K logit indices q에 대해서는 exact KL(c KLtrun)을 계산하고, q에 속하지 않는 샘플링된 토큰에 대해서는 importance-weighted sampled KL(c KLsampled)을 더해 unbiased 추정을 구성하며, importance ratio s를 [smin, smax]로 clip
• 이를 GRPO에 결합하여 EMA-PG 알고리즘을 구성하고, 수학 reasoning(OlympiadBench 등) 및 검색엔진 기반 agentic Q&A(HotpotQA, 2WikiMultiHopQA, Bamboogle 등 7개 데이터셋)에서 실험 진행

Originality

• 고정 anchor policy 대신 EMA target network 개념을 RL for LLM의 KL anchor에 도입하고, 이를 token-level KL과 결합해 기존 sequence-level EMA 연구(TR-DPO, WARP)와 차별화
• KL estimator 설계에서 "KL 값의 unbiasedness"와 "gradient의 unbiasedness"를 동시에 만족시키는 조건을 명시적으로 분석하고, 이를 만족하는 새로운 estimator(K3++, K4, K5, Top-k 버전)를 체계적으로 유도
• Top-k logit 기반 exact KL과 sampled KL을 결합해 knowledge distillation의 이점을 취하면서도 메모리 효율(O(k))을 확보하는 새로운 방식 제안

Limitation & Further Study

• 실험이 Qwen 계열(1.5B, 3B)의 비교적 작은 모델 규모에 국한되어 있어, 더 큰 규모 모델이나 다른 아키텍처(예: Llama, Mistral)에서도 동일한 이득이 유지되는지 추가 검증이 필요하다.
• EMA anchor의 hyperparameter(η)와 Top-k의 k 값 선택에 대한 민감도 분석이 제한적일 수 있어, 다양한 태스크·도메인에서의 최적 설정 가이드라인이 더 필요하다.
• reasoning과 agentic search 두 도메인에 집중되어 있어, 코딩이나 멀티모달 등 다른 RL 응용 도메인으로의 일반화 가능성은 추가 연구가 필요하다.
• 논문 자체가 ICML 워크숍(3rd AI for Math Workshop) 제출본으로, 정식 학회/저널 심사를 거친 완성도 높은 버전과는 실험 규모나 ablation의 깊이에서 차이가 있을 수 있다.

Evaluation

총평: EMA anchor와 Top-k KL이라는 간단하면서도 이론적으로 뒷받침된 두 가지 수정만으로 GRPO 대비 수학 reasoning과 agentic search RL 모두에서 눈에 띄는 성능 향상을 보인 실용적이고 완성도 높은 연구이다.