Better Estimation of the KL Divergence Between Language Models

• URL: http://arxiv.org/abs/2504.10637v1
• Date: Mon, 14 Apr 2025 18:40:02 GMT
• Title: Better Estimation of the KL Divergence Between Language Models
• Authors: Afra Amini, Tim Vieira, Ryan Cotterell,
• Abstract summary: Estimating the Kullback--Leibler (KL) divergence between language models has many applications.<n>We introduce a Rao--Blackwellized estimator that is also unbiased and provably has variance less than or equal to that of the standard Monte Carlo estimator.
• Score: 58.7977683502207
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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 the use of 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 also 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 in practice. 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.

Related papers

• Sequential Monte Carlo for Inclusive KL Minimization in Amortized Variational Inference [3.126959812401426]
  We propose SMC-Wake, a procedure for fitting an amortized variational approximation that uses sequential Monte Carlo samplers to estimate the gradient of the inclusive KL divergence. In experiments with both simulated and real datasets, SMC-Wake fits variational distributions that approximate the posterior more accurately than existing methods.
  arXiv  Detail & Related papers  (2024-03-15T18:13:48Z)
• Nearest Neighbour Score Estimators for Diffusion Generative Models [16.189734871742743]
  We introduce a novel nearest neighbour score function estimator which utilizes multiple samples from the training set to dramatically decrease estimator variance. In diffusion models, we show that our estimator can replace a learned network for probability-flow ODE integration, opening promising new avenues of future research.
  arXiv  Detail & Related papers  (2024-02-12T19:27:30Z)
• Low-variance estimation in the Plackett-Luce model via quasi-Monte Carlo sampling [58.14878401145309]
  We develop a novel approach to producing more sample-efficient estimators of expectations in the PL model. We illustrate our findings both theoretically and empirically using real-world recommendation data from Amazon Music and the Yahoo learning-to-rank challenge.
  arXiv  Detail & Related papers  (2022-05-12T11:15:47Z)
• Causal KL: Evaluating Causal Discovery [0.0]
  Two most commonly used criteria for assessing causal model discovery with artificial data are edit-distance and Kullback-Leibler divergence. We argue that they are both insufficiently discriminating in judging the relative merits of false models. We propose an augmented KL divergence, which takes into account causal relationships which distinguish between observationally equivalent models.
  arXiv  Detail & Related papers  (2021-11-11T02:46:53Z)
• Learning to Estimate Without Bias [57.82628598276623]
  Gauss theorem states that the weighted least squares estimator is a linear minimum variance unbiased estimation (MVUE) in linear models. In this paper, we take a first step towards extending this result to non linear settings via deep learning with bias constraints. A second motivation to BCE is in applications where multiple estimates of the same unknown are averaged for improved performance.
  arXiv  Detail & Related papers  (2021-10-24T10:23:51Z)
• Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient Estimator [93.05919133288161]
  We show that the variance of the straight-through variant of the popular Gumbel-Softmax estimator can be reduced through Rao-Blackwellization. This provably reduces the mean squared error. We empirically demonstrate that this leads to variance reduction, faster convergence, and generally improved performance in two unsupervised latent variable models.
  arXiv  Detail & Related papers  (2020-10-09T22:54:38Z)
• Nonparametric Estimation of the Fisher Information and Its Applications [82.00720226775964]
  This paper considers the problem of estimation of the Fisher information for location from a random sample of size $n$. An estimator proposed by Bhattacharya is revisited and improved convergence rates are derived. A new estimator, termed a clipped estimator, is proposed.
  arXiv  Detail & Related papers  (2020-05-07T17:21:56Z)
• SUMO: Unbiased Estimation of Log Marginal Probability for Latent Variable Models [80.22609163316459]
  We introduce an unbiased estimator of the log marginal likelihood and its gradients for latent variable models based on randomized truncation of infinite series. We show that models trained using our estimator give better test-set likelihoods than a standard importance-sampling based approach for the same average computational cost.
  arXiv  Detail & Related papers  (2020-04-01T11:49:30Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.