Functional-level Uncertainty Quantification for Calibrated Fine-tuning on LLMs

Additionally, beyond the lack of comprehensiveness, the paper offers very little empirical or theoretical insight into why this approach is effective. I also find that the paper lacks depth overall.

For empirical insights, we evaluate UQ4CT and the baselines with a diverse set of tasks, covering both common-sense reasoning (i.e., OBQA, BOOLQ) and domain-specific tasks(i.e., ClimateQA). The strong empirical results across tasks showing over 25% reductions in Expected Calibration Error (ECE) compared to baselines as shown in Tables 1 and 2 underscores the generalizability of UQ4CT and validates its applicability across different domains and contexts.

These results collectively provide robust empirical insights into the effectiveness and practicality of UQ4CT. Additionally, we have added experimental results with Mistral-7B and multiple ablation studies to enhance comprehensiveness.

We have demonstrated theoretical insights in Section 3.1, 3.3 and also Figure 1 in our paper. The calibration loss effectively promotes expert exploitation for correct predictions and expert exploration for incorrect predictions. Ideally, when the $K$ LoRA experts together have captured all the functional relationships across the data distribution after fine-tuning (via cross entropy loss), our proposed loss \mathcal{L}_\{cal} calibrates the mixture of LoRA experts conditioned on the input $x$ across the training set. This ensures that the model effectively selects the appropriate experts to produce outputs aligned with the underlying data distribution corresponding to $x$.

Incorporating the revision suggestions and additional experiments contributes to the depth of our work, as the new evaluations with Mistral-7B and comprehensive ablation studies ensure a robust foundation for our claims.

In equation (9), what is the term $\mathcal{L}_b$? I don't think this is a good practice to put it in the reference. It should be at least written in the appendix.

We appreciate the suggestion. We have included the definition of $\mathcal{L}_b$ in the appendix in the revised version. The term $\mathcal{L}\_b$ is the load balancing term proposed in [1] to address the phenomenon that some experts tend to be chosen more frequently by the top-k routers during fine-tuning.

Given $N$ experts indexed by $i=1$ to $N$ and a batch $B$ with $T$ tokens, the auxiliary loss is computed as: $\mathcal{L}\_{aux}=a \cdot N \cdot \sum^{N}\_{i=1}\mathcal{F}\_i \cdot \mathcal{P}\_i$,

where:

$\mathcal{F}\_i = \frac{1}{T}\sum_{x\in B} \mathbb{1}\\{argmax_k \mathcal{R}(x)_k=i\\}, \mathcal{P}_i = \frac{1}{T}\sum\_{x\in B}\mathcal{R}(x)_i$.

Here, $\mathcal{R}(\cdot)$ is the top-k router, $\mathcal{F}_i$ is the fraction of tokens dispatched to expert $i$ and $\mathcal{P}_i$ is the fraction of the router probability allocated for expert $i$. The final loss is multiplied by the expert count $N$ to keep the loss constant as the number of experts varies, and the constant term $a$ is set to $10^{-2}$as a multiplicative coefficient, which is large enough to ensure load balancing while remaining small enough not to overwhelm the primary objective.

In table 1 and 2, what is the numeric values in the subscript?

We apologize for the confusion. We report the mean and standard deviations across 3 random runs for each experiment. The subscript corresponds to the standard deviations. We have added explanations in the revised version.

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

We thank the reviewer for the insightful questions and helpful suggestions for experimental comprehensiveness. We address the concerns as follows:

The biggest weakness of this paper is the limited experimental effort. Specifically, the experiments are conducted on only one LLM, and there is very little ablation or other study to investigate the additional properties of this UQ method. While I don’t consider the number of pages a deterministic measure of effort, the brevity of the experimental section does suggest a lack of thoroughness in the design of experiments.

We appreciate the suggestions on empirical evaluations and have added multiple experiments and ablation studies to assess our method (UQ4CT) more comprehensively. For all experiments, we evaluate performance with 3 random runs and report the mean and standard deviations.

