Calibrating LLMs with Information-Theoretic Evidential Deep Learning

We sincerely thank the reviewer for recognizing our contributions and for the constructive feedback. We address the concerns as follows:

W1.1: Scope of the work.

We thank the reviewer for pointing this out. Our work focuses on calibrating fine-tuned LLMs and enhancing their uncertainty awareness. We agree that emphasizing this focus will improve the clarity of our work. We will explicitly emphasize the scope of our work by revising the first bullet point in the list of contributions in the paper's introduction to “We introduce IB-EDL, an information-theoretic approach to improve calibration in fine-tuned LLMs.

W1.2: Applicability of IB-EDL to pre-training.

IB-EDL can theoretically be applied to pre-training. However, we conducted experiments on fine-tuning LLMs due to two primary reasons: (i) pre-training LLMs demands immense computational resources and has a substantial carbon footprint, and (ii) compared to pre-trained LLMs, fine-tuned LLMs present a more pressing need for improved calibration [1, 2]. That said, we agree that exploring the application of IB-EDL to pre-training represents an exciting and valuable avenue for future research. We will add a corresponding sentence to our conclusion. We thank the reviewer for highlighting this potential extension of our work.

We thank the reviewer again for the constructive feedback. Please let us know if our response addresses the concerns and if there is anything else we need to clarify.

[1] Yang, Adam X., et al. "Bayesian Low-rank Adaptation for Large Language Models." The Twelfth International Conference on Learning Representations.

[2] Achiam, Josh, et al. "Gpt-4 technical report." arXiv preprint arXiv:2303.08774 (2023)

Dear Reviewer TTkc,

Thank you once again for your thoughtful and constructive feedback on our manuscript. We are truly encouraged by your recognition of the soundness of our theoretical approach, the relevance of our problem, and the robustness of our experimental results.

As the discussion period will conclude in approximately 72 hours, we wanted to check if our responses thus far have adequately addressed your concerns and clarified your questions.

In summary, we have clarified the scope of our method and that it is specifically designed to improve uncertainty calibration in fine-tuned LLMs. We have revised our introduction to explicitly emphasize this focus, ensuring the scope is clear from the outset.

If there are any remaining concerns or questions, we would be happy to address them.

Best regards,

Authors of IB-EDL

Q1: 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.

Following the implementation of LA [7], we use n_bins = 15 by default for computing ECE. To address this concern, we provide a sensitivity analysis using Llama3-8B on OBQA, as shown in Table F. The results show that while ECE values increase slightly as the number of bins increases, the relative rankings of the methods basically remain consistent.

Table F: Sensitivity Analysis of the Number of Bins

—

Thank you again for your constructive feedback. Please let us know if our response addresses your concerns, and if there is anything else we need to clarify.

[2] Deng, Danruo, et al. "Uncertainty estimation by fisher information-based evidential deep learning." International Conference on Machine Learning. PMLR, 2023.

[3] Chen, Mengyuan, Junyu Gao, and Changsheng Xu. "R-EDL: Relaxing Nonessential Settings of Evidential Deep Learning." The Twelfth International Conference on Learning Representations.

[4] Alemi, Alexander A., et al. "Deep variational information bottleneck." arXiv preprint arXiv:1612.00410 (2016).

[5] Wieczorek, Aleksander, and Volker Roth. "On the difference between the information bottleneck and the deep information bottleneck." Entropy 22.2 (2020): 131.

[6] Kawaguchi, Kenji, et al. "How does information bottleneck help deep learning?." International Conference on Machine Learning. PMLR, 2023.

[7] Yang, Adam X., et al. "Bayesian Low-rank Adaptation for Large Language Models." The Twelfth International Conference on Learning Representations.

W4: 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

In Section 4.5, we provided a comparison of FLOPs and model parameters. Here, we offer a more comprehensive analysis, including training and inference time, as well as memory usage. Specifically, we evaluate IB-EDL using Llama3-8B on the OBQA dataset with a single NVIDIA H100 GPU. As shown in Table A, IB-EDL’s training and inference speeds, as well as its memory consumption, are comparable to MAP and other EDL methods. This confirms the computational efficiency of our approach.

Table A: Additional Complexity Analysis

W5: The paper also lacks the calibration performance analysis in out-of-distribution scenarios.

