Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Calibrating LLMs with Information-Theoretic Evidential Deep Learning

Yawei Li, David Rügamer, Bernd Bischl, Mina Rezaei

TL;DR

Fine-tuned LLMs often exhibit overconfidence on small datasets, undermining trust in their uncertainty estimates. The authors introduce IB-EDL, which regularizes Evidential Deep Learning with an information bottleneck objective to suppress spurious evidence while preserving predictive information, achieving calibration gains with minimal overhead. By selecting the pre-evidence $Z=\tilde{\mathbf{e}}$ and using a Gaussian prior over $Z$, they derive a tractable IB objective that includes an $\ell_2$-regularization effect on the Dirichlet parameters, and show VB-based EDL methods as special cases within IB-EDL. Extensive experiments across multiple LLMs, datasets, and scenarios (ID calibration, OOD detection, and label-noise robustness) demonstrate substantial improvements in ECE and NLL without sacrificing accuracy, enabling more trustworthy deployment of LLMs in high-stakes domains.

Abstract

Fine-tuned large language models (LLMs) often exhibit overconfidence, particularly when trained on small datasets, resulting in poor calibration and inaccurate uncertainty estimates. Evidential Deep Learning (EDL), an uncertainty-aware approach, enables uncertainty estimation in a single forward pass, making it a promising method for calibrating fine-tuned LLMs. However, despite its computational efficiency, EDL is prone to overfitting, as its training objective can result in overly concentrated probability distributions. To mitigate this, we propose regularizing EDL by incorporating an information bottleneck (IB). Our approach IB-EDL suppresses spurious information in the evidence generated by the model and encourages truly predictive information to influence both the predictions and uncertainty estimates. Extensive experiments across various fine-tuned LLMs and tasks demonstrate that IB-EDL outperforms both existing EDL and non-EDL approaches. By improving the trustworthiness of LLMs, IB-EDL facilitates their broader adoption in domains requiring high levels of confidence calibration. Code is available at https://github.com/sandylaker/ib-edl.

Calibrating LLMs with Information-Theoretic Evidential Deep Learning

TL;DR

Fine-tuned LLMs often exhibit overconfidence on small datasets, undermining trust in their uncertainty estimates. The authors introduce IB-EDL, which regularizes Evidential Deep Learning with an information bottleneck objective to suppress spurious evidence while preserving predictive information, achieving calibration gains with minimal overhead. By selecting the pre-evidence and using a Gaussian prior over , they derive a tractable IB objective that includes an -regularization effect on the Dirichlet parameters, and show VB-based EDL methods as special cases within IB-EDL. Extensive experiments across multiple LLMs, datasets, and scenarios (ID calibration, OOD detection, and label-noise robustness) demonstrate substantial improvements in ECE and NLL without sacrificing accuracy, enabling more trustworthy deployment of LLMs in high-stakes domains.

Abstract

Fine-tuned large language models (LLMs) often exhibit overconfidence, particularly when trained on small datasets, resulting in poor calibration and inaccurate uncertainty estimates. Evidential Deep Learning (EDL), an uncertainty-aware approach, enables uncertainty estimation in a single forward pass, making it a promising method for calibrating fine-tuned LLMs. However, despite its computational efficiency, EDL is prone to overfitting, as its training objective can result in overly concentrated probability distributions. To mitigate this, we propose regularizing EDL by incorporating an information bottleneck (IB). Our approach IB-EDL suppresses spurious information in the evidence generated by the model and encourages truly predictive information to influence both the predictions and uncertainty estimates. Extensive experiments across various fine-tuned LLMs and tasks demonstrate that IB-EDL outperforms both existing EDL and non-EDL approaches. By improving the trustworthiness of LLMs, IB-EDL facilitates their broader adoption in domains requiring high levels of confidence calibration. Code is available at https://github.com/sandylaker/ib-edl.

Paper Structure

This paper contains 29 sections, 1 theorem, 26 equations, 1 figure, 14 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. MLE fine-tuning of LLMs
  4. Uncertainty-aware modeling via EDL
  5. Training of EDL networks
  6. Information bottleneck-regularized EDL
  7. A Careful examination of the information bottleneck criterion
  8. Regularizing EDL with IB
  9. Practical implementation
  10. Variational Bayes-based EDL as a special case of IB-EDL
  11. Experiments
  12. Experimental setups
  13. In-distribution calibration
  14. Out-of-distribution detection
  15. Fine-tuning with noisy labels
  16. ...and 14 more sections

Key Result

Proposition 1

The VB-based EDL methods, which minimize $\mathbb{E}_{p(\bm{x}, \bm{y})}\left[D_{\text{KL}}\left(p(\bm{\pi} | \bm{x} ; \bm{\theta}) || p(\bm{\pi} | \bm{y}) \right) \right]$, are a special case of IB-EDL when the hidden variable is chosen as $\bm{z} = \bm{\pi}$ (i.e. token probabilities) and the prio

Figures (1)

  • Figure 1: Ablation study. IB-EDL reduces ECE and NLL compared to MAP across a broad range of $\beta$ and $K$ values. $\beta$ controls the regularization strength and balances the calibration and accuracy.

Theorems & Definitions (3)

  • Proposition 1
  • Remark 1
  • proof