Research map / Probabilistic and scientific modeling / Variational and generative sampling
Better Estimation of the Kullback--Leibler Divergence Between Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a Rao-Blackwellized Monte Carlo estimator for KL divergence between language models (and its gradient), which is unbiased, non-negative, and has no higher variance than the standard MC estimator, and shows more stable RLHF training.
Problem
Sampling-based estimates of KL divergence between language models have high variance and can be negative, which is problematic for RLHF and other applications that need stable KL estimates and gradients.
Approach
Applies Rao-Blackwellization to the step-wise decomposition of the MC KL estimator, computing the exact per-position KL between next-token distributions conditioned on sampled prefixes; this yields an unbiased, non-negative estimator with provably no greater variance than MC, and a matching Rao-Blackwellized gradient estimator.
Key results
- On IMDB prompts with M=1, the RB estimator reduced standard deviation to 0.11 versus 0.16 for MC (estimated KL 6.76 for unbiased estimators).
- The control variate estimator with alpha=1 showed bias (6.28) and large variance (±2.54 at M=1).
- 78% of points on the overall reward-vs-KL Pareto front came from RB-fine-tuned models (95% for KL below 5).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the MC KL estimator with Rao–Blackwellization | KL estimation on IMDB prompts · IMDB prompts, M=1 | standard deviation of KL estimate: 0.16 -> 0.11 | Helped |
| Replace the MC KL estimator with Rao–Blackwellization | KL estimation on IMDB prompts · IMDB prompts, M=5 | standard deviation of KL estimate: 0.07 -> 0.05 | Helped |
| Replace the MC KL estimator with Rao–Blackwellization | KL estimation on IMDB prompts · IMDB prompts, M=10 | standard deviation of KL estimate: 0.05 -> 0.03 | Helped |
| Use the Horvitz–Thompson estimator instead of MC | KL estimation on IMDB prompts · IMDB prompts | standard deviation of KL estimate: 0.16 -> 0.16; 0.07 -> 0.07; 0.05 -> 0.05 | No clear effect |
| Use the control variate with optimal alpha instead of MC | KL estimation on IMDB prompts · IMDB prompts | standard deviation of KL estimate: 0.16 -> 0.16; 0.07 -> 0.07; 0.05 -> 0.05 | No clear effect |
| Set the control-variate alpha to 1 | Control-variate KL estimator · IMDB prompts | standard deviation of KL estimate: 0.16 -> 2.54; 0.07 -> 1.13; 0.05 -> 0.79 | Hurt |
| Set the control-variate alpha to 1 | Control-variate KL estimator · IMDB prompts | estimated KL: 6.76 -> 6.28 | Hurt |
| Replace the MC KL-gradient estimator with the Rao–Blackwellized gradient estimator | RLHF fine-tuning · Sentiment-controlled RLHF, 40 prompts | variance of gradient norm: 59.90 -> 45.44; 24.6% reduction | Helped |
| Replace the MC KL-gradient estimator with the Rao–Blackwellized gradient estimator | RLHF fine-tuning · RLHF reward–KL Pareto frontier | share of overall Pareto-front points: 78% | Helped |
| Replace the MC KL-gradient estimator with the Rao–Blackwellized gradient estimator | RLHF fine-tuning · RLHF reward–KL Pareto frontier, KL values below 5 | share of overall Pareto-front points: 95% | Helped |
| Replace the MC KL-gradient estimator with the Rao–Blackwellized gradient estimator | RLHF fine-tuning · Sentiment-controlled generation | training stability | Helped |
Surprising findings
- The control variate estimator with alpha=1 (used in prior RLHF libraries) fails to reduce variance and can exhibit bias and unbounded variance for some prompts.
- Horvitz-Thompson estimation gave no significant variance reduction over MC in the experiments.
Limitations
- Experiments limited to GPT-2 on a sentiment-controlled task due to compute; larger models and other RL alignment methods not tested.
- Significance tests required training many models (36), making variance evaluation costly.
- Exact KL between arbitrary LMs is intractable; the estimator still relies on sampling from the policy.
Experimental setup
Datasets: IMDB
Benchmarks: IMDB sentiment-controlled generation
Metrics: estimator standard deviation bias average reward KL divergence Pareto frontier fraction gradient norm variance
Baselines: Monte Carlo KL estimator Horvitz-Thompson estimator control variate estimator (alpha=1) control variate estimator (optimal alpha) MC gradient estimator
Builds on: Rao-Blackwell theorem control variate KL estimator (Schulman k3) Horvitz-Thompson estimator DPO REINFORCE
Closest papers by meaning
- Correcting the Mythos of KL-Regularization: Direct Alignment without Overoptimization via Chi-Squared Preference Optimization (ICLR 2025 · similarity 0.88)
- Learning Parametric Distributions from Samples and Preferences (ICML 2025 · similarity 0.87)
- Theoretical Limits of Language Model Alignment (NeurIPS 2026 · similarity 0.87)
- BLoB: Bayesian Low-Rank Adaptation by Backpropagation for Large Language Models (NeurIPS 2024 · similarity 0.86)
- Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo (ICML 2024 · similarity 0.86)
- Statistical and Geometrical properties of the Kernel Kullback-Leibler divergence (NeurIPS 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.