On the Design of KL-Regularized Policy Gradient Algorithms for LLM Reasoning
Yifan Zhang, Yifeng Liu, Rina Hughes, Yang Yuan, Quanquan Gu, Andrew C. Yao
Abstract
Policy gradient algorithms have been successfully applied to enhance the reasoning capabilities of large language models (LLMs). KL regularization is ubiquitous, yet the design surface, choice of KL direction (forward vs. reverse), normalization (normalized vs. unnormalized), and estimator (), is scattered across the literature and often intertwined with off-policy estimation. We ask a focused question: under the off-policy setting, what weighting is required for each KL variant so that the surrogate we optimize yields the exact gradient of the intended KL-regularized objective? We answer this with a compact, unified derivation we call the Regularized Policy Gradient (RPG) view. RPG (i) unifies normalized and unnormalized KL variants and shows that the widely-used penalty is exactly the unnormalized KL; (ii) specifies conditions under which REINFORCE-style losses with stop-gradient are gradient-equivalent to fully differentiable surrogates; (iii) identifies and corrects an off-policy importance-weighting mismatch in GRPO's KL term; and (iv) introduces RPG-Style Clip, a clipped-importance-sampling step within RPG-REINFORCE that enables stable, off-policy policy-gradient training at scale. On mathematical reasoning benchmarks (AIME24, AIME25), RPG-REINFORCE with RPG-Style Clip improves accuracy by up to absolute percentage points over DAPO. We extend our experiments to 8K context length, and RPG-REINFORCE with RPG-Style Clip achieves 52% accuracy on AIME25, surpassing the official Qwen3-4B-Instruct model (47%). Notably, RPG is a stable and scalable RL algorithm for LLM reasoning, realized via (a) a KL-correct objective, (b) clipped importance sampling, and (c) an iterative reference-policy update scheme. Project Page: https://github.com/complex-reasoning/RPG.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 99df550b-94a0-4eff-9a7c-d6c97f48908fCited by top-tier papers9
- Learning to Discover at Test TimeMert Yuksekgonul, Daniel Koceja, Xinhao Li, Federico Bianchi et al.ICML 2026 · 73 citations
- Proximal Diffusion Neural SamplerWei Guo, Jaemoo Choi, Yuchen Zhu, Molei Tao et al.ICLR 2026 · 19 citations
- Rethinking the Design Space of Reinforcement Learning for Diffusion Models: On the Importance of Likelihood Estimation Beyond Loss DesignJaemoo Choi, Yuchen Zhu, Wei Guo, Petr Molodyk et al.ICML 2026 · 15 citations
- Achieving Logarithmic Regret in KL-Regularized Zero-Sum Markov GamesAnupam Nayak, Tong Yang, Osman Yagan, Gauri Joshi et al.ICML 2026 · 9 citations
- Unifying Stable Optimization and Reference Regularization in RLHFLi He, Qiang Qu, He Zhao, Stephen Wan et al.ICLR 2026 · 6 citations
Builds on16
- Training language models to follow instructions with human feedbackLong Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida et al.NeurIPS 2022 · 24,707 citations
- Direct Preference Optimization: Your Language Model is Secretly a Reward ModelRafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning et al.NeurIPS 2023 · 10,924 citations
- DAPO: An Open-Source LLM Reinforcement Learning System at ScaleQiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan et al.NeurIPS 2025 · 2,828 citations
- SimPO: Simple Preference Optimization with a Reference-Free RewardYu Meng, Mengzhou Xia, Danqi ChenNeurIPS 2024 · 1,203 citations
- Efficient Memory Management for Large Language Model Serving with PagedAttentionWoosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng et al.SOSP 2023 · 1,016 citations
Related papers
- GPG: A Simple and Strong Reinforcement Learning Baseline for Model ReasoningXiangxiang Chu, Hailang Huang, Xiao Zhang, Fei Wei et al.ICLR 2026 · 168 citations
- From łog π to π: Taming Divergence in Soft Clipping via Bilateral Decoupled Decay of Probability Gradient WeightXiaoliang Fu, Jiaye Lin, Yangyi Fang, Chaowen Hu et al.ACL 2026 · 5 citations
- Efficient KL Divergence Estimation via Truncated Top-K Integration for Large Language ModelsXinyuan Wang, Luozhijie Jin, Bo Wang, Yuan Li et al.ACL 2026
- MRPO: Magnitude-Regularized Policy Optimization via L1 ConstraintsWei Han, Yuanxing Liu, Mingda Li, Ruiyu Xiao et al.ICML 2026
- Single-stream Policy OptimizationZhongwen Xu, Zihan DingICLR 2026 · 29 citations