Baseline Comparisons with Mistral-7B. The experimental results demonstrates that our method is applicable to different LLMs, outperforming existing methods in uncertainty calibration while maintaining accuracy:

Ablation Study: # of Experts to Select. We conduct an ablation study on the number of experts to select on OBQA dataset with Llama2-7B. This explores how the number of active experts can impact the functional relationships utilized by the model for inference:

An theoretical insight of using the top 2 experts out of 8 in total is that one expert alone cannot capture complicated functional relationships, while more than 2 experts introduce redundant functional bases to the model, which deviates the output distribution more from data distribution, thus worsening predictive and calibration performance.

Ablation Study: Varying Loss Coefficients. We also conduct an ablation study on varying loss terms where we fix alpha (coefficient of the load balancing term) to 1 and compare model performance with different beta values (0, 0.2, 0.5, 0.8, 1) on OBQA dataset with Llama2-7B:

The results demonstrate the effectiveness of our proposed calibration loss. When $\beta = 0$, the model is optimized without calibration on parameter mixtures, resulting in high ECE value. Even with small $\beta=0.2$ or $\beta=0.5$, the ECE scores drastically improved compared to no calibration setting. Finally, when $\beta=1$, the calibration term effectively optimizes the conditional parameter mixtures to generate outputs that fit data distribution well, resulting in lower ECE scores and higher accuracies.

Thanks for the response. My concerns are partially addressed and I will raise the score from 3 to 5.

We appreciate the constructive feedback and thoughtful suggestions. We address the concerns as follows:

The ablation study/studies could be more elaborated, i.e. e.g. examining what the influence of the different terms in the loss is, or how sensitive performance is to the hyperparameters alpha and beta in the loss. I think such examinations are important to understand what’s going on and how brittle such triple losses are.

This is a great suggestion, we have added an ablation study on varying loss terms where we fix alpha (coefficient of the load balancing term) to 1 and compare model performance with different beta values (0, 0.2, 0.5, 0.8, 1) on OBQA dataset with Llama2-7B. Results are given as following, we report the mean and standard deviation of the performance under 3 random runs:

The results demonstrate the effectiveness of our proposed calibration loss. When $\beta = 0$, the model is optimized without calibration on parameter mixtures, resulting in high ECE value. Even with small $\beta=0.2$ or $\beta=0.5$, the ECE scores drastically improved compared to no calibration setting. Finally, when $\beta=1$, the calibration term effectively optimizes the conditional parameter mixtures to generate outputs that fit data distribution well, resulting in lower ECE scores and higher accuracies. We will add this ablation study in the revised version.

The approach is specifically tailored to the nature of the LoRA MoE architecture as it critically depends on leveraging the weights of the Top-K Routing layers of the MoE. While this seems to work pretty well here, one might also consider this a limitation of the approach in general since it is not generally applicable.

Although our approach depends on the LoRA MoE architecture for LLMs, the core ideas of functional-level epistemic uncertainty quantification and calibration can indeed be generalized and applied to broader settings beyond LoRA. For example, [1] explores the MoE architecture with various deep learning networks such as CNN/ResNet, demonstrating the broad applicability of MoE architecture. The MoE mechanism leverages router distribution to dynamically assign inputs to different experts based on the hidden states. This is not exclusive to LoRA and can be applied to other neural networks, including non-transformer architectures.

In Equation (5): What exactly do M’ and a’ refer to? I believe this is not explicitly specified anywhere

We appreciate the feedback. We have added demonstrations in the paper accordingly. M’ refers to the optimal mixture of experts with model M, and a’ refers to the ground truth answer generated by model M’.

Why did you place the “Related work” section after your methodology? I think it introduces important concepts to the reader who would thus benefit from reading this before your explanations in the “Methodology” section.

We presented the 'Preliminaries' section prior to the 'Method' section, demonstrating the necessary and context-specific background knowledge required to understand our proposed method. After 'Method' section, we then present the 'Related Work' section for a broader overview of recent advancements in relevant fields, offering insights that inform the selection of experimental baselines.

