PaperlayerView exact revision on arXiv

Source-linked AI summary

Calibrating LLMs with Information-Theoretic Evidential Deep Learning

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

arXiv:2502.06351v2cs.LG

TL;DR

Fine-tuned LLMs can be overconfident on small datasets, and EDL objectives can drive predicted distributions toward overly concentrated forms. IB-EDL regularizes EDL with an information bottleneck; experiments report improved calibration across LLMs and datasets while maintaining model performance.

  • Problem

    Fine-tuned LLMs can be overconfident on small datasets, and EDL objectives can drive predicted distributions toward overly concentrated forms.

    Source: §1 · p. 1, §2.3 · p. 3

  • Method

    IB-EDL regularizes EDL with an information bottleneck to suppress spurious evidence while preserving model performance.

    Source: §1 · p. 1

  • Results

    Across various LLMs and datasets, IB-EDL scales effectively; it shows the lowest ECE and NLL in reported Llama2-7B and Llama3-8B evaluations.

    Source: §1 · p. 2, §4.2 · p. 8

  • Takeaways & Limitations

    IB-EDL improves calibration while maintaining model performance, including in the presence of substantial label noise.

    Source: §6 · p. 11

    Evaluations primarily cover conventional classification, and the method assumes uncorrelated pre-evidences; uncertainty metrics for generative tasks remain under development.

    Source: §6 · p. 11

Abstract

from arXiv · show

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.

1 INTRODUCTION

Fine-tuned LLMs can be overconfident, undermining their reliability in critical domains. The paper proposes IB-EDL, an information-bottleneck regularizer for EDL, and reports broad experiments on fine-tuned LLMs.

1 INTRODUCTION

  • Fine-tuned LLMs often show overconfidence; conventional uncertainty-aware methods can require multiple forward passes, while EDL can still generate excessive evidence and inaccurate uncertainty estimates.Excessively large evidence can produce overly confident class probabilities.Source: §1 · p. 1
  • IB-EDL adaptively distorts generated evidence to suppress spurious or uninformative evidence while preserving model performance, with less than 2% computational overhead.The authors report improved calibration and theoretically motivate penalizing disproportionately large evidence.Source: §1 · pp. 1–2
  • IB-EDL frames EDL regularization information-theoretically and is evaluated across LLMs and datasets, where the results demonstrate effective scaling.The framework also unifies several existing EDL methods as special cases.Source: §1 · p. 2

2 BACKGROUND

The background contrasts deterministic MLE fine-tuning with uncertainty-aware EDL, then describes EDL’s training objectives and their vulnerability to overconfident predictions. IB regularization addresses that vulnerability without changing EDL’s assumptions about the target or Dirichlet parameters.

2.1 MLE FINE-TUNING OF LLMS

  • MLE fine-tuning maps LLM logits through Softmax to token probabilities and minimizes negative log-likelihood, but small datasets can yield overfitting and overconfident, deterministic predictions.The paper uses “tokens” for target classes and treats sequence targets as multiple single-token targets in its theoretical analysis.Source: §2.1 · p. 2

2.2 UNCERTAINTY-AWARE MODELING VIA EDL

  • EDL predicts pre-evidence, converts it to nonnegative evidence and Dirichlet parameters, then uses the resulting distribution to estimate token probabilities and prediction uncertainty in one forward pass.Belief masses represent support for individual tokens, while uncertainty mass captures overall uncertainty; their sum is normalized.Source: §2.2 · pp. 2–3

2.3 TRAINING OF EDL NETWORKS

  • EDL training minimizes expected loss under the predicted Dirichlet distribution, with cross-entropy or MSE objectives and a KL regularizer that suppresses non-target evidence.The regularization term is defined after removing evidence associated with the target token.Source: §2.3 · p. 3
  • EDL objectives can drive the Dirichlet distribution toward a point mass and extreme target-token evidence; IB instead regularizes the model against irrelevant correlations without changing EDL’s assumptions on y or α.The existing non-target-evidence penalty does not fully address excessive target evidence.Source: §2.3 · p. 3

