Large Language Models (LLMs) often suffer from overconfidence during inference, particularly when adapted to downstream domain-specific tasks with limited data. Previous work addresses this issue by employing approximate Bayesian estimation after the LLMs are trained, enabling them to quantify uncertainty. However, such post-training approaches' performance is severely limited by the parameters learned during training. In this paper, we go beyond post-training Bayesianization and propose Bayesian Low-Rank Adaptation by Backpropagation (BLoB), an algorithm that continuously and jointly adjusts both the mean and covariance of LLM parameters throughout the whole fine-tuning process. Our empirical results verify the effectiveness of BLoB in terms of generalization and uncertainty estimation, when evaluated on both in-distribution and out-of-distribution data.
Paper
Similar papers
Peer review
Summary
The paper proposes an algorithm to learn a (low-rank) variational posterior distribution over (a subset of) LLM weights during fine-tuning by combining Low-Rank Adaptation (LoRA) and Bayes by Backprop (BBB). To make the algorithm work in practice, several nontrivial modifications are introduced, such as certain parameterisations and optimisation strategies. Finally, the proposed method is empirically evaluated on a variety of in-distribution and out-of-distribution benchmarks and compared to relevant baselines.
Strengths
- The paper is well written and the proposed algorithm is contextualised into existing work, with Figure 1 providing a great visual overview. - Concepts are introduced in an adequate pace and order, and a clean notation makes it easy for the reader to follow. - Elaborate empirical evaluation seems to suggest (somewhat) consistent gains over baseline approaches. - While the proposed algorithm seems like a straightforward combination of LoRA and BBB, the paper mentions that certain modifications are necessary for the algorithm to work in practice. These modifications are a core part of the paper's contribution, and they are discussed thoroughly and justified with a mix of theoretical and empirical arguments.
Weaknesses
- The proposed method requires additional compute time and memory compared to standard LoRA. While this information is provided in Appendix C.2 (and seems to be small), I encourage the authors to be transparent about this by briefly mentioning it in the main paper (if that is not already the case, perhaps I missed it).
Questions
- Appendix B.1 mentions weights for the KL divergence term. How important is this? Please correct me if I am wrong, but I believe this was not explicitly discussed in the main paper. If this is important for the overall performance of the algorithm (and typically these kind of adjustments tend to be...), I suggest to at least mention it in the main paper (even if details remain in the appendix). - Is there any particular reason why $\mathbf{A}$ is modelled in a probabilistic way instead of $\mathbf{B}$? In other words, would it also be possible to do the reverse, i.e. assume a prior over $\mathbf{B}$ and treat $\mathbf{A}$ deterministically? The choice seems arbitrary to me, but perhaps there is an underlying reason to it.
Rating
7
Confidence
3
Soundness
4
Presentation
4
Contribution
3
Limitations
Theoretical assumptions are stated as part of the provided theorems. As far as I am aware, general limitations of the proposed method (beyond the ones which are being addresses by the proposed method itself) are not explicitly discussed and could at least be briefly mentioned in the conclusion.
Summary
This paper proposes a method called Blob, which uses mean-field variational inference on LoRA parameters to obtain Bayesian estimation on LLM. This will result in a richer posterior structure on the weight than diagonal, while maintaining the same computational cost as diagonal covariance. During training, a new parameterization is used for posterior covariance to accelerate convergence and Flipout is used to make training more stable. Empirical results show Blob obtains competitive performance compared with other Bayesian methods.
Strengths
- By putting a diagonal prior and approximate posterior on LoRA parameters, it leads to a richer posterior covariance structure while maintaining the computational cost of diagonal covariance. - The new parameterization of standard deviation seems promising for all VI training. - Good performance compared with other Bayesian methods.
Weaknesses
- Not clear about the method's limitations. Since it uses VI, common disadvantages of VI also apply: (1) it needs to do sampling when making predictions, which can be computationally expensive given the size of LLM. (2) Given the complex posterior structure, there will be a very high chance of getting noise samples, which leads to suboptimal performance. This is already demonstrated in Table 1 that using MC estimation decreases the ACC compared with only using the mean. In additional, it seems in practice many extra efforts are needed to get the method to work, e.g., the complex KL re-weighting and rescale size of training data mentioned in Appendix B.1. - The whole paper feels a bit overselling: (1) I don't see the need to make Theorem 3.1 and Theorem 3.2 theorems, nothing surprising is stated there. (2) In Line 318, the authors claim their method "demonstrates superior generalization and uncertainty estimation capabilities", I don't think the results justify the "superior" claim. (3) In essence, Blob is mean-field VI on LoRA parameters trained with better strategies to make it suffer less from VI's high variance. I have no problem at explaining methods in great detail, but the whole writing on page 4 and page 5 feels like the authors are trying to make it sound as fancy as possible. I like the idea of the paper, but I think it can be presented in a more honest way.
Questions
Major: - I'm not sure what it mean for the mean matrix of q(A) to be an output of a neural network (Line 186-187). Does this mean you are using another neural network to fit the mean of posterior? - In Eq. (70), it seems different mini-batch has different KL re-weighting hyperparamers lambda? How sensitive is the training to the value of lambda? - What is the difference between BBB and Blob? From my understanding, the only difference is how you do the training. For BBB you use original parameterization and for Blob you use new parameterization plus Flipout? Also how many samples did you use to evaluate BBB in Table 1? - I would be interested to see a figure plot of model performance versus the number of MC samples. Specifically, the y-axis is accuracy/ECE/NLL, the x-axis is the number of samples used during inference. This helps the reader to understand the sufficient sampling size for inference. - I would be interested to see Laplace added to Table 8 to show a fair comparison of the computation cost compared with Blob. - Why is Laplace performing so poorly on ARC-C dataset? The original Bayesian LoRA paper it seems to perform much better on that dataset. Minor: - Line 236, MLE and MAP are not uncertainty estimation methods. - I don't entirely agree with the claim made in Lines 46-47. The possible suboptimal estimation of Laplace approximation comes from (1) LA, which assumes we reach MAP estimation while in practice you never know if the complex NN has really converged to MAP; (2) setting prior precision. It doesn't come from the post-training procedures. - the notation "n" refers to two different things, at Eq. (4) it's the number of parameters, and at Line 153 it's the dimension of the weight matrix. - Bolding is missing on ACC in Table 7.
Rating
5
Confidence
4
Soundness
2
Presentation
2
Contribution
2
Limitations
See weaknesses and questions above.
Remaining Questions for Reviewer Adn3
**W1. Not clear about the method's limitations.** Please refer to **Q1. [Discussion of Limitations]** in **General Response**. **Q1. What it mean for the mean matrix of q(A) to be an output of a neural network (Line 186-187)?** We are sorry for the typo. $q(A)$ is not the output of a neural network. It is directly modeled as a Gaussian distribution $\mathcal{N}(M, \Omega)$, with $M$ and $\Omega$ as learnable parameters. **Q2. In Eq. (70), it seems different mini-batch has different KL re-weighting hyperparamers lambda? How sensitive is the training to the value of lambda?** Yes, different mini-batches has different $\lambda$ values. When we use the typical settings of $\lambda_i = \frac{2^{M-i}}{2^M-1}$ [1] or $\lambda_i = \frac{1}{M}$ [2], the model struggles to fit the data distribution while converging to the prior distribution, resulting in an "NaN" NLL loss. Our setting of $\lambda_i = \frac{2^i}{2^M-1}$ allows the model to fit to the data in the early stages of training and achieve a reasonable level of complexity cost in the later stages. This re-balance scheme of data likelihood and complexity cost has proven effective across multiple datasets and can be potentially applied to other VI methods. We will include the discussion above in the revision as suggested. [1] Blundell, Charles, et al. "Weight uncertainty in neural network." International conference on machine learning. PMLR 2015. [2] Graves, Alex. "Practical variational inference for neural networks." NeurIPS 2011. **Q4. I would be interested to see a figure plot of model performance versus the number of MC samples. ...** This is a good suggestion. Following your suggestion, we have run additional experiments for model performance versus the number of MC samples on WG-S dataset. The results are presented in the table below. The corresponding figure is provided as **Figure 1** in the PDF attached to the **General Response**. | | **N=0** | **N=1** | **N=2** | **N=3** | **N=4** | **N=5** | **N=10** | **N=20** | **N=40** | **N=80** | **N=160** | | ---------------------- | ------- | ------- | ------- | ------- | ------- | ------- | -------- | -------- | -------- | -------- | ------ | | **ACC ($\uparrow$)** | 71.44 | 65.20 | 66.28 | 66.91 | 67.63 | 68.20 | 68.31 | 68.04 | 68.15 | 68.31 | 68.15 | | **ECE ($\downarrow$)** | 19.71 | 19.72 | 14.60 | 12.82 | 11.27 | 10.58 | 9.51 | 9.47 | 9.25 | 8.75 | 8.75 | | **NLL ($\downarrow$)** | 0.84 | 0.8617 | 0.7355 | 0.6971 | 0.6766 | 0.6620 | 0.6395 | 0.6313 | 0.6233 | 0.6213 | 0.6215 | We can see that in general, the performance improves with more MC samples, but saturates after around $N=10$. **Q6. Why is Laplace performing so poorly on ARC-C dataset? The original Bayesian LoRA paper it seems to perform much better on that dataset.** Please refer to **Q2. [Reproducing LAP's Results on the ARC-C Dataset]** in **General Response**. **Q7. Line 236, MLE and MAP are not uncertainty estimation methods.** We apologize for the confusion. We will remove "uncertainty estimation" when describing MLE and MAP in the revision as suggested. **Q8. I don't entirely agree with the claim made in Lines 46-47. The possible suboptimal estimation of Laplace... never know if the complex NN has really converged to MAP; (2) setting prior precision...** This is a good point. We agree that these are also major reasons that lead to Laplace approximation's poor performance. We will include these two additional reasons in our revision as suggested. **Q9. Notation issue and missing Bolding** We are sorry for these typos and will correct them in the revision as suggested.
Thank you for the rebuttal, it addresses most of my concerns. Still, I am not fully convinced that a Theorem is needed to state results in Theorem 3.1 and Theorem 3.2, maybe this is because I come from a more Bayesian background and thus the results do not look surprising. On a side note, I recently came across an ICML 2024 paper "Variational Learning is Effective for Large Deep Networks". In essence, they do natural gradient descent in natural parameter space for mean-field VI, and they make it work for large-scale neural networks including GPT2. This is highly relevant to the paper so I strongly suggest the author to discuss it in the paper.
Thank You for Your Further Feedback
Thank you for your continued feedback and for keeping the communication channel open. We are glad that you found our response helpful and that it addressed most of your concerns. Below we address your remaining comments in detail. **W3.1 [the Term "Theorem"]** We appreciate your insightful feedback. In light of your suggestion, we will revise Theorem 3.1 and Theorem 3.2 to Proposition 3.1 and Proposition 3.2, respectively. We believe this change can more accurately reflect the nature of our contributions. Nonetheless, we maintain that the discussion presented here offers valuable insights to the broader community. The theoretical foundation of Asymmetric Bayesianization laid out in this work can serve as a springboard for the development of new methods in the field. In addition, inspired by your comments, we plan to reorganize Section 3.1 to enhance clarity. We will start with the calculation of the full weight matrix $W\_{ij}=W\_{0,ij} + \sum\_{k=1}^r B\_{ik}A\_{kj}$ (i.e., Eqn. 5 of the paper) and subsequently introduce the advantages of Asymmetric Bayesianization over bayesianizing both $A$ and $B$. **[Related Work]** Thank you for bringing the ICML 2024 paper "Variational Learning is Effective for Large Deep Networks" to our attention. It was not available at the time of our submission. We agree that it is highly relevant to our work, and we will ensure it is cited and discussed in the revised version. Finally, we would like to express our sincere gratitude once again for your insightful and constructive comments, which have greatly helped improve our paper.
Summary
Proposes to have a Bayesian version of LoRA by placing priors over the low rank matrices. This evolves into placing prior over one of the low rank matrices. Inference is via variational inference. Results are mixed, though it shows the method is competitive with others.
Strengths
- The eventual simplicity of the approach by placing prior over only one of the low rank matrices. - Use of independent Gaussians in prior and posterior for efficiency. - The method is competitive, based on the results.
Weaknesses
The flow of the paper is rather poor, and the presentation of the method overly complicated. I prefer the authors to simply start from (5) and then say the deltas to the pre-trained weights are Gaussians because they are sum of Gaussians. The consideration to also put prior on $B$ can be deferred to discussions, and the trade-offs discussed there. This make the model presentation cleaner, and avoids the problem of having to involve improper priors in Theorem 3.1 (since B is low rank, $\Sigma_{q}$ is low rank, the Gaussian is degenerate) and Theorem 3.2, and the subsequent problem of reader having to think about how the KL with generate Gaussians would work out in theory. In short, the theorems are unnecessary and need fixing because of improper distributions. 1. Line 137: What is it that requires accurate estimation? 2. Section 4: How is more than one sample (for N>0) used in generating a single output for evaluation and then influence the evaluation metric? *Minor* - Eq (3) to remove the $\min_{\theta}$ prefix, and to introduce the prior and likelihood. - Section 4: Shouldn't ECE and NLL be the metrics in focus because this is about "Bayesian"? - Section 4: The numbers in the table does not justify saying BLoB is the best. - "Bayesianization" is a mouthful --- can we think of a better derived word?
Questions
Q1 and Q2 above (in weaknesses)
Rating
6
Confidence
5
Soundness
4
Presentation
1
Contribution
3
Limitations
No. Paper need to include some discussion on the limitations of their approach/paper. I do not see any potential negative societal impact to be addressed.
Degenerate Gaussians are Proper for Probabilistic Inference
In this separate Official Comment, we will elaborate on two key reasons why degenerate Gaussian distributions are valid for probabilistic inference: 1. **Their probability density is well-defined.** According to Schoeman et al. [1], the *probability density* of a degenerate Gaussian distribution can be factorized into the product of the density of a non-degenerate Gaussian distribution in a lower-dimensional linear subspace and a Dirac delta function (Eq. 12 in [1]). This density function has an alternative expression in the form of limit (Eq. 16-18 in [1]), which aligns with our main theorem 3.1 and 3.2: $$ p(x) = \lim_{a\rightarrow 0}\mathcal{N}(x | \mu, Q(\Lambda^{-1} - aI)Q^T + a I) = \mathcal{N}(x | \mu, \Sigma=\lim_{a\rightarrow 0} Q(\Lambda^{-1} - aI)Q^T + a I). $$ Here $\Lambda^{-1}$ is the precision matrix of the lower-dimensional non-degenerate Gaussian distribution, and $Q$ is the linear transformation matrix that expands this to high dimensions. In our main paper, we use the simplified notation $\mathcal{N}(x | \mu, \Sigma)$ for readability, omitting the limit. The full limit notation and proof appear in Appendix A.1. We will clarify this simplification in our revision. 2. **The KL divergence between two degenerate Gaussians is computable under certain conditions.** Secondly, the *KL divergence* between two degenerate Gaussians can be computed under the condition that they "have support on the same lower-dimensional manifold" (Eq. 44 in Schoeman et al. [1]). This condition has also been clearly stated as satisfied in our Theorem 3.2 ($RR^{T} = BB^{T}$, line 177), and we also include a detailed derivation of how we reach this condition in Appendix A.1 (line 676-678). [1] Schoeman, Johannes Cornelius, Corné E. van Daalen, and Johan A. du Preez. "Degenerate Gaussian factors for probabilistic inference." IJAR 2022.
W1 + W2. I understand that the method is sound because your degenerate Gaussians are on the same subspace. It is now clear to me that you also understand this constraint. I'll increase the soundness score, and the overall score. However, I still believe that the presentation can be vastly simplified (and hence the paper more accessible) by simply starting from (5). The relation to full Bayesianization can be deferred to discussion. W5. It is "generally better" than the other methods. You can also say it is "the best on average" if you do some averaging either over the score or by counting wins. You cannot just say it is "the best". Q2. I see that the "experimental setup strictly adheres to the original LoRA [2] framework". It is good to say this explicitly in the paper. For ACC, do you generate, say N=10, and then count if *any* is correct? Is there some discount because you have N chances?
Thank You for Your Further Feedback
Thank you for your further feedback and for keeping the communication channel open. We are glad that you found our response helpful and that it addressed your concerns. Below we address your remaining comments in detail. **W1+W2 [Presentation]**: Thank you again for your suggestion regarding the presentation. We will re-organize the presentation to start with Eqn. 5 in the revision as you suggested. **W5 ["Generally Better"]**: Thank you for your suggestion. We will adjust the claim from "the best" to "generally better" to more accurately reflect our contribution. **Q2 [ACC and N=10 Samples]**: Thank you for your recognition and insightful question. For ACC, we first draw N=10 samples, i.e., N=10 output probabilities, from our BLoB, and then compute the average of these N samples as *one* single final prediction. This approach aligns with the method we previously described: "approximate this expectation through sampling: $E_{q(W|\theta)}[P(Y|W,X)] \approx \frac{1}{N} \sum_{n=1}^{N} P(Y|W_n, X)$." Based on this averaged output probability (i.e., the approximation of the expectation), we then calculate the ACC. The ACC is *not* calculated based on whether *any* individual sample is correct. Instead, N=10 samples are aggregated to make *only one prediction*. Therefore there is *no discount* needed. We will follow your suggestion to provide further clarification on the output of variational inference and clarify that "experimental setup strictly adheres to the original LoRA [2] framework" in our revised version. Last but not least, we would like to thank you again for your insightful and constructive comments. They have greatly help improve our paper.
Improved Outline of the Method
Thank you for your valuable suggestions on enhancing the presentation quality of our work. Following your suggestions, we have refined Sections 3.1 and 3.2. Below is the outline of our updated version: **Section 3.1: Low-Rank Variational Approximate Posterior Distribution** 1. Main Method: Introduction of Asymmetric Bayesianization Scheme - Calculating each entry of the full weight matrix: $W\_{ij}=W\_{0,ij} + \sum\_{k=1}^r B\_{ik}A\_{kj}$ (i.e., Eqn. 5 of the paper) - $A$ modeled by independent Gaussians: $P(A|\theta)=\mathcal{N}(A|M,\Omega^2)$ - $B$ modeled as deterministic values 2. Presentation of Theorem 3.1 - Statement: Asymmetric Bayesianization corresponds to assuming a low-rank Gaussian variational distribution for the full weight $W$ 3. Discussion on the Choice of Asymmetric Bayesianization - Advantages over Bayesianizing both $A$ and $B$: + Stable training at the early stage + Faster convergence of parameters during training + Lower memory cost **Section 3.2: Low-Rank Prior Distribution** 1. Presentation of the Prior Distribution on the Low-Rank Component $A$ 2. Corresponding Prior Distribution on the High-Dimensional Space of the Full Weight Matrix $W$ 3. Presentation of Theorem 3.2 - Statement: Full-weight KL divergence can be efficiently computed on the low-rank component 4. Dedicated Remark on the Legitimacy of using Degenerate Gaussians for Probabilistic Inference (as in Our Previous Discussion) Thank you once again for your suggestions. We believe this revision will make the paper more accessible for our audience and help address any potential areas of confusion.
Summary
This work introduces a Bayesian Deep Learning framework for finetuning LLMs with probabilistic LoRA. Unlike existing work by Yang et al. [1] which uses a Laplace approximation, the variational distribution is parameterized using diagonal Gaussians as done in [2] (”Bayes by Backprop”). Specifically, the authors design the variational distribution as follows: Given the standard LoRA updates $W = W_0 + BA$, only $A$ is modeled probabilistically while $B$ is deterministic. In combination with an appropriate choice of prior, this allows for a a low-rank parameterization of the variational posterior and fast KL-computation in the lower dimensional parameter space. Experiments compare the methods to existing baselines in terms of accuracy of the mean predictions, calibration (ECE) and NLL on a range of datasets. [1] https://arxiv.org/pdf/2308.13111 [2] https://proceedings.mlr.press/v37/blundell15.pdf ____ I have updated my overall score +1 after the rebuttal.
Strengths
The problem of assigning principled uncertainties to (fine-tuned) LLM outputs is important, and it makes sense to compare different Bayesian Deep Learning strategies in combination with LoRA. The exposition of the method and paper in general is clear.
Weaknesses
Novelty is limited compared to Yang et al. [1], which implements a very similar idea via Laplace approximations instead of the Bayes-by-Backprop variational posterior. Additionally, the method by Yang et al. ("LAP") is not replicated successfully in Table 1. For example, for the LAP method on ARC-C datasets the results in this work fall behind the results in [1] by 50%: - This work: Accuracy: 29.73±12.02, ECE: 14.24±1.65 and NLL: 1.53±0.01 - [1]: Accuracy: 66.91±1.1, ECE: 7.5±1.2 and NLL: 0.86±0.02, respectively. This makes it difficult to gauge how much the approach improves compared to [1].
Questions
Questions: - Table 1: As $N \rightarrow \infty$, would we expect the accuracy to converge to those obtained with $N=0$ (mean prediction)? Is BLoB an unbiased estimator of the mean prediction? - l. 240: “For MCD, we use an ensemble of 10 LoRAs with a dropout rate of p = 0.1”. Why is an ensemble used here? Shouldn’t the uncertainties come from the dropout variational posterior? Why are multiple models necessary? - Will you open source the code on publication? Suggestions & Typos: - l. 33: I.e., unable to … (sentence is not grammatical, delete) - l. 110: yielded by parameterization —> yielded by reparameterization - Algorithm 1, line 11: Updater —> Update - l. 237: Maximize A Posteriori —> Maximum A Posteriori - l. 239: “For MAP, we use a weight decay rate of 1e-5.” I suggest to formulate this in terms of prior variance instead. - l. 638 covaraince —> covariance
Rating
6
Confidence
3
Soundness
3
Presentation
3
Contribution
2
Limitations
No, I recommend a more explicit treatment of the limitations, perhaps in a dedicated section.
Thank you
Thank you very much for the encouraging comments and for acknowledging our contribution. We are glad that our response has been helpful and convincing. We will include the discussion and ablation study in the revision based on your suggestion.
Response to rebuttal
Thanks for the additional details, especially the derivation and experiment in Q1. I see that $$E_{q(W|\theta)}[P(Y | W, X)] \neq P(Y | E_{q(W|\theta)} [W], X).$$ Follow up: What's the intuition behind the mean prediction (RHS) achieving higher accuracy than the (approximate) Bayesian prediction (LHS), also for large $N$? Overall, the rebuttal partly addresses my concerns around novelty (Q2 and Q3 in "General response"), hence I increase my overall score by one point.
Thank You for Your Further Feedback
Thank you for your further feedback and for maintaining an open line of communication. We are glad that you found our response helpful and that it addressed your concerns. Below we address your follow-up comment on Q1. **[Higher Accuracy of Mean Prediction Compared to Bayesian Prediction]** This is a good question. By abandoning the modeling of the posterior distribution, mean prediction sacrifices some degree of calibration in exchange for improved accuracy. This empirical trade-off between accuracy and calibration has been noted in [3]. We will include this discussion in our revision. Lastly, we would like to express our gratitude once again for your insightful and constructive comments. They have significantly contributed to the improvement of our paper. [3] Stengel-Eskin, Elias, and Benjamin Van Durme. "Calibrated interpretation: Confidence estimation in semantic parsing." Transactions of the Association for Computational Linguistics 2023
Decision
Accept (poster)