Thank you for raising this point. We provide calibration results on OOD datasets in Tables B, C, and D. These results demonstrate that IB-EDL achieves the best ECE and NLL on two out of three OOD datasets. This indicates that IB-EDL’s calibration abilities generalize well to distribution shifts, further validating its robustness.

Table B: Calibration in the setting OBQA -> ARC-C

Table C: Calibration in the setting OBQA -> ARC-E

Table D: Calibration in the setting OBQA -> CSQA

We thank the reviewer for the constructive comments and feedback. We address the concerns as follows:

S1: The LoRA-based Mixture of Experts idea [1] could also be relevant.

We appreciate the reviewer for bringing this work to our attention. We will include it in the Related Work in a revised version of the manuscript, which we will upload later during the discussion phase.

W1: 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.

Our work is not a simple combination of these methods; instead, it tackles key unresolved issues in EDL (see also [2, 3]) and addresses the emerged challenge when regularizing EDL with IB. Importantly, implementing IB on EDL in a naïve manner, such as inserting IB at an arbitrary layer, can compromise theoretical soundness. Our contributions go beyond mere implementation and address these challenges through the following advancements:

1. Uncovering theoretical limitations: We identify a fundamental issue overlooked in prior literature [4, 5]: applying IB at arbitrary layers in an LLM can result in an optimization objective that no longer serves as a valid lower bound of $I(Z, Y)$. In Section 3.1, we provide a detailed theoretical analysis (Eq. 11) that highlights this limitation.

2. Proposing a novel solution: In Section 3.2, we propose a unique solution by selecting the pre-evidence $\tilde{e}$ as the IB variable. This choice ensures that the optimization objective remains theoretically sound. Furthermore, as shown in Eq. (14), our method introduces L2 regularization on evidence values, reducing overconfidence by penalizing extreme outputs. This directly addresses critical issues in EDL identified in prior works [2, 3].

3. Unifying prior approaches: Through Proposition 1, we provide a cohesive framework that unifies prior approaches in EDL, offering deeper insights and a clearer understanding of their relationship to IB. This unified perspective is novel.

W2: 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.

Overparameterized LLMs often exhibit overconfidence when fine-tuned on small datasets or limited human feedback, primarily due to overfitting to spurious correlations in the training data. We utilize IB to address this issue by encouraging the model to focus on truly predictive features and discard irrelevant or spurious correlations. This, in turn, helps the model to learn more generalizable features and reduces overconfident predictions, leading to improved calibration. We will make the connection more clear in a revised manuscript version and thank the reviewer for pointing this out.

W3: 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.

1. Theoretical guarantees in our work: We provide several theoretical guarantees: (i) Through our analysis in Eq. (11) and Section 3.2, we ensure that our optimization objective remains a valid lower bound of $I(Y, Z)$; (ii) Through Eq. (14), we show that our approach penalizes overconfident predictions by introducing L2 regularization on evidence, preventing extreme values.

2. On calibration error bounds: Although deriving a formal bound on calibration error would further strengthen our contributions, it requires further analysis and is beyond the scope of this work. In particular, theoretical bounds for calibration error have not been studied in EDL or IB literature so far. Nonetheless, using Theorem 2 from [6], we can at least infer a generalization bound for IB-EDL. For any $\delta > 0$, the theorem guarantees that with probability at least $1 - \delta$, the generalization bound $\Delta$ for $n$ data points is $\Delta(n) \leq \min_l Q_l$, where $Q_l$ is a bounded term dependent on the $l$-th layer and number of observations $n$. The theorem by [6] is the tightest guarantee currently known with linear growth rates depending on $I(X,Z|Y)$, which is indirectly minimized by our objective. So we could apply the convexity of $\min$ function to derive a generalization bound $\Delta(n) \leq Q_L$, where $L$ is the layer index of the pre-evidence $\tilde{e}$ used in IB-EDL.

Dear Reviewer 4Vtx,

Thank you for providing valuable feedback on our manuscript. We greatly appreciate your recognition of our effort in the important task of calibrating LLMs and our strong experiment results.

As the discussion period is set to conclude in approximately 72 hours, we wanted to check if our responses thus far have sufficiently addressed your concerns and clarified your questions.

In summary, we addressed the points raised by you as follows:

• Incorporating relevant related work: We have added MixLoRA in the Related Work section of a revised manuscript.