3 INFORMATION BOTTLENECK-REGULARIZED EDL

The section develops an information-bottleneck objective for EDL and applies it to pre-evidence, aiming to retain predictive information while compressing irrelevant input information. It also explains the practical training procedure and shows how variational Bayes-based EDL fits within the framework.

3.1 A CAREFUL EXAMINATION OF THE INFORMATION BOTTLENECK CRITERION

  • The IB objective balances information Z retains about target Y against information it retains about input X, with β controlling the relevance–compression trade-off.The section also warns that replacing p(y|z) with a learned q(y|z) at an intermediate LLM layer can make the proposed lower bound invalid when their gap is negative.Source: §3.1 · pp. 4–5

3.2 REGULARIZING EDL WITH IB

  • Choosing pre-evidence ˜e as Z lets IB-EDL compute p(y|z) directly through the EDL pipeline, avoiding a learned approximation q(y|z) and preserving a valid lower bound on I(Z, Y).The resulting objective imposes ℓ2 regularization on the mean of ˜e and thus on α, penalizing large α associated with overconfident predictions.Source: §3.2 · pp. 5–6
  • IB-EDL uses a standard Gaussian prior for pre-evidence ˜e and models its conditional distribution with predicted mean µ and variance σ.The authors motivate the prior by the typical -2 to 2 range of pre-evidence and reported similarities between later-layer activations and Gaussian distributions.Source: §3.2 · pp. 5–6

3.3 PRACTICAL IMPLEMENTATION

  • Training samples K pre-evidence values using reparameterization and averages their losses; inference also samples K values and averages pre-evidence.The paper gives K = 20 as an example and describes the extra inference time as minimal relative to the LLM.Source: §3.3 · p. 6

3.4 VARIATIONAL BAYES-BASED EDL AS A SPECIAL CASE OF IB-EDL

  • VB-based EDL methods that minimize the stated divergence are a special case of IB-EDL when Z is token probabilities and the prior is Dir(z; y ⊙ α + (1 − y)).Source: §3.4 · p. 6

4 EXPERIMENTS

Across calibration, OOD detection, and noisy-label evaluation, IB-EDL performs favorably against the compared methods while maintaining accuracy comparable to MAP. Ablations show gains across a broad range of β and K values, with low reported computational overhead.

4.1 EXPERIMENTAL SETUPS

  • Experiments use fine-tuned Llama2-7B, Llama3-8B, and Mistral-7B on six multiple-choice datasets, comparing IB-EDL with MAP, uncertainty-aware methods, and EDL baselines.The main text primarily reports Llama2-7B and Llama3-8B results; Mistral-7B results are in the appendix.Source: §4.1 · pp. 6–7

4.2 IN-DISTRIBUTION CALIBRATION

  • IB-EDL achieves the lowest ECE and NLL while maintaining accuracy on par with or higher than MAP on Llama2-7B and Llama3-8B.The authors attribute its calibration advantage over other EDL methods to ℓ2 regularization discouraging excessively large evidence.Source: Table 1 · p. 7, §4.2 · p. 8, Table 2 · p. 8

4.3 OUT-OF-DISTRIBUTION DETECTION

  • IB-EDL achieves the highest AUROC across all tested OOD datasets using both max probability and uncertainty mass, surpassing EDL and non-EDL competitors.Evaluation trains on OBQA and tests on ARC-C, ARC-E, and CSQA; non-EDL methods without uncertainty mass are assessed with max probability only.Source: §4.3 · pp. 8–9, Table 3 · p. 9

4.4 FINE-TUNING WITH NOISY LABELS

  • With 30% of training labels perturbed, IB-EDL achieves the highest overall accuracy among the compared methods.Models are evaluated on clean test sets; the authors suggest the information bottleneck filters spurious signals while retaining predictive information.Source: §4.4 · p. 9, Table 4 · p. 10