L. 326: I personally think it’s not optimal to call something “significant” that you actually did not formally test with a statistical test. Maybe it’s better to use a formulation like “notable” or “substantial”?

L 359: Just to be sure, what do you mean by “bin size of 15”? I think it is more common to communicate this in terms of “number of bins (of equal size usually”?

We apologize for the ambiguity. We have revised the wording in the paper. Here it should be “number of bins” instead of “bin size”.

L. 447 More of a comment than a question, but I think a sole subsubsection (5.4.1) does not really make sense, probably you should rather just add this subsubheading to the heading of subsection 5.4.

Thanks for the suggestion, it is reasonable to do so with only one ablation study. Since we have added multiple ablation studies now, we are allocating them into multiple subsections.

[1] Chen, Zixiang, et al. "Towards Understanding the Mixture-of-Experts Layer in Deep Learning."

There are a lot of hyper-parameters, such as the number of top gate routers (why use 2 in this paper), and N, s, L. How to set them in practice and is there any theoretical insight for these choices?

Generally, we follow the hyperparameter setup in previous works regarding cutoff length and total number of LoRA experts. For the optimizer hyperparameters, we perform a hyperparameter search and use the optimal value found. For the number of top gate routers, we performed an exploration from 1 to 5 experts and found that 2 is the optimal. The performance comparison on the OBQA dataset with Llama2-7B is given as following, we report the mean and standard deviation of the performance under 3 random runs:

We will also add this as an ablation study in the revised version. An theoretical insight of using the top 2 experts out of 8 in total is that one expert alone cannot capture complicated functional relationships, while more than 2 experts introduce redundant functional bases to the model, which deviates the output distribution more from data distribution, thus worsening predictive and calibration performance.

We appreciate the constructive feedback and thoughtful suggestions and address the concerns as follows:

The novel calibration loss function is mentioned as one of the contributions. However, the paper lacks detailed theoretical analysis to illustrate why this specific design especially the loss is effective for uncertainty calibration.

Can the authors add some theoretical analysis?

We demonstrate the theoretical insight of our proposed method as following:

When the $K$ LoRA experts together have captured all the functional relationships across the data distribution after fine-tuning (via cross entropy loss), and the training and testing data share the same distribution, our proposed loss \mathcal{L}_\{cal} calibrates the mixture of LoRA experts conditioned on the input $x$ across the training set. This ensures that the model effectively selects the appropriate experts to produce outputs aligned with the underlying data distribution corresponding to $x$.

We will add detailed proof in the revised version. Empirically, our results show that UQ4CT effectively aligns the model's predictive confidence with the underlying data distribution by selecting the appropriate functional bases conditioned on the input $x$ after calibration.

Although the paper compares UQ4CT with other PEFT-based uncertainty quantification techniques, a more comprehensive comparison with non-PEFT-based methods might provide a clearer view of its advantages and limitations.

Thanks for the suggestion, we have explored recent non-PEFT-based uncertainty quantification methods on LLMs while none of them applies to our setup. Current non-PEFT-based uncertainty quantification methods focus on open-ended questions while current PEFT-based methods are limited to multiple-choice questions. Out of three recent works on non-PEFT-based methods [1,2,3], only [1] is applicable to multiple-choice questions with the question-paraphrasing approach. [2] requires multiple high-quality clarification phrases for each prompt, while [3] relies on multiple samples given the same question, which does not fit into the multiple-choice scenario.

Although [1] is applicable, they conducted the experiment with models that are orders of magnitude bigger than our setup (ChatGPT3.5/ChatGPT4, etc). When extending their method to smaller models without fine-tuning, their approach of perturbing the questions and the temperature significantly worsens the predictive accuracy, thus making the calibration measurements meaningless. We provide the results with [1] in the following table just for demonstration purposes. The results are obtained with Llama2-7b-hf:

The whole framework is built on LoRA, and can it be applied to more general settings?

Absolutely, while the framework described in the paper is built on the LoRA Mixture-of-experts (MoE) architecture, the core ideas of functional-level epistemic uncertainty quantification and calibration can indeed be generalized and applied to broader settings beyond LoRA. For example, [4] explores the MoE architecture with various deep learning networks such as CNN/ResNet, demonstrating the broad applicability of MoE architecture. The MoE mechanism leverages router distribution to dynamically assign inputs to different experts based on the hidden states. This is not exclusive to LoRA and can be applied to other neural networks, including non-transformer architectures.

[1] Gao, Xiang, et al. "SPUQ: Perturbation-Based Uncertainty Quantification for Large Language Models."

[2] Hou, Bairu, et al. "Decomposing Uncertainty for Large Language Models through Input Clarification Ensembling."

[3] Kuhn, Lorenz, et al. "Semantic Uncertainty: Linguistic Invariances for Uncertainty Estimation in Natural Language Generation."

[4] Chen, Zixiang, et al. "Towards Understanding the Mixture-of-Experts Layer in Deep Learning."

Thanks for answering my questions but I still think that some theoretical analysis is necessary

We appreciate the constructive feedback and thoughtful suggestions. We address the concerns as follows:

The use of MoE for dynamic expert selection, while effective, could lead to scalability issues with larger models or multi-task scenarios due to increased computational demands. Optimizing expert selection with pruning or adaptive sparsity methods could help make UQ4CT more efficient and feasible for resource-limited applications.

Optimizing with pruning would definitely help with reducing computational resources needed. However, the LoRA Mixture of Experts (MoE) architecture is highly unlikely to experience scalability issues because of its parameter-efficient nature. As reported in [2], the MoE architecture only introduces 2.9% additional parameters regarding the Llama2-7B model, which has 7 billion parameters. For multi-task scenarios, some existing applications [3,4] with similar LoRA MoE architecture have shown promising results without scalability issues. This is because the number of experts is independent from different problem difficulties and, according to the empirical results in [3,4], does not scale up even for multi-task scenarios.

Computational wise, LoRA MoE only inherently activates a small subset of experts for each input (i.e. 2 out of 8 LoRA experts in our setup), which reduces the computational overhead compared to dense approaches. This sparsity ensures that even with a large number of experts, the computational cost scales more efficiently than with a dense approach of equivalent size.

Since UQ4CT’s calibration relies on binary correctness, how might this approach be adapted for tasks where “correctness” is less definitive, such as open-ended text generation or summarization?

For open-ended text generation / summarization problems, there indeed lacks a definitive optimization objective. We can adapt UQ4CT to open-ended text generation or summarization tasks with non-definitive signals, such as human preference feedback for RLHF tasks or semantic uncertainty [1], which defines and utilizes semantic entropy in detecting hallucinations in open-ended question answering tasks.

Could the authors elaborate on any experiments or initial explorations in applying UQ4CT to such tasks?

Exploring approaches that rely on alternative correctness signals, such as soft scoring or feedback-based calibration, might broaden the applicability.

This is a great point. We are actively exploring methods to apply UQ4CT for open-ended question-answering datasets. The semantic entropy method in [1] is a suitable alternative to the predictive correctness for open-ended tasks. Our current approach is to compute the semantic entropy following [1] and use it as the target score to calibrate the functional level uncertainty. Considering the rebuttal time limitation and implementation complexity, we would not be able to show the performance of this approach.

[1] Kuhn, Lorenz, et al. "Semantic Uncertainty: Linguistic Invariances for Uncertainty Estimation in Natural Language Generation."

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

[3] Liu, Qidong, et al. “When MOE Meets LLMs: Parameter Efficient Fine-tuning for Multi-task Medical Applications.”

[4] Feng, Wenfeng, et al. “Mixture-of-LoRAs: An Efficient Multitask Tuning for Large Language Models.”