# On a few pitfalls in KL divergence gradient estimation for RL

arXiv: https://arxiv.org/abs/2506.09477 · HTML: https://mlsanity.com/paper/2506.09477/

## Key facts

- Cited by 29 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 37 references that are in the map.
- First posted on arXiv 11 Jun 2025.

## Builds on (most cited first, 26 of 26)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning](https://mlsanity.com/paper/2501.12948/) (2025)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Fine-Tuning Language Models from Human Preferences](https://mlsanity.com/paper/1909.08593/) (2019)
- [Understanding R1-Zero-Like Training: A Critical Perspective](https://mlsanity.com/paper/2503.20783/) (2025)
- [Scaling Laws for Reward Model Overoptimization](https://mlsanity.com/paper/2210.10760/) (2022)
- [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://mlsanity.com/paper/2306.13649/) (2023)
- [A General Theoretical Paradigm to Understand Learning from Human Preferences](https://mlsanity.com/paper/2310.12036/) (2023)
- [Black Box Variational Inference](https://mlsanity.com/paper/1401.0118/) (2013)
- [d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning](https://mlsanity.com/paper/2504.12216/) (2025)
- [Equivalence Between Policy Gradients and Soft Q-Learning](https://mlsanity.com/paper/1704.06440/) (2017)
- [Gradient Estimation Using Stochastic Computation Graphs](https://mlsanity.com/paper/1506.05254/) (2015)
- [Variational inference for Monte Carlo objectives](https://mlsanity.com/paper/1602.06725/) (2016)
- [Correcting the Mythos of KL-Regularization: Direct Alignment without Overoptimization via Chi-Squared Preference Optimization](https://mlsanity.com/paper/2407.13399/) (2024)
- [Monte-Carlo Tree Search as Regularized Policy Optimization](https://mlsanity.com/paper/2007.12509/) (2020)
- [On the Design of KL-Regularized Policy Gradient Algorithms for LLM Reasoning](https://mlsanity.com/paper/2505.17508/) (2025)
- [Contrastive Policy Gradient: Aligning LLMs on sequence-level scores in a supervised-friendly fashion](https://mlsanity.com/paper/2406.19185/) (2024)
- [Soft Policy Optimization: Online Off-Policy RL for Sequence Models](https://mlsanity.com/paper/2503.05453/) (2025)
- [RL-finetuning LLMs from on- and off-policy data with a single algorithm](https://mlsanity.com/paper/2503.19612/) (2025)

## Built on by (most cited first, 29 of 29)

- [Self-Distillation Enables Continual Learning](https://mlsanity.com/paper/2601.19897/) (2026)
- [Skill-SD: Skill-Conditioned Self-Distillation for Multi-turn LLM Agents](https://mlsanity.com/paper/2604.10674/) (2026)
- [Outcome-based Exploration for LLM Reasoning](https://mlsanity.com/paper/2509.06941/) (2025)
- [Principled RL for Diffusion LLMs Emerges from a Sequence-Level Perspective](https://mlsanity.com/paper/2512.03759/) (2025)
- [EXAONE 4.0: Unified Large Language Models Integrating Non-reasoning and Reasoning Modes](https://mlsanity.com/paper/2507.11407/) (2025)
- [KL-Regularized Reinforcement Learning is Designed to Mode Collapse](https://mlsanity.com/paper/2510.20817/) (2025)
- [KL for a KL: On-Policy Distillation with Control Variate Baseline](https://mlsanity.com/paper/2605.07865/) (2026)
- [GenEvolve: Self-Evolving Image Generation Agents via Tool-Orchestrated Visual Experience Distillation](https://mlsanity.com/paper/2605.21605/) (2026)
- [A Comedy of Estimators: On KL Regularization in RL Training of LLMs](https://mlsanity.com/paper/2512.21852/) (2025)
- [OPD+: Rethinking the Advantage Design for On-Policy Distillation](https://mlsanity.com/paper/2606.01039/) (2026)
- [Neglected Free Lunch from Post-training: Progress Advantage for LLM Agents](https://mlsanity.com/paper/2606.26080/) (2026)
- [Stabilizing Policy Gradients for Sample-Efficient Reinforcement Learning in LLM Reasoning](https://mlsanity.com/paper/2510.00819/) (2025)
- [EMA Policy Gradient: Taming Reinforcement Learning for LLMs with EMA Anchor and Top-k KL](https://mlsanity.com/paper/2602.04417/) (2026)
- [f-Trajectory Balance: A Loss Family for Tuning GFlowNets, Generative Models, and LLMs with Off- and On-Policy Data](https://mlsanity.com/paper/2605.15417/) (2026)
- [Rethinking Large Language Model Distillation: A Constrained Markov Decision Process Perspective](https://mlsanity.com/paper/2509.22921/) (2025)
- [Fast Rates for Offline Contextual Bandits with Forward-KL Regularization under Single-Policy Concentrability](https://mlsanity.com/paper/2605.09214/) (2026)
- [Scaling Reward Modeling without Human Supervision](https://mlsanity.com/paper/2603.02225/) (2026)
- [Sample Efficient Generative Molecular Optimization with Joint Self-Improvement](https://mlsanity.com/paper/2602.10984/) (2026)
- [MHPO: Modulated Hazard-aware Policy Optimization for Stable Reinforcement Learning](https://mlsanity.com/paper/2603.16929/) (2026)
- [Beyond Distribution Sharpening: The Importance of Task Rewards](https://mlsanity.com/paper/2604.16259/) (2026)
- [Local Guidance, Global Impact: Gaussian-Reshaped Trust Region Unlocks Behavior Transitions](https://mlsanity.com/paper/2606.03382/) (2026)
- [Interpreting Language Model Hidden States at Scale](https://mlsanity.com/paper/2608.10260/) (2026)
- [Suan: Rectifying Direct Preference Safety Alignment in Large Language Models](https://mlsanity.com/paper/2609.08634/) (2026)
- [Scaling Properties of Same-Family On-Policy Distillation](https://mlsanity.com/paper/2609.32722/) (2026)
- [On-Policy or Off-Policy Learning? A Systematic Study of Distillation Dynamics](https://mlsanity.com/paper/2609.35259/) (2026)
- [The Weakest Link: Distilling LLM Reasoning with Worst-Case Constrained Reinforcement Learning](https://mlsanity.com/paper/2610.00332/) (2026)
- [Ontology Concept Overlap as a Training Signal: Knowledge-Grounded Reinforcement Learning for Clinical Question Answering](https://mlsanity.com/paper/2610.06360/) (2026)
- [Reward-Driven Learning under Prompt-Level Differential Privacy](https://mlsanity.com/paper/2610.07212/) (2026)
- [When KL Regularization Misfires in Group Policy Optimization](https://mlsanity.com/paper/2610.12161/) (2026)

## Abstract (arXiv)

We point out a few pitfalls in implementing gradient estimation for KL divergence in RL training for LLM, as seen in a number of open source projects and papers. The first major pitfall is to differentiate through the KL estimate as loss functions to minimize KL divergence. We show that such implementations are generally incorrect and do not produce the desired KL gradient. Secondly, we show that some implementations do not account for the sequential nature of the estimation problem and produce a partial gradient at best. We demonstrate the impact of such issues with illustrative tabular and LLM experiments, and show the correct way to implement the KL gradient.