4.5 ABLATION STUDY

  • IB-EDL reduces ECE by at least 40% and NLL by at least 50% versus MAP for all tested β values, and outperforms MAP across a broad range of β and K.Increasing K slightly improves accuracy but does not necessarily improve ECE or NLL; β balances regularization and predictive performance.Source: Figure 1 · p. 9, §4.5 · p. 10
  • For Llama2-7B with the full vocabulary as target space, IB-EDL adds 1.95% parameters and 1.98% of pretrained-model GFLOPs.The reported overhead comes from an additional linear head and evidence averaging.Source: §4.5 · p. 10

5 RELATED WORK

Related work covers EDL methods and their regularization strategies, other uncertainty-estimation approaches, and information bottleneck applications. IB-EDL is distinguished by regularizing without changing EDL assumptions or requiring uniformity in non-target Dirichlet parameters.

EDL

  • EDL methods predict Dirichlet priors and use varied losses, regularizers, or formulation changes; IB-EDL instead regularizes without requiring non-target uniformity or altering EDL assumptions.The cited approaches include variational Bayesian EDL and Posterior Networks using Normalizing Flows.Source: §5 · p. 10

Other uncertainty-aware methods

  • Other uncertainty-aware approaches include variational Bayesian neural networks, MC-Dropout, stochastic gradient MCMC, and Laplace approximations, including extensions to LoRA-fine-tuned LLMs.The cited LLM extensions include Yang et al. (2024), Wang et al. (2024a), and Li et al. (2024).Source: §5 · p. 10

IB

6 CONCLUSION

The conclusion presents IB-EDL as an information-theoretic regularizer that improves calibration with minimal computational overhead and maintains performance under substantial label noise. It identifies uncorrelated pre-evidences and conventional classification evaluations as limitations, with generative-task evaluation left for future work.

6 CONCLUSION

  • IB-EDL combines information-theoretic regularization with minimal computational overhead and maintains model performance under substantial label noise.The conclusion presents it as a promising approach to more trustworthy LLMs.Source: §6 · p. 11
  • The method assumes pre-evidences are uncorrelated to reduce covariance-prediction complexity, although the assumption could be relaxed.The evaluations primarily cover conventional classification, and generative tasks remain future work amid ongoing uncertainty-metric research.Source: §6 · p. 11

A DERIVATION OF THE VARIATIONAL BOUNDS

This appendix reproduces derivations for variational bounds on I(Z, X) and I(Z, Y), including the assumptions and information quantities used.

A DERIVATION OF THE VARIATIONAL BOUNDS

B PROOF OF PROPOSITION 1

The proof connects the IB objective to the target optimized by variational Bayes-based EDL methods. Under specified latent-variable and prior choices, the IB objective upper-bounds that target.

B PROOF OF PROPOSITION 1

C DETAILED IMPLEMENTATION

The implementation fine-tunes three LLMs with LoRA and specifies shared training settings. IB-EDL uses jointly trained prediction heads and a β-selection guideline.

C DETAILED IMPLEMENTATION

  • The experiments fine-tune Llama2-7B, Llama3-8B, and Mistral-7B-v0.1.LoRA adaptors are applied to the models' q head modules; the passage also reports dropout p = 0.1, LoRA α = 16, and rank r = 8.Source: C DETAILED IMPLEMENTATION · p. 20
  • Training runs for 30,000 steps on CSQA and 10,080 steps on other datasets, using learning rate 0.00005 with cosine annealing.The passage also specifies token-length limits, bfloat16 precision, and method-specific settings for MCD, Ens, and other baselines.Source: C DETAILED IMPLEMENTATION · p. 20
  • IB-EDL jointly trains LoRA-equipped linear heads for μ and σ with transformer-layer LoRA weights, sampling K = 20 pre-evidences by default.The heads are initialized from the pretrained LLM's linear head, and sampling is used during training and inference.Source: C DETAILED IMPLEMENTATION · p. 20
  • The authors suggest larger β when MAP has lower accuracy and higher ECE or with label noise; otherwise, smaller β can pass more input information through pre-evidences.β is optimized by grid search; OOD experiments reuse the OBQA parameters.Source: C DETAILED IMPLEMENTATION · p. 20

