Calibrating LLMs with Information-Theoretic Evidential Deep Learning

Authors: Yawei Li, David Rügamer, Bernd Bischl, Mina Rezaei
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
Keywords: evidential deep learning; information bottleneck; calibration; large language models

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.

Decision: accept (poster)
Reviewer scores: 8, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730687297847

thread_type:
official_review

<thread_object>
creator: reviewer_4vtx

<content>
summary:
The paper presents the Information Bottleneck (IB) - Evidential Deep Learning (EDL) method, designed to regularize EDL for better calibration of large language models (LLMs). The empirical experimental results demonstrate that IB-EDL enhances calibration in fine-tuned LLMs.

strengths:
The paper is well-written and easy to understand, trying to address the important task of calibrating fine-tuned large language models. The empirical experiments are comprehensive, testing across multiple LLMs including Llama2-7B, Llama3-8B, and Mistral-7B. The paper covers most of the baselines. The LoRA-based Mixture of Experts idea [1] could also be relevant and serve as a baseline for comparison. The empirical performance appears promising.

[1] Li, Dengchun, et al. "Mixlora: Enhancing large language models fine-tuning with lora based mixture of experts."

weaknesses:
The novelty of the paper is limited, as both EDL and the IB method are well studied in the literature. The primary contribution of the paper appears to be the implement IB on EDL. The authors explain that the IB objective allows the latent variables to retain the most predictive information about the target variable while discarding irrelevant information from X. However, a more detailed discussion on how this relates to improved calibration of LLMs is missing. Additionally, there is a lack of theoretical guarantees supporting this approach. It is beneficial to provide a theoretical bound on calibration error or a proof of convergence.

In terms of the empirical analysis, the paper highlights computational efficiency as a key contribution of the proposed method, yet it lacks experimental comparisons of computational time and memory costs to validating this claim. The paper also lacks the calibration performance analysis in out-of-distribution scenarios.

questions:
Please refer to the strengths and weaknesses. Additionally, how many bins did the authors use to compute the ECE? It is beneficial to conduct a sensitivity analysis on the number of bins.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730690845847

thread_type:
official_review

<thread_object>
creator: reviewer_qw6t

<content>
summary:
In this work, authors propose a new loss regularization term to reduce overconfidence in LLMs when training via the evidential deep learning (EDL) framework. This regularization is based on the information bottleneck loss, which maximizes mutual information between token logits and ground truth prediction distribution while minimizing M.I. between predictions and inputs. This is achieved by minimizing a KL divergence between a learned conditional logit distribution and a Gaussian prior. This work demonstrates that fine-tuning via this loss produces better calibration (ECE, AUC) with equivalent accuracy to baselines. Authors also illustrate this method is robust to noise in the training data and distribution shifts.

strengths:
* This work is well-motivated theoretically, pointing out deficiencies with previous EDL losses (e.g. paragraph *Challenges when applying IB to an internal layer of an LLM*). 
* The final regularization term $\mathcal{L}_{IB-Info}$ seems relatively simple, although I do have some questions about practical implementation (see questions)
* Strong experimental results; the OOD detection and noise injection experiments (tables 3,4) in particular seem support the argument for this additional regularization during calibration

weaknesses:
* This work seems only applicable for a discrete set of mutually exclusive classes. I'm not sure how this extends to the case of open-ended generation where multiple responses may be semantically equivalent or have different logical relationships.
* While experiments show improvements in OOD detection when moving from OBQA--> ARC,CSQA datasets, I'd be interested to see the performance when the datasets are 'futher apart' semantically (e.g. moving from a reading comprehension to a math task).

questions:
* While it appears the training methods/baselines tested are normalized for number of training steps, I am curious about comparison in terms of the number of training flops. Do the new mean and covariance prediction heads require separate training? 
* Is there a benefit to having two parameters estimating the distributional shift (mean and variance) of the token predictions vs. just one?
* What happens to calibration if we finetune on multiple datasets simultaneously?
* How does this fine-tuning procedure effect LLM performance on open-ended generation tasks?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730722790963

thread_type:
official_review

<thread_object>
creator: reviewer_tujz

<content>
summary:
This paper proposes IB-EDL, combining Evidential Deep Learning (EDL) with information bottleneck (IB) for uncertainty estimation in large language models (LLMs). This approach aims to enhance model calibration by reducing overconfidence, which is a common issue in fine-tuned LLMs. The experimental results show that IB-EDL achieves improved calibration with minimal computational overhead.

strengths:
The paper is well-presented and easy to follow.

The paper proposes a unified perspective on EDL methods by framing several existing EDL methods as special cases within the IB-EDL framework.

The experimental results are strong and covers multiple LLMs, showing that IB-EDL addresses the overconfidence problem well.

weaknesses:
The pipeline of applying IB to EDL ($x → f(x; θ) → e˜ → e →
α → π → y$) seems computational heavy, it would be informative to include an inference time comparison of IB-EDL against other methods. 

