Title: Better Estimation of the Kullback--Leibler Divergence Between Language Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: KL divergence

Monte Carlo estimation

RLHF

post training

alignment
Abstract: Estimating the Kullback--Leibler (KL) divergence between language models has many applications, e.g., reinforcement learning from human feedback (RLHF), interpretability, and knowledge distillation. However, computing the exact KL divergence between two arbitrary language models is intractable. Thus, practitioners often resort to sampling-based estimators. While it is easy to fashion a simple Monte Carlo (MC) estimator that provides an unbiased estimate of the KL divergence between language models, this estimator notoriously suffers from high variance and can even result in a negative estimate of the KL divergence, a non-negative quantity. In this paper, we introduce a Rao--Blackwellized estimator that is unbiased and provably has variance less than or equal to that of the standard Monte Carlo estimator. 
In an empirical study on sentiment-controlled fine-tuning, we show that our estimator provides more stable KL estimates and reduces variance substantially. Additionally, we derive an analogous Rao--Blackwellized estimator of the gradient of the KL divergence, which leads to more stable training and produces models that more frequently appear on the Pareto frontier of reward vs. KL compared to the ones trained with the MC estimator of the gradient.


Peer reviews:
comments: SUMMARY:
This paper propose an estimator to the KL divergence and its gradient during RLHF, based on the decomposition of conditional expectation. The proposed estimation is shown to have lower variance, compared with the classicial Montecarlo exstimation, both in KL divergence itself and its gradient. The authors also use the proposed estimation to conduct RLHF and verify it induce a more stable training process.

STRENGTHS AND WEAKNESSES:
Strengths:

This paper is well written and discuss an important question. 

This paper has a strong theoretical guarantees of the proposed techniques. 

Weakness:

My major concerns are two folds, respectively about the improvement over other varaince reduction techniques and the empirical evidance on the varaince reduction of gradient on KL divergence. 

For the first, though the author compare their estimation with the vanilla MC method, but the improvement over the other variance reduction method is neglected. For example to comparison with  "Schulman. 2020. Approximating KL divergence."

For the second, the empirical results only verify the estimated KL divergence has lower variance, but the verification on the gradient of KL divergence is unexplored, which actually decides the dynamics of RLHF.

QUESTIONS:
See Weankess, I will raise my score if the author can address my concerns in above.

LIMITATIONS:
See Weakness.
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-30

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the problem of estimating the Kullback-Leibler (KL) divergence between language models, which is crucial for applications like reinforcement learning from human feedback (RLHF). The authors identify that while Monte Carlo (MC) estimation is commonly used for this task, it suffers from high variance and can even produce negative estimates of KL divergence, which is theoretically non-negative. To address these issues, the paper introduces a Rao-Blackwellized estimator that applies the classical variance reduction technique of Rao-Blackwellization to KL divergence estimation. The key insight is to decompose the language model estimation into step-wise components and compute exact conditional KL divergences at each step, which are then combined to form the final estimate. The authors prove that their estimator is unbiased and has variance less than or equal to the standard MC estimator, while maintaining the same computational complexity and guaranteeing non-negative estimates.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths：\n\n1. The paper provides rigorous theoretical analysis with formal proofs showing that the Rao-Blackwellized (RB) estimator is unbiased and has variance less than or equal to the Monte Carlo estimator. \n\n2. A key strength is that the proposed method achieves variance reduction without additional computational overhead compared to the standard MC estimator. \n\n3. The experiments systematically compare multiple estimators across different metrics (bias, variance, consistency) and demonstrate practical benefits in RLHF training, including improved stability and better Pareto frontier performance.\n\nWeakness：\n\n1. The core contribution applies an existing statistical technique (Rao-Blackwellization) to a specific domain. While the application is novel and valuable, the methodological innovation is somewhat incremental.\n\n2. The empirical evaluation is restricted to relatively small models (GPT-2) and a single task (sentiment control). While this is acknowledged as a limitation due to computational constraints, it raises questions about scalability to modern large language models and diverse applications.\n\n3. The theoretical results assume finite variance for both estimators, but the paper doesn't deeply analyze when these assumptions might be violated in practice or provide guidance for detecting such cases.\n\nQUESTIONS:\nIs there any assumption about the variance?\n\nLIMITATIONS:\nSee the weakness part.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: SUMMARY:
This paper introduces a Rao-Blackwellized estimator for KL divergence between language models, offering unbiased estimates with lower variance than standard Monte Carlo methods. It improves training stability in RLHF and yields models closer to the reward–KL Pareto frontier. To facilitate fine-tuning practices, this paper also derives a more stable and effective version of the estimator. The arguments are well-supported by both theoretical analysis and experimental results.

STRENGTHS AND WEAKNESSES:
Strength:
1. This paper has proposed an RB estimator for KL divergence estimation between LLMs. Lower variance and unbiasedness are theoretically justified, despite the higher estimation cost.
2. The idea is straightforward, originating from the context of divergence estimation; a code snippet for implementing the RB estimator is provided.
3. Evaluation is mostly limited to sentiment control.

Weakness:
1. Broader alignment tasks and certain well-benchmarked datasets such as UltraFeedback remain to be validated.

QUESTIONS:
Q1, It is out of the scope for this paper, but it would be much helpful if the dependency w.r.t. the reward function is clarified. 

In many real-world datasets, the binary reward models, e.g., the Bradley-Terry(BT) model, have been known to be exposed to 'intransitivity' risk because it relies on scalar variables, where all preferences are transitive by assumption. 
- The literature below studied representative preference datasets in the real world, where the 'transitive' relationship between preference annotations may not always hold. 
- https://arxiv.org/abs/2409.19325 (Duan et al, 2017) 

What is the implication or relationship between this 'intransitivity' and the proposed RB estimator?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 4
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new estimator of KL divergence between language models through Rao-Blackwellization. The proposed estimator is unbiased and has less variance than the baseline Monte Carlo estimator that is currently well used in machine learning applications. The paper introduces the Rao–Blackwellized Monte Carlo method and analyzes the gradient updates. The paper also presents experiments on preference finetuning with GPT-2 on the IMDB dataset. The experiments present bias and variances of different baseline estimators along with averaged rewards using different KL estimators. Results show the proposed RB estimator has low variance and achieves pareto fronts in evaluation.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n- The paper gives detailed preliminaries and derivations of the proposed method such as Control Variate Monte Carlo and Rao–Blackwellization.\n- The paper also analyzes gradient computation and also off-policy gradient for real-world settings.\n- The paper implements and compares with a number of baselines thus making results convincing in experiments.\n\nWeaknesses\n- The experiments only include one preference finetuning task on one small-scale dataset IMDB-sentiment. The experiments are done with GPT-2 model. It would be more helpful to show the advantage of the proposed estimator in more real-world applications. These could be done with more diverse variety of experiments like experiments with other applications such as knowledge distillation which also use KL divergence a lot, or experiments with different tasks in preference finetuning domain, or experiments with different model families.\n\nQUESTIONS:\n- I wonder if there is anyway to quantify the improvements in numerical stability. It could be a plus to show this in the experiment sections.\n- Does the method apply to other divergences in the f-divergence family like Jensen–Shannon divergence?\n- In the authors' opinion, what downstream tasks benefit most from lower variances of KL divergence estimation? It would be helpful to elaborate the potential impact of the proposed method.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-23"}