REVIEW 3 major objections 5 minor 2 cited by
On a few pitfalls in KL divergence gradient estimation for RL
T0 review · 3 major / 5 minor · reviewed 2026-08-07 · deepseek-v4-flash
Pith's one-line read Differentiating Monte Carlo KL estimates as RL losses yields biased or zero gradients, not the KL gradient; the paper proves this and gives correct estimators.
desk verdict The core gradient identities are correct and the bug diagnosis is real, but the paper overreaches in applying them to PPO-style RLHF with stale rollouts, and the LLM evidence is too thin to carry the practical claims. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is the decomposition of a KL gradient into a path-wise derivative and a score-function derivative. For an unbiased estimator $\hat{d}_{\mathrm{KL}}$ of $\mathrm{KL}(\pi, \pi_{\mathrm{ref}})$, the full gradient is $\nabla \mathrm{KL} = \mathbb{E}[\nabla \hat{d}_{\mathrm{KL}}] + \mathbb{E}[\hat{d}_{\mathrm{KL}} \nabla \log \pi(y)]$; implementations that differentiate the estimator as a loss compute only the first term. The named estimators are the vanilla estimate $\hat{d}_{\mathrm{KL}} = \log(\pi(y)/\pi_{\mathrm{ref}}(y))$, the variance-reduced estimate $\hat{d}_{\mathrm{KL}} = \log(\pi(y)/\pi_{\mathrm{ref}}(y)) + \pi_{\mathrm{ref}}(y)/\pi(y) - 1$, and the squared estimate $\hat{d}_{\mathrm{KL}} = \tfrac{1}{2}(\log(\pi(y)/\pi_{\mathrm{ref}}(y)))^2$, which happens to produce the unbiased gradient when differentiated. For sequences, the machinery is a REINFORCE-style expansion showing that the sequence-level KL gradient equals $\mathbb{E}\big[\sum_t \nabla \log \pi(y_t \mid y_{<t}) \sum_{s \ge t} \log \rho_s\big]$ with $\rho_s = \pi(y_s \mid y_{<s}) / \pi_{\mathrm{ref}}(y_s \mid y_{<s})$, which yields the cumulative estimator that credits past tokens for their effect on future KL terms.
What would settle it
On a small tabular policy (for example a 100-arm bandit) with known $\pi_{\mathrm{ref}}$, draw many samples $y \sim \pi$ and compute the empirical mean of the gradient of $\log(\pi(y)/\pi_{\mathrm{ref}}(y))$ with gradients stopped on the sampling distribution; if this mean is not close to zero for $\pi \neq \pi_{\mathrm{ref}}$, equation (2) is false, and if the empirical mean of the variance-reduced loss gradient differs from $\nabla \mathrm{KL}(\pi_{\mathrm{ref}}, \pi)$, the reverse-KL claim is false.
Extended reading notes
Core claim
Minimizing KL divergence by differentiating an unbiased Monte Carlo estimate of the KL as a loss is generally incorrect: unbiasedness of the estimate does not carry over to unbiasedness of its gradient. Because $\mathrm{KL}(\pi, \pi_{\mathrm{ref}})$ depends on $\pi$ both through the sampling distribution and the integrand, the true gradient has a score-function component $\mathbb{E}[\log(\pi(y)/\pi_{\mathrm{ref}}(y)) \nabla \log \pi(y)]$ as well as the path-wise component $\mathbb{E}[\nabla \log(\pi(y)/\pi_{\mathrm{ref}}(y))]$. Differentiating the vanilla estimate keeps only the path-wise part, whose expectation is zero; differentiating the variance-reduced estimate $\log(\pi(y)/\pi_{\mathrm{ref}}(y)) + \pi_{\mathrm{ref}}(y)/\pi(y) - 1$ gives $\mathbb{E}[\nabla \hat{d}_{\mathrm{KL}}] = \nabla \mathrm{KL}(\pi_{\mathrm{ref}}, \pi)$, reversing the arguments of the divergence, with the regularization effect coming entirely from the control-variate term. For autoregressive sequences, summing token-level KL losses ignores the influence of earlier tokens on later token-level KLs, so the aggregate is at best a partial gradient of the sequence-level KL. The paper supplies correct estimators—the vanilla score-function estimate $\log(\pi(y)/\pi_{\mathrm{ref}}(y)) \nabla \log \pi(y)$ and its leave-one-out and cumulative sequence-level variants—and shows in tabular and LLM experiments that they minimize the target KL, whereas the incorrect implementations plateau or converge to a different policy.
Load-bearing premise
The load-bearing premise is that samples are drawn on-policy from the current policy $\pi$ and the reference policy $\pi_{\mathrm{ref}}$ is held fixed, so the score-function identity $\mathbb{E}[\nabla \log \pi(y)] = 0$ and the importance-sampling step in the appendix hold; if data come from a stale behavior policy or $\pi_{\mathrm{ref}}$ is also being trained, the claimed biases and the proposed correct estimators do not apply.
Editorial extensions
If this is right
- RLHF systems that use the vanilla KL estimate as a differentiable loss are not enforcing KL at all: in expectation the gradient is zero, so the KL term acts as a zero-mean noise term on the policy update.
- RLHF systems that use the variance-reduced KL estimate are incidentally minimizing the reverse KL $\mathrm{KL}(\pi_{\mathrm{ref}}, \pi)$; when combined with reward maximization this converges to a different optimal policy than the intended KL regularizer.
- Token-level KL losses in LLM training are partial gradients of the sequence-level KL, because they omit the dependence of future token-level KLs on earlier tokens.
- The correct score-function estimator and its sequence-level cumulative variant minimize the target KL and improve KL efficiency, with the largest gains in the high-KL regime such as on-policy distillation.
- In the low-KL regime, where $\mathrm{KL}(\pi, \pi_{\mathrm{ref}}) \approx \mathrm{KL}(\pi_{\mathrm{ref}}, \pi)$, the choice of gradient estimator makes little practical difference to the reward-KL trade-off.
Reading between the lines
- A practical audit of existing RLHF code: if the KL term is implemented as a loss that differentiates through $\log \pi(y)$ with gradients flowing into the sampling distribution, the paper's account predicts that term is either inert or reverse-KL; wrapping the log-ratio in a stop-gradient and using a score-function term would restore the intended gradient.
- The incidental reverse-KL behavior may explain why some RLHF systems work despite the bug: reverse KL also anchors the policy near the reference, and in low-KL regimes the two divergences are close; the failure only becomes visible under large policy drift.
- The same reasoning suggests a design rule for any parameterized divergence regularizer: write the surrogate loss with a stop-gradient on the density ratio, or explicitly separate the score-function estimator, so that unbiasedness of the gradient holds by construction.
- A testable extension: in a long RLHF run with weak KL regularization, compare the final policy against the optimum of $\mathrm{KL}(\pi, \pi_{\mathrm{ref}})$-regularized and $\mathrm{KL}(\pi_{\mathrm{ref}}, \pi)$-regularized objectives; the paper's theory predicts the variance-reduced estimator will track the reverse-KL optimum.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. This short note studies gradient estimation for KL-divergence regularization in RL and LLM fine-tuning. The central claim is that differentiating an unbiased Monte Carlo estimate of KL as a loss does not yield an unbiased gradient of the KL objective. The authors show that the vanilla estimate has zero expected gradient under on-policy sampling, that the variance-reduced estimate has expectation equal to the gradient of the reverse KL, and that the squared estimate, when differentiated, yields the correct forward-KL gradient. They further argue that token-level losses give at most a partial gradient of the sequence-level KL and propose sequence-level estimators, including leave-one-out and cumulative variants. Tabular experiments verify the bias and optimization effects; LLM experiments on reward maximization and distillation illustrate the practical consequences.
Significance. The paper identifies a real and frequently misunderstood issue in RLHF implementations: the difference between estimating a scalar objective and estimating its gradient. The derivations in Section 3 and Appendices A, B, and D are clean and correct under the stated on-policy assumption, and the tabular experiments substantiate the bias claims. The proposed sequence-level estimators and the cumulative variant are useful practical contributions. However, the practical scope is narrower than the abstract and conclusion suggest: every central identity requires samples from the current policy with a frozen reference policy, which is not the setting of PPO-style rollouts. The LLM experiments are single-run and lack statistical support. If the scope is clarified and the empirical claims are appropriately qualified, the paper would be a valuable cautionary note for the RLHF community.
major comments (3)
- [Section 3, Eqs. (2)-(3), Appendix A, Section 6] The identities in Section 3, Eqs. (2)-(3), and Appendix A assume y~π and a frozen π_ref. Under a stale behavior policy π_old, E_{y~π_old}[∇log π(y)] is generally nonzero, so the 'inert vanilla loss' and 'incidentally reverse-KL' conclusions do not transfer to PPO-style implementations that optimize over multiple epochs on a rollout buffer. Section 6 acknowledges off-policy learning but only proposes the squared regularizer; it does not analyze the bias of the criticized estimators under π_old. Since the paper targets open-source RLHF implementations that often use stale rollouts, this is a load-bearing gap. The authors should either restrict the practical claims explicitly to on-policy sampling or provide an off-policy analysis of the estimators.
- [Section 5, Figures 7-8] The high-KL distillation claim in Section 5.2 (that the variance-reduced incorrect estimate makes little progress on the target KL) is supported only by a single trajectory with no error bars or multiple seeds, and the same is true for the reward-maximization results in Figure 7. Without repetitions, the empirical support for the conclusion that sequence-level estimators 'impact performance significantly in high-KL regime' is not statistically assessable. Please report multiple seeds with confidence intervals, or clearly label these curves as illustrative single runs and temper the corresponding conclusions.
- [Section 4.1 and Appendix D] The statement that token-level variance-reduced losses 'do not produce even the partial gradient of any existing sequence-level divergence' is stated without a derivation in the paper; Appendix D derives the decomposition of the full sequence gradient but does not compute the expectation of the variance-reduced token-level loss itself. The claim is plausible under the on-policy assumption, but the paper currently delegates the proof to Tang et al. (2025). Please provide the explicit computation or state precisely which assumptions are needed and which part of the claim is proven elsewhere.
minor comments (5)
- [References] The in-text citation '(Paszke et al., 1912)' should be '(Paszke et al., 2019)'; the reference entry also mixes the year 1912 and 2019.
- [Section 4.2 vs. Appendix D] The notation for ρ_t is inconsistent: Section 4.2 defines ρ_t as the ratio π(yt|y1:t−1)/πref(yt|y1:t−1), while Appendix D defines ρ_t as the log ratio; please unify the notation.
- [Section 3, Eq. (3)] There is a bracket typo in Eq. (3): 'E h [∇KLvar-reduced i' should be 'E[∇d̂_KL^var-reduced]' or similar.
- [Appendix C] The sentence 'Throughput we adopt a learning rate of η = 1' contains a typo: 'Throughput' should be 'Throughout'.
- [Figures 7-8 captions] The captions of Figures 7 and 8 should state the number of runs and explicitly note that the curves are single trajectories, so that readers do not infer error bars or statistical significance that are not present.
Circularity Check
No significant circularity: the paper's central gradient identities are derived from standard score-function and importance-sampling calculations, not from its own conclusions.
full rationale
The paper's load-bearing claims are Eqs. (2) and (3): differentiating the vanilla KL estimate has zero expectation and differentiating the variance-reduced estimate yields the reverse-KL gradient. Both are derived in the text and Appendix A directly from the definitions of the estimators and the classical identity E_{y~pi}[nabla log pi(y)] = 0. No fitted parameter is renamed as a prediction, and no result is imported from the authors' prior work as the basis for the main derivation. The unbiased gradient formula in Appendix B and the squared-estimate identity follow by routine product-rule and score-function manipulations, and the sequence-level results in Section 4 and Appendix D are derived from REINFORCE-type score-function arguments with the stated on-policy sampling assumption. The self-citations that do occur, such as Tang et al. (2025) for extended sequence-level discussion and the off-policy 'unique regularization' remark in Section 6, support secondary or motivational claims and are not load-bearing for the central zero-gradient or reverse-KL identities. The paper explicitly assumes samples are drawn from pi, which is a stated hypothesis of the estimator analysis rather than a circular import; whether that assumption holds in practical stale-rollout RLHF pipelines is a correctness-scope concern, not a circularity concern. Overall, the derivation chain is self-contained and the paper does not reduce by construction to its inputs.
Assumptions & free parameters
assumptions (5)
- standard math Score function identity E[∇ log π(y)] = 0 for y ~ π
- standard math Importance sampling identity E_{y~π}[π_ref(y)/π(y) g(y)] = E_{y~π_ref}[g(y)]
- domain assumption π_ref is fixed and does not receive gradient
- domain assumption Samples are drawn on-policy from π
- standard math Auto-regressive factorization of the sequence distribution
Cite this review
Pith. "Pith review of On a few pitfalls in KL divergence gradient estimation for RL." pith.science (2026). https://pith.science/paper/GRKZZCHM
@misc{pith2026250609477,
author = {Pith},
title = {Pith review of: On a few pitfalls in KL divergence gradient estimation for RL},
year = {2026},
howpublished = {\url{https://pith.science/paper/GRKZZCHM}},
note = {Machine review of arXiv:2506.09477}
}
read the original abstract
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.
Figures
Figures from the paper (3 more)
Forward citations
Cited by 2 Pith papers
-
Outcome-based Exploration for LLM Reasoning
Outcome-based exploration bonuses (UCB-Con and Batch) improve pass@1 and pass@32 for LLM math reasoning while slowing diversity collapse, supported by a bandit model with a strong generalization assumption.
-
ARMOR: Stabilizing On-Policy LLM RL with Off-Policy Anchor Samples
Injecting correct reference-policy samples and optimizing a mixed importance-sampling ratio prevents validation collapse and raises asymptotic math-reasoning scores beyond reverse-KL baselines.
Reference graph
Works this paper leans on
-
[8]
Aaron Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, et al. Openai o1 system card. arXiv preprint arXiv:2412.16720,
-
[9]
Understand- ing r1-zero-like training: A critical perspective
Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin. Understand- ing r1-zero-like training: A critical perspective. arXiv preprint arXiv:2503.20783,
-
[11]
Gradient estimation using stochastic com- putation graphs
John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel. Gradient estimation using stochastic com- putation graphs. Advances in neural information process- ing systems, 28, 2015a. John Schulman, Sergey Levine, Pieter Abbeel, Michael Jor- dan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pa...
-
[13]
Rl-finetuning llms from on-and off- policy data with a single algorithm
Yunhao Tang, Taco Cohen, David W Zhang, Michal Valko, and R´emi Munos. Rl-finetuning llms from on-and off- policy data with a single algorithm. arXiv preprint arXiv:2503.19612,
-
[15]
d1: Scaling reasoning in diffusion large lan- guage models via reinforcement learning
Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover. d1: Scaling reasoning in diffusion large lan- guage models via reinforcement learning. arXiv preprint arXiv:2504.12216,
-
[17]
9 A. Details on theoretical results of variance reduced gradient estimate We provide some more details on the theoretical results of the variance reduced gradient estimate. It follows that Ey∼π h [∇KLvar-reduced i =(a) Ey∼π − πref(y) π(y)2 ∇π(y) =(b) Ey∼π − πref(y) π(y) ∇ log π(y) =(c) −Ey∼πref [−∇ log π(y)] =(d) ∇KL(πref, π). Here, (a) is due to the fact...
work page 2004
-
[18]
of size 8B and 70B. The reward maximization and on-policy distillation experiments are both carried out on the 7500 MATH training prompts (Hendrycks et al., 2021). We use n = 4samples per gradient estimate. 11
work page 2021
-
[1992]
On the design of kl- regularized policy gradient algorithms for llm reasoning
Yifan Zhang, Yifeng Liu, Huizhuo Yuan, Yang Yuan, Quan- quan Gu, and Andrew C Yao. On the design of kl- regularized policy gradient algorithms for llm reasoning. arXiv preprint arXiv:2505.17508,
Show all 18 references
-
[1999]
Generalized preference optimization: A unified approach to offline alignment
Yunhao Tang, Zhaohan Daniel Guo, Zeyu Zheng, Daniele Calandriello, R´emi Munos, Mark Rowland, Pierre Har- vey Richemond, Michal Valko, Bernardo ´Avila Pires, and Bilal Piot. Generalized preference optimization: A unified approach to offline alignment. arXiv preprint arXiv:2402.05749,
-
[2004]
The llama 3 herd of models
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhi- nav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783,
-
[2008]
Fine-tuning language models from human preferences
Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv:1909.08593,
1909 arXiv
-
[2017]
Soft policy opti- mization: Online off-policy rl for sequence models
Taco Cohen, David W Zhang, Kunhao Zheng, Yunhao Tang, Remi Munos, and Gabriel Synnaeve. Soft policy opti- mization: Online off-policy rl for sequence models. arXiv preprint arXiv:2503.05453,
-
[2018]
Measuring mathematical problem solving with the math dataset
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Ja- cob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874,
-
[2019]
arXiv preprint arXiv:1912.01703, 10,
1912 arXiv
-
[2020]
Deepseek-r1: Incentivizing reason- ing capability in llms via reinforcement learning
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-r1: Incentivizing reason- ing capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948,
-
[2021]
Correcting the mythos of kl-regularization: Direct align- ment without overoptimization via chi-squared preference optimization
Audrey Huang, Wenhao Zhan, Tengyang Xie, Jason D Lee, Wen Sun, Akshay Krishnamurthy, and Dylan J Foster. Correcting the mythos of kl-regularization: Direct align- ment without overoptimization via chi-squared preference optimization. arXiv preprint arXiv:2407.13399,
-
[2024]
Mohammad Gheshlaghi Azar, Zhaohan Daniel Guo, Bilal Piot, Remi Munos, Mark Rowland, Michal Valko, and Daniele Calandriello
URL https://arxiv.org/abs/2306.13649. Mohammad Gheshlaghi Azar, Zhaohan Daniel Guo, Bilal Piot, Remi Munos, Mark Rowland, Michal Valko, and Daniele Calandriello. A general theoretical paradigm to understand learning from human preferences. In Interna- tional Conference on Arti...
-
[2025]
Contrastive policy gradient: Aligning llms on sequence-level scores in a supervised-friendly fashion
Yannis Flet-Berliac, Nathan Grinsztajn, Florian Strub, Bill Wu, Eugene Choi, Chris Cremer, Arash Ahmadian, Yash Chandak, Mohammad Gheshlaghi Azar, Olivier Pietquin, et al. Contrastive policy gradient: Aligning llms on sequence-level scores in a supervised-friendly fashion. arX...
Reviewed August 7, 2026 · model on record in the stance chip above.
Discussion (0). Sign in to comment.