Although the training overhead of IB-EDL is relatively small compared with the pretraining step, the paper lacks comparison of training time compared with LoRA, which is the fine-tuning backbone of IB-EDL.

For out-of-distribution (OOD) scenarios, the paper only compares the OOD detection capabilities with AUROC. The expected calibration error (ECE) is missing, which is also an indicator of OOD performance as the calibrated model should maintain relatively low ECE score on OOD datasets.

The novelty is not clear for combining IB with EDL, which are two well established methods.

questions:
Please see weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730731248405

thread_type:
official_review

<thread_object>
creator: reviewer_ttkc

<content>
summary:
The paper present a method to effectively apply Evidential Deep Learning (EDL), a parametric uncertainty estimation approach, to LLMs. The approach introduces two features to standard EDL that make it practical and effective for application to LLMs; i) the introduction of an information bottleneck cost function, which mitigates over-fitting and ii) the use of the LLM outputs as the intermediate "layer" Z where to impose the bottleneck, which simplifies direct computation of certain terms in the mutual information lower bounds, part of the objective function.

strengths:
- Relevance: The paper addresses an important problem; the uncertainty calibration of LLM outputs (at the token level), which is an important component for providing safe deployment and error estimation for LLMs.

- Theoretical Soundness: The paper takes a very principled approach, identifying precise issues in the application of EDL to LLMs and addressing them with mathematically derived solutions based on reasonable starting assumptions.

- Experimental Soundness: The experiments are more than sufficiently extensive: they target the right settings and metrics to validate the claims of the paper and comprehend great breadth of models, tasks and mode of evaluation (acc+ECE, OOD detection). The claims of improved uncertainty calibration are strongly backed by the experiments.

weaknesses:
The only main weakness of the paper is the clarity of scope:

From the experiments, it is clear that this is a method for uncertainty calibration for a specific task at the fine-tuning stage of an LLM. e.g., an LLM is fine-tuned for summarisation and the proposed method offers a way to obtain superior uncertainty calibration *for the fine-tuned task*. This is indeed very useful, but needs to be stated more clearly in the abstract and intro, as from all sections up until the experiments it is unclear if this is a method for pre-training and/or fine-tuning. In principle, I can see the method may be applied to pre-training too, but since this is not tested in the experiments it should not be included in the claims.

questions:
See and clarify the weakness point above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732299267388

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Key Updates of Manuscript

comment:
**Dear Reviewers**,

As the deadline for the discussion phase approaches, we have submitted an improved manuscript incorporating your constructive comments and feedback. Below, we summarize the key updates made:

### Key Updates

**Additional Experiments**: We appreciate your recognition of the extensive experiments included in our original submission. Building on this foundation and incorporating your valuable feedback, we have added the following additional experiments: (1) OOD detection under large distribution shifts, (2) Calibration performance on OOD datasets, (3) Analysis of training and inference speed, as well as memory consumption, (4) Sensitivity analysis of the number of ECE bins.

These updates are detailed in Appendix D.2–D.6, with corresponding references and discussions added in the main text.

**Highlighting Novelty**: To better articulate our contributions, we have revised the introduction to emphasize the novelty of our approach. Specifically:

Our method goes beyond a straightforward implementation of IB on EDL. Importantly, it introduces several key advancements:

* **New theoretical insights**: We uncover a theoretical issue overlooked in previous IB literature.

* **Innovative solution**: We address this issue in the context of EDL by altering the IB's latent variable parametrization.

* **Effective regularization**: Using regularization on the evidence, we mitigate the problem of overconfident evidence values reported in previous EDL studies.

* **Unified framework**: Our IB-EDL framework unifies several previous EDL methods, offering a comprehensive perspective and advancements to the field.

* **Scalability to LLMs**: We scale EDL methods to LLMs, extending their applicability beyond the medium-sized neural networks traditionally used in EDL literature. This significantly enhances the potential for developing reliable, modern large-scale applications.

**Additional Related Work**: In response to Reviewer 4Vtx, we have expanded the related work section to include MixLoRA.

**Clarification of Research Scope**: As suggested by Reviewer TTkc, we have revised the introduction to clarify the research scope and focus of our work.

---

If you have any further concerns or suggestions, please do not hesitate to let us know. We deeply appreciate your valuable insights and guidance throughout the review process. Thank you again for your time and effort in reviewing our work.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper present a method to apply evidential deep learning, a parametric uncertainty estimation approach, to LLMs. The proposed approach introduces an information bottleneck cost function over the LLM outputs as well as the model parameter. The reviewers raised points about the applicability to more general settings, computataional complexity, and clarity of the presentation. During the rebuttal, the authors have addressed these questions by revising the paper and by providing additional experiments. All the reviewers are positive about the paper, after the rebuttal period.