• Novelty of our contributions: We clarified that our method goes beyond a naive implementation of IB on EDL, we uncovered theoretical challenges overlooked by previous IB literature (Section 3.1), and we address these challenges through a novel IB variable parameterization to maintain theoretical soundness (Section 3.2).

• Connection between IB and LLM calibration: We elaborated on why EDL is suitable for calibrating LLMs and why IB can further enhance the calibration of EDL methods.

• Theoretical guarantees: While deriving explicit bounds on ECE is beyond the scope of our work, we discussed existing theoretical guarantees related to generalization bounds that indirectly support our approach.

• Computational efficiency: We provided additional analyses of training and inference time, as well as memory usage (Table A), which demonstrate that IB-EDL’s computational requirements are comparable to MAP and other EDL baselines.

• Calibration in OOD scenarios: We addressed this concern by presenting additional experiments on OOD datasets (Tables B, C, and D), showing that IB-EDL achieves superior calibration performance under distribution shifts.

• Sensitivity to bin count in ECE: We conducted a sensitivity analysis on the number of bins for ECE calculation (Table F), showing that the rankings of methods remain consistent across different configurations.

If there are any remaining questions or concerns, please let us know.

Should you find that our responses sufficiently address your feedback, we would sincerely appreciate it if you consider raising the score.

Best regards,

Authors of IB-EDL

Q2: Is there a benefit to having two parameters estimating the distributional shift (mean and variance) of the token predictions vs. just one?

For implementation, we employ a single linear head with double the output neurons, instead of two separate linear modules. This design allows us to predict the mean and variance in a single forward pass. As discussed in Q1, both parts of the common linear head are initialized with the pre-trained LLM’s linear head.

Note that if we only estimate a mean parameter, we would have to assume an equal variance across different fine-tuning tasks, which would be a rather restrictive assumption.

Q3: What happens to calibration if we finetune on multiple datasets simultaneously?

This is an intriguing question. Incorporating multiple datasets during fine-tuning reduces the likelihood of the model overfitting to spurious correlations in a single dataset, thereby potentially improving calibration. While this perspective has not been extensively explored in prior work, we believe that enhancing calibration by increasing the diversity of training data represents an interesting and promising direction for future research.

Q4: How does this fine-tuning procedure effect LLM performance on open-ended generation tasks?

As discussed in response to W1, this question remains open and warrants further exploration. The answer will become clearer as the community develops well-established metrics for evaluating calibration and uncertainty in open-ended generation tasks. Testing IB-EDL under these scenarios is an exciting future direction.

Thank you again for your insightful comments. We will incorporate your suggestions in a revised version of the manuscript. We warmly invite you to engage in further interactive discussions or share additional suggestions.

[1] Kuhn, Lorenz, Yarin Gal, and Sebastian Farquhar. "Semantic Uncertainty: Linguistic Invariances for Uncertainty Estimation in Natural Language Generation." The Eleventh International Conference on Learning Representations.

We thank the reviewer for the constructive comments and thoughtful suggestions. We address the concerns as follows:

W1: 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.

Theoretically, our method is applicable to both multi-choice tasks and open-ended generation. This is because fine-tuning the model for open-ended generation involves predicting tokens within the vocabulary, corresponding to a case where the number of classes equals the vocabulary size. This scenario is already captured by our theoretical framework, so there is no theoretical barrier here. However, as we discussed in Section 6, we did not test our method on generative tasks due to the lack of well-established evaluations for model calibration in generative tasks. This remains an active research topic [1]. Testing IB-EDL on open-ended generation is indeed an exciting direction for future work. In this paper, addressing the overconfidence of LLMs in multi-choice tasks represents a significant step towards reliable LLM applications.

W2: 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).

To address this concern, we evaluate the performance on the MMLU-Math [3] dataset (covering topics such as “college_mathematics”, “high_school_mathematics”, and “abstract_algebra”) as the OOD test set. We use OBQA as the in-distribution dataset, consistent with Table 3 of the paper. Additionally, we employ Llama2-7B for this experiment.

As shown in Table E, we observe that: (i) IB-EDL consistently achieves the best OOD detection performance. (ii) In general, the AUROCs of our method and other baselines improve as the distribution shift increases.

Table E: OOD Detection AUROC on OBQA -> MMLU-Math

Q1: 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?

