Better Estimation of the Kullback--Leibler Divergence Between Language Models

We thank the reviewer for their thoughtful feedback.

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."

We want to clarify that the control variate estimator we discuss (CV with alpha = 1) is the same as the estimator proposed by Schulman (2020), known as the K3 estimator. We will make sure to emphasize this point further in the final version of the paper. In our experiments, we observe that this estimator (CV with alpha=1) can suffer from unbounded variance (figure 1), which limits its reliability compared to our proposed method.

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.

In our RLHF experiments, we compare the variance of the MC and RB estimators at each training step. Figure 2 (right) shows the standard deviation of both estimators’ KL estimates throughout fine-tuning. We run RLHF five times with each estimator and plot the mean KL values along with their standard deviations (shown as shaded regions). The results clearly show that the MC estimator exhibits higher variance across training steps, while the RB estimator provides more stable and reliable KL estimates. Importantly, this improved stability translates to better downstream performance. As demonstrated in Figure 3, models trained using the RB estimator more consistently appear on the Pareto front balancing reward and KL divergence, highlighting the practical benefits of reduced variance in the training dynamics.

Thank you very much for the clarification. We now better understand your concern and have conducted additional experiments to directly assess the variance of the gradient estimators for KL divergence.

Experimental Setup

To ensure a fair and controlled comparison, we directly measured the gradient variance as follows:

• Model: GPT-2 fine-tuned on the IMDB sentiment classification task
• Evaluation Protocol: For each estimator, we ran 5 independent trials, each using 5 batches of 8 prompts (generating 40 responses). All estimators used the same prompts and model initializations to ensure a fair comparison. In each trial, we computed gradients based on the model’s responses to the prompts, with new responses sampled in every run. We then measured the variance of the gradient norms across the 5 runs.

Results

Key Observations

1. Unbiased KL Estimation: Both estimators yield comparable average KL values (7.81–7.96), confirming that they are unbiased.

2. Reduced Gradient Variance: The RB estimator demonstrates a 24.6% reduction in gradient variance compared to the conventional MC estimator (variance ratio = 0.75), providing direct empirical evidence in line with the theoretical motivation.

While we are unfortunately unable to include an additional figure at this stage due to NeurIPS formatting constraints, we will incorporate a detailed visualization and extended discussion in the final camera-ready version.

We hope this addresses your concern, and we would be happy to clarify further if needed.

We thank the reviewer for their thoughtful feedback.

• Regarding the comment that our work applies an existing technique and is methodologically incremental, it is true that our method is based on Rao–Blackwellization, a well-known variance reduction technique. Our contribution lies in demonstrating that this simple, principled technique, requiring only minimal changes and reusing logits already computed during the forward pass, can substantially reduce variance at almost no additional cost. In fact, this simplicity could facilitate the integration of this method in RLHF libraries.

• Regarding other model families and datasets, we refer the reviewer to the response to the reviewer #Q7xD.

• Regarding the assumption of bounded variance, we thank the reviewer for raising this important point. We assume that the true KL divergence is finite. This assumption is generally realistic when p is a model fine-tuned to approximate q, often initialized from a checkpoint of q. However, in cases where q is close to zero while p remains positive, the KL divergence and consequently the variance of our estimators can become unbounded. We agree with the reviewer that this aspect deserves further clarification, and we will address it in the final version of the paper.

We thank the reviewer for their thoughtful feedback. For the comment regarding other model families and datasets, we refer the reviewer to the response to the reviewer #Q7xD.

Questions:

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. What is the implication or relationship between this 'intransitivity' and the proposed RB estimator?

That is an interesting question. As discussed in the cited work, intransitivity in preference data can introduce significant challenges when learning a reliable reward model. In particular, it can lead to a reward function that is noisy or poorly calibrated. When optimizing with a KL-constrained RL objective in this setting, we suspect that the model may compensate for the weak reward signal by relying more heavily on the KL term, since it provides a more consistent optimization direction. However, if the KL estimator used in this process has high variance, this over-reliance can destabilize training. Our proposed RB estimator could help in this case by providing a more stable estimate of the KL divergence.

We thank the reviewer for their thoughtful feedback.

Experiments on another model family and real-world benchmarks

We thank all the reviewers for acknowledging our limited computational resources to conduct large-scale experiments. That said, a recurring concern was the lack of evaluation on a different model family and a real-world benchmark (for example, reviewer #5fbx specifically suggested experiments on the UltraFeedback dataset)

To address this, we conducted an additional experiment to evaluate all the discussed KL divergence estimators on a random subset of prompts from the UltraFeedback dataset. We compute the KL divergence between Zephyr-7B-Beta and its reference model, Mistral-7B-v0.1. Zephyr is fine-tuned from Mistral using DPO on UltraFeedback, and as part of this fine-tuning, it is desirable not to diverge significantly from the base model.

We randomly sampled 512 prompts and generated 100 responses per prompt. For estimation, we used subsets of 1, 5, and 10 samples, reserving the remaining samples to estimate each method’s standard deviation. The table below reports the KL estimate ± standard deviation for each estimator. Unlike our GPT-2 experiments, we had to use half-precision to perform inference and forward passes on a single GPU, which introduces bias in the HT and CV estimators. Consistent with our findings on the IMDB dataset, our proposed RB estimator consistently achieves the lowest standard deviation across all settings, reaffirming its stability and reliability.

Questions

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.

Thank you for this suggestion. The numerical instability arises when computing g(Y), and it's partially captured in Table 1 and Figure 1. However, it's more clearly illustrated in Appendix G (Figure 4). This issue typically occurs with samples from the tail of p (often negative movie reviews) that are still likely under q but suppressed during DPO, making p very small. As a result, the ratio q/p can overflow and bias the estimates. Figure 4 shows that this becomes more likely as the sample size increases (since tail samples are more likely to appear), especially for neutral or negative prompts. We will include a more detailed discussion in the final version of the paper.

Does the method apply to other divergences in the f-divergence family like Jensen–Shannon divergence?

Yes. While our focus is on KL divergence due to its common use, the proposed estimator can also be applied to both forward and backward KL, which together can be used to estimate the Jensen–Shannon divergence.

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.

This estimator is particularly useful when the number of samples available for estimating the divergence is limited, i.e., cases where standard Monte Carlo estimates of KL become highly unreliable. One important example is in RLHF, where GPU memory constraints the batch size, and therefore, the number of samples used for KL estimation is limited. In such settings, a lower-variance estimator can lead to more stable training and better-aligned models.