📊 statistics

Toward Efficient Uncertainty in LLMs through Evidential Knowledge Distillation

This paper proposes an efficient method for uncertainty quantification in large language models by using LoRA-based knowledge distillation to transfer uncertainty-aware capabilities from computationally expensive, multi-pass teachers to single-pass students, achieving comparable performance without the heavy inference latency of traditional sampling methods.

Original authors: Lakshmana Sri Harsha Nemani, P. K. Srijith, Tomasz Kuśmierczyk

Published 2026-07-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Lakshmana Sri Harsha Nemani, P. K. Srijith, Tomasz Kuśmierczyk

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, but incredibly slow, expert consultant (the Teacher). This consultant is amazing at guessing the right answer to a question, but they are also very cautious. Before giving you an answer, they don't just guess once; they ask ten different versions of themselves, or they run the same question through ten different "lenses," to see how much their answers vary.

If all ten versions agree, the consultant says, "I'm 100% sure."
If the ten versions are all over the place, the consultant says, "I'm not sure at all."

This process of asking ten versions of themselves is how the paper describes sampling-based methods (like Bayesian models or Ensembles). It's great for knowing how sure the model is, but it's painfully slow. If you ask the consultant a question, you have to wait for them to run ten simulations before they give you an answer. In the world of AI, this is like waiting for a slow train when you just wanted a quick bus ride.

The authors of this paper asked: "Can we train a fast, single-pass student to learn from this slow, cautious expert?"

Here is how they did it, broken down into simple concepts:

1. The Goal: The Fast Student

They wanted to create a Student model that is just as smart as the Teacher but doesn't need to run ten simulations. The Student should be able to look at a question once, give an answer, and tell you how confident it is, all in a single instant.

To make the Student fast, they used a technique called LoRA (Low-Rank Adaptation). Think of this as giving the Student a set of "training wheels" or a lightweight backpack. Instead of rebuilding the Student's entire brain (which is huge and expensive), they just tweak a tiny, specialized part of it to learn from the Teacher.

2. The Two Types of Students

The researchers tried two different ways to teach the Student how to express uncertainty:

The "Average" Student (Softmax):
This student learns the average of what the Teacher's ten versions said. If the Teacher's versions were 50% "Yes" and 50% "No," this student learns to say "50% Yes."
- The Problem: This student can tell you the average, but it forgets the story behind the average. It doesn't know if the Teacher was confused (all versions disagreed) or if the data was just noisy. It collapses all that complex uncertainty into a single number.
The "Evidence" Student (Dirichlet):
This is the paper's main innovation. Instead of just learning an average, this student learns to think in terms of evidence.
- Imagine the Teacher's ten versions are like ten witnesses in a courtroom.
- The Softmax student just counts the votes: "6 said Guilty, 4 said Not Guilty."
- The Dirichlet student looks at the confidence of the witnesses. Did all 10 witnesses shout "Guilty" with high confidence? Or did they all mumble "Maybe Guilty" with low confidence?
- This student outputs a "Dirichlet distribution," which is a mathematical way of saying, "Here is my best guess, and here is a measure of how much evidence I have to back it up." If the evidence is low, the student knows it's uncertain, even if the guess looks clear.

3. The Magic Trick: Knowledge Distillation

The process of teaching the Student is called Knowledge Distillation.

The researchers let the slow Teacher run its ten simulations on a bunch of practice questions.
They took the results of those ten simulations (the average guess and the spread of opinions) and fed them to the Student.
The Student was then trained to mimic those results using only one forward pass.
They used a "stop-watch" method (Early Stopping) to ensure the Student didn't just memorize the Teacher's mistakes but actually learned the right patterns.

4. The Results: Speed vs. Smarts

The paper tested this on text classification tasks (like sorting reviews as "Positive" or "Negative").

Speed: The Teacher was slow because it had to run many simulations. The Student was 11 to 36 times faster because it only ran once.
Accuracy: The Student was just as good at getting the right answer as the Teacher.
Uncertainty: The Dirichlet Student was particularly good at knowing when it was unsure. It could tell the difference between "I don't know because the data is confusing" (Epistemic uncertainty) and "I don't know because the question is vague" (Aleatoric uncertainty).

The Bottom Line

The paper claims they successfully built a "fast lane" for uncertainty. They took a slow, heavy, sampling-based AI that is great at knowing what it doesn't know, and distilled its wisdom into a lightweight, single-pass model.

The Teacher: A slow, cautious expert who checks their work ten times.
The Student: A fast, efficient worker who learned to check their work once but still knows exactly how confident they should be.

The authors emphasize that this works for text classification (sorting text into categories). They do not claim this works for generating long stories or complex reasoning tasks yet, but they believe it proves the concept that you can have fast, uncertainty-aware AI without the heavy computational cost.

Technical Summary: Toward Efficient Uncertainty in LLMs through Evidential Knowledge Distillation

