Efficient Epistemic Uncertainty Estimation for Large Language Models via Knowledge Distillation

TL;DR

The paper tackles the prohibitive cost of estimating epistemic uncertainty (EU) in large language models by leveraging a family of lightweight draft models. It formalizes a bias–variance decomposition where EU is approximated by a variance proxy (the Jensen–Shannon divergence among drafts) plus a bias proxy (the KL divergence between the draft mixture and the target posterior), and introduces Online Stochastic Distillation (OSD) to learn a proxy mean $p_{mix}$ that converges to the target predictive distribution. To maximize posterior diversity, it proposes Data-Diverse Drafts (DDD) and analyzes three distillation variants, demonstrating substantial improvements in EU estimation (e.g., RMSE reductions up to ~37%) and strong hallucination-detection performance with much lower inference cost than heavy perturbation baselines like TokUR. Experiments on GSM8K show that DDD achieves superior uncertainty fidelity across model sizes while maintaining efficient computation, enabling practical uncertainty-aware deployment of LLMs. Overall, the work provides a principled, scalable framework for accurate EU estimation through draft-based inference, combining theory with empirical validation and actionable guidance for implementation.

Abstract

Quantifying uncertainty in Large Language Models (LLMs) is essential for mitigating hallucinations and enabling risk-aware deployment in safety-critical tasks. However, estimating Epistemic Uncertainty(EU) via Deep Ensembles is computationally prohibitive at the scale of modern models. We propose a framework that leverages the small draft models to efficiently estimate token-level EU, bypassing the need for full-scale ensembling. Theoretically grounded in a Bias-Variance Decomposition, our approach approximates EU via Jensen-Shannon divergence among drafts (variance proxy) and KL divergence between the draft mixture and the target (bias proxy). To further ensure accuracy without significant overhead, we introduce Online Stochastic Distillation (OSD) to efficiently approximate target aggregation and the Data-Diverse Drafts (DDD) strategy to enhance draft diversity for better target approximation. Extensive experiments on GSM8K demonstrate that our method reduces the estimation error (RMSE) by up to 37% compared to baselines. Crucially, our approach achieves Hallucination Detection performance competitive with heavy perturbation-based methods like TokUR while incurring negligible inference costs, offering a practical solution for uncertainty-aware LLM deployment.

Figures (3)

• Figure 1: RMSE vs. Spearman correlation for 8B (target model) $\rightarrow$ 3B (draft model). Each point represents one of 5 independent runs. Lower RMSE and higher Spearman indicate better uncertainty estimation. Our proposed DDD achieves the best trade-off among all methods.
• Figure 2: Stability analysis comparing $p_{\text{orig}}$ (baseline) and $p_{\text{mix}}$ (Ours). Box plots show RMSE distribution over 5 runs under $K=3$ and $K=10$ perturbations. Lower RMSE indicates more stable estimation.
• Figure 3: Stability analysis for Llama-3.2-3B-Instruct comparing $p_{\text{orig}}$ (baseline) and $p_{\text{mix}}$ (Ours). Box plots show RMSE distribution over 5 runs under $K=3$ and $K=10$ perturbations. Lower RMSE indicates more stable estimation.

Theorems & Definitions (8)

• Lemma 3.1
• Theorem 3.2: General Upper Bound
• proof
• Theorem 3.4: Bias-Variance Decomposition
• proof
• Lemma 1.1
• proof
• proof