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

arXiv: https://arxiv.org/abs/2504.10637 · HTML: https://mlsanity.com/paper/2504.10637/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Placed in LLM preference optimization methods, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 15 of its 41 references that are in the map.
- First posted on arXiv 14 Apr 2025.

## Placed in

- LLM alignment / RLHF / preference optimisation / [LLM preference optimization methods](https://mlsanity.com/alignment/llm-preference-optimization-methods/)
- Knowledge distillation / [Knowledge distillation for LLMs](https://mlsanity.com/distill/knowledge-distillation-for-llms/)

## Builds on (most cited first, 15 of 15)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [HybridFlow: A Flexible and Efficient RLHF Framework](https://mlsanity.com/paper/2409.19256/) (2024)
- [On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes](https://mlsanity.com/paper/2306.13649/) (2023)
- [Zephyr: Direct Distillation of LM Alignment](https://mlsanity.com/paper/2310.16944/) (2023)
- [UltraFeedback: Boosting Language Models with Scaled AI Feedback](https://mlsanity.com/paper/2310.01377/) (2023)
- [Activation Scaling for Steering and Interpreting Language Models](https://mlsanity.com/paper/2410.04962/) (2024)
- [Context versus Prior Knowledge in Language Models](https://mlsanity.com/paper/2404.04633/) (2024)
- [Using Large Ensembles of Control Variates for Variational Inference](https://mlsanity.com/paper/1810.12482/) (2018)
- [What Languages are Easy to Language-Model? A Perspective from Learning Probabilistic Regular Languages](https://mlsanity.com/paper/2406.04289/) (2024)
- [Can Transformers Learn n-gram Language Models?](https://mlsanity.com/paper/2410.03001/) (2024)
- [The Role of n-gram Smoothing in the Age of Neural Networks](https://mlsanity.com/paper/2403.17240/) (2024)

## Built on by (most cited first, 11 of 11)

- [Reinforcement Learning via Self-Distillation](https://mlsanity.com/paper/2601.20802/) (2026)
- [Self-Distillation Enables Continual Learning](https://mlsanity.com/paper/2601.19897/) (2026)
- [Privileged Information Distillation for Language Models](https://mlsanity.com/paper/2602.04942/) (2026)
- [Composer 2 Technical Report](https://mlsanity.com/paper/2603.24477/) (2026)
- [A Comedy of Estimators: On KL Regularization in RL Training of LLMs](https://mlsanity.com/paper/2512.21852/) (2025)
- [Hybrid Policy Distillation for LLMs](https://mlsanity.com/paper/2604.20244/) (2026)
- [Learning to Explore with Parameter-Space Noise: A Deep Dive into Parameter-Space Noise for Reinforcement Learning with Verifiable Rewards](https://mlsanity.com/paper/2602.02555/) (2026)
- [PrivacyAlign: Contextual Privacy Alignment for LLM Agents](https://mlsanity.com/paper/2606.21710/) (2026)
- [Total Variation Distance Estimation in Autoregressive Models](https://mlsanity.com/paper/2607.19510/) (2026)
- [The Alignment Curse: Modality Alignment Supercharges Audio Attacks via Text Transfer](https://mlsanity.com/paper/2602.02557/) (2026)
- [Function-Space Temporal Self-Distillation for Policy Optimization](https://mlsanity.com/paper/2610.19687/) (2026)

## Abstract (arXiv)

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.