Efficient Uncertainty Estimation via Distillation of Bayesian Large Language Models

TL;DR

This work tackles the inefficiency of test-time uncertainty estimation in Bayesian LLMs by introducing Efficient Uncertainty Estimation via Distillation (EUD), which transfers a Bayesian teacher's predictive distribution to a deterministic student. By minimizing a KL-based distillation loss between the teacher's predictive distribution $P(\boldsymbol{y}|\boldsymbol{x}, \boldsymbol{\phi})$ and the student's $P(\boldsymbol{y}|\boldsymbol{x}, \boldsymbol{\theta})$, and employing a simple loss schedule $\alpha_t$, EUD enables a single forward pass at inference while preserving uncertainty calibration. The method is augmented by training-data expansion and uses a LoRA-based student to maintain parameter efficiency. Experiments on Meta-Llama-3.1-8B across six in-distribution tasks and cross-domain OOD scenarios show that EUD achieves competitive accuracy and superior uncertainty calibration compared with sampling-free baselines, while delivering substantial test-time speedups over sampling-based Bayesian methods. These results highlight the practicality of distilling Bayesian LLMs into deterministic models for reliable, scalable uncertainty estimation in real-world applications.

Abstract

Recent advances in uncertainty estimation for Large Language Models (LLMs) during downstream adaptation have addressed key challenges of reliability and simplicity. However, existing Bayesian methods typically require multiple sampling iterations during inference, creating significant efficiency issues that limit practical deployment. In this paper, we investigate the possibility of eliminating the need for test-time sampling for LLM uncertainty estimation. Specifically, when given an off-the-shelf Bayesian LLM, we distill its aligned confidence into a non-Bayesian student LLM by minimizing the divergence between their predictive distributions. Unlike typical calibration methods, our distillation is carried out solely on the training dataset without the need of an additional validation dataset. This simple yet effective approach achieves N-times more efficient uncertainty estimation during testing, where N is the number of samples traditionally required by Bayesian LLMs. Our extensive experiments demonstrate that uncertainty estimation capabilities on training data can successfully generalize to unseen test data through our distillation technique, consistently producing results comparable to (or even better than) state-of-the-art Bayesian LLMs.

Figures (7)

• Figure 1: Overview of Bayesian LLMs.a. Point estimation (e.g., LoRA hu2022lora) overfits to limited data, causing unreliable responses. b. Bayesian Low-Rank Adaptation (BLoB) wang2024blob uses variational Gaussian distributions for better generalization and confidence alignment. c. Training-Free Bayesianization (TFB) shi2024training enables direct conversion to Bayesian LLMs without complex training. d. Efficient Uncertainty Estimation via Distillation (EUD) distills uncertainty estimation to point estimation. Our approach serves as an important puzzle of test-time efficient deployment for Bayesian LLMs.
• Figure 2: Impact of Dataset Size on Student Model Performance on BoolQ boolq. Reducing the size of training data leads to higher ECE and NLL, and generally lower ACC, highlighting the importance of sufficient data for effective uncertainty distillation for EUD.
• Figure 3: Effect of Teacher Model Sample Size on EUD, evaluated on WG-S wg. Increasing the number of samples from the Bayesian teacher model during training significantly enhances the uncertainty estimation performance of the distilled deterministic student model of our EUD. Importantly, this sampling is performed only once during training and incurs no additional cost at test time, making EUD highly practical for real-world deployment of Bayesian models.
• Figure 4: Training dynamics of EUD on OBQA obqa.
• Figure 5: Training dynamics of EUD on WG-S wg.