D A POST-HOC CALIBRATION TECHNIQUE FOR IB-EDL

The post-hoc technique adjusts evidence to account for pre-evidence variance, while noisy-data experiments report IB-EDL's accuracy and calibration advantage among EDL methods.

D A POST-HOC CALIBRATION TECHNIQUE FOR IB-EDL

  • The post-hoc adjustment subtracts ζ·σ_j from α_j, making belief more conservative when predicted pre-evidence is uncertain.The adjustment is intended to prevent excessively large α_j and improve calibration; it also incorporates pre-evidence variance into uncertainty quantification.Source: D A POST-HOC CALIBRATION TECHNIQUE FOR… · p. 21
  • The sign of ζ is chosen from the calibration curve: positive for overconfidence and negative for underconfidence.The technique is optional and was omitted from the main-text experiments.Source: D A POST-HOC CALIBRATION TECHNIQUE FOR… · p. 21

E.1 ID CALIBRATION AND OOD DETECTION RESULTS OF MISTRAL-7B

  • For fine-tuned Mistral-7B, Table 6 covers ID calibration and Table 7 covers OOD detection.The supplied passages identify these tables' scope but do not report their results.Source: Table 6 · p. 21, §E.1 · p. 22, Table 7 · p. 22

E.2 ADDITIONAL RESULTS OF FINE-TUNING MODELS ON NOISY DATASETS

  • With label noise, IB-EDL has the highest accuracy and lowest ECE and NLL among EDL approaches.Non-EDL methods tend to underperform EDL methods in accuracy in this setting.Source: §E.2 · p. 22

E.3 ADDITIONAL OOD DETECTION RESULTS ON MMLU-MATH DATASET

On the OBQA-to-MMLU-Math shift, IB-EDL leads OOD detection with both scores; on other OOD test sets, it has the best ECE and NLL on two of three.

E.3 ADDITIONAL OOD DETECTION RESULTS ON MMLU-MATH DATASET

  • The MMLU-Math OOD test set uses math-related questions from MMLU, while OBQA is the in-distribution dataset.The authors characterize this as a substantially larger distribution shift than OBQA→ARC-C.Source: §E.3 · p. 22
  • IB-EDL achieves the best OOD detection performance with both MP and UM on the OBQA-to-MMLU-Math setting.The setting shifts from common-sense reasoning to math-focused questions and is described as substantially larger than OBQA→ARC-C.Source: §E.3 · p. 22

E.4 ADDITION CALIBRATION RESULTS ON OOD TEST SETS

  • On two of three OOD test sets, IB-EDL achieves the best ECE and NLL after fine-tuning Llama3-8B on OBQA.The evaluation tests whether calibration generalizes to OOD datasets.Source: §E.4 · p. 23

E.5 ANALYSIS ON TRAINING AND INFERENCE SPEED, AND MEMORY CONSUMPTION

IB-EDL has speeds and memory use comparable to MAP and other EDL baselines, with faster inference than methods requiring multiple passes or inference-time gradients; ECE rankings are stable across bin counts.

E.5 ANALYSIS ON TRAINING AND INFERENCE SPEED, AND MEMORY CONSUMPTION

  • IB-EDL has comparable training and inference speeds and memory consumption to MAP and other EDL baselines, and faster inference than MCD, Ens, and LA.MCD and Ens require multiple forward passes, while LA uses gradient computation during inference.Source: §E.5 · p. 23, Table 13 · p. 26
  • The efficiency comparison uses Llama3-8B in mixed precision on OBQA, with experiments conducted on one NVIDIA H100 GPU.OBQA is used as both the training and test dataset.Source: §E.5 · p. 23

E.6 SENSITIVITY ANALYSIS ON THE NUMBER OF BINS OF ECE

  • With 10, 15, 25, or 35 ECE bins on Llama3-8B and OBQA, ECE rises slightly as bins increase, while method rankings remain largely consistent.The experiment uses 15 bins by default.Source: §E.6 · p. 24, Table 14 · p. 27

Related papers

Loading 2502.06351v2…