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.

本研究聚焦于语言模型之间KL散度估计的高方差问题，尤其是在传统的采样基础估计方法中的不足。我们提出了一种Rao-Blackwell化估计器，它在保持无偏性的同时，具有低于标准蒙特卡洛估计器的方差。实证研究表明，该估计器在情感控制的微调中显著提高了KL估计的稳定性，同时我们还推导了KL散度梯度的类似估计器，进一步改善了训练稳定性。

更好的语言模型间KL散度估计