Efficient Uncertainty in LLMs through Evidential Knowledge Distillation

TL;DR

This paper presents a novel evidential knowledge distillation method that enables large language models to efficiently estimate uncertainty with a single forward pass, maintaining high performance without costly sampling.

Contribution

The paper introduces a new evidential distillation approach for LLMs that achieves accurate uncertainty quantification using only one forward pass, reducing computational costs significantly.

Findings

Students match or outperform teachers in uncertainty estimation.

Evidential distillation achieves comparable predictive performance.

Single-pass uncertainty estimation is effective in LLMs.

Abstract

Accurate uncertainty quantification remains a key challenge for standard LLMs, prompting the adoption of Bayesian and ensemble-based methods. However, such methods typically necessitate computationally expensive sampling, involving multiple forward passes to effectively estimate predictive uncertainty. In this paper, we introduce a novel approach enabling efficient and effective uncertainty estimation in LLMs without sacrificing performance. Specifically, we distill uncertainty-aware teacher models - originally requiring multiple forward passes - into compact student models sharing the same architecture but fine-tuned using Low-Rank Adaptation (LoRA). We compare two distinct distillation strategies: one in which the student employs traditional softmax-based outputs, and another in which the student leverages Dirichlet-distributed outputs to explicitly model epistemic uncertainty via evidential learning. Empirical evaluations on classification datasets demonstrate that such students can achieve comparable or superior predictive and uncertainty quantification performance relative to their teacher models, while critically requiring only a single forward pass. To our knowledge, this is the first demonstration that immediate and robust uncertainty quantification can be achieved in LLMs through evidential distillation.