Problem Statement
Accurate uncertainty quantification is a critical challenge for Large Language Models (LLMs). While theoretically grounded methods such as Bayesian Neural Networks (BNNs) and ensemble-based approaches (e.g., Bayesian Prompt Ensembles) effectively model epistemic uncertainty and improve calibration, they suffer from prohibitive computational costs. These methods typically require multiple forward passes (sampling) during both training and inference to approximate predictive distributions via Monte Carlo sampling. This multi-pass requirement significantly increases inference latency, limiting the practical deployment of uncertainty-aware LLMs in real-time or large-scale applications.

Methodology
The authors propose a framework for Evidential Knowledge Distillation to transfer uncertainty estimates from computationally expensive, sampling-based "teacher" models to efficient, single-pass "student" models.

Teacher Models: The study utilizes two distinct sampling-based teachers:
- Bayesian Prompt Ensembles (BayesPE): An ensemble formulation over input prompts.
- Laplace LoRA: A Bayesian approach using Laplace approximation over LoRA weights.
  Both teachers generate predictive distributions by aggregating multiple samples (prompts or weight perturbations) to capture total uncertainty (aleatoric + epistemic).
Student Architecture: The students are pre-trained LLMs (Mistral 7B Instruct) fine-tuned using Low-Rank Adaptation (LoRA). This allows for efficient parameter updates without retraining the entire model. The distillation process initializes the student's LoRA matrices from the teacher (if applicable) and fine-tunes them using a modified classification head.
Distillation Strategies: The paper compares two student output parameterizations:
- Softmax Student: A baseline model that outputs a single categorical distribution (mean probabilities) via softmax. It minimizes the negative log-likelihood of the teacher's average predictions. While fast, it collapses all uncertainty into Shannon entropy, failing to distinguish between aleatoric and epistemic components.
- Evidential (Dirichlet) Student: This model employs Evidential Deep Learning, outputting the parameters ( $\alpha$ ) of a Dirichlet distribution. Instead of predicting class probabilities directly, it predicts concentration parameters derived from logits ( $\alpha_c = 1 + \text{softplus}(z_c)$ ). The total evidence ( $\alpha_0 = \sum \alpha_c$ ) serves as an explicit measure of confidence. This formulation allows the model to separate aleatoric uncertainty (inherent data noise) from epistemic uncertainty (model ignorance) within a single forward pass.
Training Objective: The students are trained to match the teacher's predictive distribution. The loss function for the evidential student is the negative log-likelihood of the teacher's sampled predictions under the predicted Dirichlet distribution.
- Regularization: The authors investigate various regularization strategies for the concentration parameter $\alpha_0$ , including fixing it to a constant, pulling it toward a uniform prior, or anchoring it to a learnable dataset-wide mean, to prevent overconfidence and stabilize training.
- Early Stopping: Training is halted based on the negative log-likelihood (NLL) on the ground-truth labels to prevent overfitting to teacher noise.

Key Results
Experiments were conducted on four text classification datasets (Amazon, SST2, Yahoo, YouTube) using the Mistral 7B backbone.

Inference Efficiency: The distilled students achieve a massive reduction in inference time compared to their teachers. For BayesPE teachers, the Dirichlet student is 11x to 36x faster (e.g., 36x speed-up on Yahoo), as it requires only one forward pass versus the multiple passes required by the teacher.
Predictive Performance: The distilled students match or occasionally exceed the predictive accuracy of their teachers. For instance, the Dirichlet student outperformed the BayesPE teacher on Yahoo (+1.7 pp) and YouTube (+2.5 pp) datasets.
Uncertainty Quantification:
- Calibration: The Dirichlet students generally achieve lower Expected Calibration Error (ECE) and Brier scores than both the teachers and the Softmax baselines.
- OOD Detection: In Out-of-Distribution (OOD) experiments, the Dirichlet student demonstrated superior ability to distinguish between in-domain and OOD samples (higher AUROC) compared to the teacher and Softmax student. It successfully amplified total predictive entropy for OOD data.
- Decomposition: Unlike the Softmax student, the Dirichlet student successfully decomposes uncertainty into aleatoric and epistemic components, providing richer uncertainty signals.
Convergence: Students distilled from BayesPE teachers converged very quickly (often within 1–5 epochs), whereas those from Laplace teachers required more epochs (20–40), suggesting a more complex optimization landscape for the latter.

Significance and Claims
The paper claims to be the first application of Evidential Deep Learning to the distillation of Large Language Models. Its primary contribution is establishing a foundation for efficient, single-pass uncertainty estimation in LLMs.

Efficiency vs. Accuracy Trade-off: The work demonstrates that it is possible to inherit the robust uncertainty modeling of computationally intensive sampling methods without incurring their inference latency.
Epistemic Modeling: By utilizing Dirichlet outputs, the method explicitly models epistemic uncertainty, a capability lost in standard softmax distillation.
Scope and Limitations: The authors modestly frame the current work as a proof-of-concept limited to discrete text classification tasks. They acknowledge that while the results are promising, extending this framework to open text generation remains a direction for future work. The study focuses on matching teacher performance rather than improving upon the teachers' absolute capabilities, though empirical results show students sometimes surpassing teachers due to the regularizing effects of the distillation process and early stopping.