FLOPs are typically measured at inference time. As shown in Section 4.5, IB-EDL introduces only 1.98% additional FLOPs compared to the pre-trained LLM. To further address the reviewer’s question, we measure the number of samples processed per second as an indicator of computational complexity per step. Using Llama3-8B on a single NVIDIA H100 GPU, Table A shows that the training and inference speeds of IB-EDL are comparable to those of MAP and other EDL methods.

The mean and covariance prediction heads are initialized using the pre-trained LLM’s linear head and equipped with LoRA adapters. Both are fine-tuned jointly with the transformer layers (i.e., backbone). Thank you for raising this point, and we will include these details in the revised manuscript.

Table A: Additional Complexity Analysis

Dear Reviewer qw6T,

Thank you for your thoughtful review and constructive feedback on our manuscript. We greatly appreciate your recognition of the theoretical motivation and the strong experimental results.

As the discussion period will conclude in approximately 72 hours, we would like to kindly check whether our responses thus far have successfully addressed your concerns and clarified your questions.

In summary, we addressed the key points raised by you as follows:

• OOD detection on semantically distant datasets: We added experiments evaluating OOD detection performance on the MMLU-Math dataset (see Table E), demonstrating that IB-EDL achieves superior OOD performance, even with increased semantic distance.

• Training FLOPs and head design: We detailed in Table A that IB-EDL introduces minimal additional computational overhead and that the mean and variance are predicted via a single linear head in one forward pass, ensuring efficiency.

If there are any additional questions or remaining concerns, we would be happy to address them.

Should you find that our responses sufficiently address your feedback, we would sincerely appreciate your consideration of raising the score, if deemed appropriate.

Thank you again for your valuable input and for helping us improve our work.

Best regards,

Authors of IB-EDL

Thank you for the detailed response and additional experiments. Overall I think the reviewer's initial response addresses most of my concerns and I am willing to raise my score from marginal accept to accept.

• The detailed distribution shift experiments (including those in response to reviewer Tujz) help to demonstrate the benefits of this method over previous baselines and satisfy my concerns about the dependence on 'semantic distance' between sets.
• Thank you for the clarification on training of the mean/covariance heads and additional results for wall-clock time of training.
• While I agree that this method can possibly be extended to a discrete set of token sequences, the number of token sequences in open-ended generation is infinite and by definition the dirichlet prior is over discrete distributions. Additionally, it seems that y is assumed to be a one-hot vector, while in reality multiple semantically equivalent token sequences may be correct for a single question. I would like to see a bit more nuanced discussion on this extension but agree that experimentally it can be left to future work.

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

Our work is not a simple combination of these methods; instead, it tackles key unresolved issues in EDL (see also [3, 4]) and addresses the challenge when regularizing EDL with IB. Importantly, implementing IB on EDL in a naïve manner, such as inserting IB at an arbitrary layer, can compromise theoretical soundness. Our contributions go beyond mere implementation and address these challenges through the following advancements:

1. New theoretical insight: We uncover a critical issue overlooked by previous variational IB literature [1, 2]: applying IB at an arbitrary internal LLM layer can result in an optimization objective that is no longer a valid lower bound of $I(Z, Y)$. In Section 3.1, we provide a detailed theoretical analysis through Eq. (11).

2. Novel solution: We propose a unique and novel solution to the identified challenge. We show that, in EDL, this challenge can be addressed by selecting the pre-evidence $\tilde{e}$ as the IB variable. This approach ensures that the optimization objective remains a valid variational bound. Furthermore, as Eq. (14) illustrates, our IB method introduces L2 regularization on the evidence, directly preventing the model from producing extremely large evidence values and thereby reducing overconfidence. This solution directly addresses the drawbacks of EDL identified in prior works [3, 4].

3. Unified perspective: Through Proposition 1, we unify previous approaches, providing a cohesive theoretical framework. This unified viewpoint is also novel and provides valuable insights into the family of EDL methods.

In addition to these theoretical contributions, we also offer practical advancements. For instance, previous EDL works only validated their methods on ResNet-sized models. In contrast, we demonstrate that our EDL method scales to billion-parameter-sized LLMs. This advancement opens new avenues for extending EDL methods to broader research and application domains.

Thank you again for the insightful comments and we will incorporate your suggestions in a revised version of the manuscript. We warmly invite you to engage in further interactive discussions to explore any additional perspectives you might have.

