Efficient Epistemic Uncertainty Estimation for Large Language Models via Knowledge Distillation

TL;DR

This paper introduces a computationally efficient framework for estimating epistemic uncertainty in large language models using knowledge distillation, significantly reducing costs while maintaining accuracy for safety-critical applications.

Contribution

It presents a novel approach that leverages small draft models and theoretical bias-variance decomposition to estimate uncertainty without full ensembling, including new strategies for draft diversity and efficiency.

Findings

Reduces estimation error (RMSE) by up to 37% on GSM8K.

Achieves competitive hallucination detection with minimal inference overhead.

Provides a practical, scalable uncertainty estimation method for large language models.

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.