[1] Alemi, Alexander A., et al. "Deep variational information bottleneck." arXiv preprint arXiv:1612.00410 (2016).

[2] Wieczorek, Aleksander, and Volker Roth. "On the difference between the information bottleneck and the deep information bottleneck." Entropy 22.2 (2020): 131.

[3] Deng, Danruo, et al. "Uncertainty estimation by fisher information-based evidential deep learning." International Conference on Machine Learning. PMLR, 2023.

[4] Chen, Mengyuan, Junyu Gao, and Changsheng Xu. "R-EDL: Relaxing Nonessential Settings of Evidential Deep Learning." The Twelfth International Conference on Learning Representations.

W3: 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.

We now present additional calibration experiments in Table B, C, and D. These tables show that IB-EDL achieves the best calibration performance on two out of three OOD datasets. This result demonstrates that IB-EDL’s calibration capabilities generalize well under distribution shifts.

Table B: Calibration in the setting OBQA -> ARC-C

Table C: Calibration in the setting OBQA -> ARC-E

Table D: Calibration in the setting OBQA -> CSQA

We thank the reviewer for the constructive comments and feedback. We address the concerns as follows:

W1: The pipeline of IB-EDL ( $x \rightarrow f(x; \theta) \rightarrow \tilde{e} \rightarrow e \rightarrow \alpha \rightarrow \pi \rightarrow y$) seems computationally heavy, it would be informative to include an inference time comparison of IB-EDL against other methods.

Compared to other baselines such as MAP, the additional complexity of IB-EDL lies in the augmented final linear layer. However, the main computational cost of the model is induced by the transformer layers rather than the linear head. Furthermore, after the forward pass through the transformer layers and the final linear layer, the subsequent steps in the pipeline ($\tilde{e} \rightarrow e \rightarrow \alpha \rightarrow \pi \rightarrow y$) are basic algebraic operations and do not involve any model parameters.

In Section 4.5, we present a complexity analysis showing that IB-EDL adds less than 2% extra FLOPs to the pre-trained model. In addition, we now include more detailed training/test time and memory analyses. Specifically, we conducted experiments with Llama3-8B on the OBQA dataset using a single NVIDIA H100 GPU. As shown in Table A, IB-EDL’s training and inference speeds are comparable to those of the MAP and EDL baselines and are significantly faster than MCD or Ens. Additionally, the extra memory cost of IB-EDL is also marginal.

Table A: Additional Complexity Analysis

W2: 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.

To address this concern, it is sufficient to compare IB-EDL and MAP in Table A. This is because the MAP setting corresponds exactly to LoRA fine-tuning of the pre-trained LLM. As Table A demonstrates, IB-EDL does not show a significant slowdown in training time compared to MAP.

Dear Reviewer tujz,

Thank you for your thorough review and constructive feedback on our manuscript. We greatly appreciate your recognition of the strengths of our work, including the strong experimental results, the clear presentation, and the unified theoretical perspective that IB-EDL provides for EDL research.

As the discussion period is set to conclude in approximately 72 hours, we would like to kindly check whether our responses thus far have successfully addressed your concerns and clarified your questions.

In summary, we addressed the points raised by you as follows:

• Inference and training time comparisons: We provided detailed complexity analyses (Section 4.5 in paper and Table A here), showing that IB-EDL introduces minimal computational overhead during training and inference, with performance comparable to MAP and EDL baselines.

• Comparison with LoRA training time: We clarified that the MAP setting corresponds directly to LoRA fine-tuning, and our experiments confirm that IB-EDL incurs negligible additional overhead compared to LoRA.

• Calibration on OOD datasets: We conducted additional experiments to evaluate ECE on OOD datasets. As shown in Tables B, C, and D, IB-EDL achieves strong calibration performance under distribution shifts.

• Novelty of IB-EDL: We highlighted that IB-EDL is not a simple combination of existing methods but addresses critical theoretical challenges in applying IB to EDL, as detailed in Section 3.1 and Section 3.2. We also provided a unified framework for EDL methods and demonstrated scalability to large LLMs, which represents a significant advancement over prior works.

If there are any additional questions or concerns, we would be happy to address them.

Should you find that our responses sufficiently address your feedback, we would sincerely appreciate your consideration of raising the score, if deemed appropriate.

Best regards,

Authors of IB-EDL

Thanks for addressing my concerns, I have raised the score accordingly.