Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Bayesian-LoRA: Probabilistic Low-Rank Adaptation of Large Language Models

Moule Lin, Shuhao Guan, Andrea Patane, David Gregg, Goetz Botterweck

TL;DR

Bayesian-LoRA reframes LoRA’s deterministic, low-rank fine-tuning as a probabilistic, flow-augmented model by introducing inducing variables $U$ and a Kronecker-structured prior, enabling end-to-end calibration during training. The approach builds a structural isomorphism between LoRA’s bilinear updates and Kronecker-factored sparse Gaussian process posteriors, with a flow $T_\phi$ enriching posterior flexibility while keeping overhead small. Empirical results across commonsense reasoning, WikiText-2 language modeling, and math reasoning on models up to 14B dense and 30B MoE demonstrate significant reductions in NLL and ECE, especially under distribution shift, with only about ${\approx}0.42$M additional parameters and ${\approx}1.2\times$ training time relative to standard LoRA. The method yields robust uncertainty estimates and favorable out-of-distribution performance, suggesting calibrated end-to-end Bayesian fine-tuning as a practical alternative to costly ensembles or post-hoc calibration.

Abstract

Large Language Models usually put more emphasis on accuracy and therefore, will guess even when not certain about the prediction, which is especially severe when fine-tuned on small datasets due to the inherent tendency toward miscalibration. In this work, we introduce Bayesian-LoRA, which reformulates the deterministic LoRA update as a probabilistic low-rank representation inspired by Sparse Gaussian Processes. We identify a structural isomorphism between LoRA's factorization and Kronecker-factored SGP posteriors, and show that LoRA emerges as a limiting case when posterior uncertainty collapses. We conduct extensive experiments on various LLM architectures across commonsense reasoning benchmarks. With only approximately 0.42M additional parameters and ${\approx}1.2{\times}$ training cost relative to standard LoRA, Bayesian-LoRA significantly improves calibration across models up to 30B, achieving up to 84% ECE reduction and 76% NLL reduction while maintaining competitive accuracy for both in-distribution and out-of-distribution (OoD) evaluations.

Bayesian-LoRA: Probabilistic Low-Rank Adaptation of Large Language Models

TL;DR

Bayesian-LoRA reframes LoRA’s deterministic, low-rank fine-tuning as a probabilistic, flow-augmented model by introducing inducing variables and a Kronecker-structured prior, enabling end-to-end calibration during training. The approach builds a structural isomorphism between LoRA’s bilinear updates and Kronecker-factored sparse Gaussian process posteriors, with a flow enriching posterior flexibility while keeping overhead small. Empirical results across commonsense reasoning, WikiText-2 language modeling, and math reasoning on models up to 14B dense and 30B MoE demonstrate significant reductions in NLL and ECE, especially under distribution shift, with only about M additional parameters and training time relative to standard LoRA. The method yields robust uncertainty estimates and favorable out-of-distribution performance, suggesting calibrated end-to-end Bayesian fine-tuning as a practical alternative to costly ensembles or post-hoc calibration.

Abstract

Large Language Models usually put more emphasis on accuracy and therefore, will guess even when not certain about the prediction, which is especially severe when fine-tuned on small datasets due to the inherent tendency toward miscalibration. In this work, we introduce Bayesian-LoRA, which reformulates the deterministic LoRA update as a probabilistic low-rank representation inspired by Sparse Gaussian Processes. We identify a structural isomorphism between LoRA's factorization and Kronecker-factored SGP posteriors, and show that LoRA emerges as a limiting case when posterior uncertainty collapses. We conduct extensive experiments on various LLM architectures across commonsense reasoning benchmarks. With only approximately 0.42M additional parameters and training cost relative to standard LoRA, Bayesian-LoRA significantly improves calibration across models up to 30B, achieving up to 84% ECE reduction and 76% NLL reduction while maintaining competitive accuracy for both in-distribution and out-of-distribution (OoD) evaluations.
Paper Structure (64 sections, 3 theorems, 53 equations, 5 figures, 17 tables)

This paper contains 64 sections, 3 theorems, 53 equations, 5 figures, 17 tables.

Table of Contents

  1. Introduction
  2. Limitations of existing approaches.
  3. Related Work
  4. Preliminaries
  5. LoRA: Low-Rank Adaptation
  6. Variational Sparse Inducing Weight Model
  7. Core idea.
  8. Prior and variational posterior on $U$.
  9. Conditional distribution of $W$ given $U$.
  10. Marginal distribution and ELBO.
  11. Methodology: Bayesian-LoRA
  12. From LoRA to Bayesian-LoRA
  13. Probabilistic low-rank update.
  14. Normalizing flow for posterior flexibility.
  15. Training Objective: Flow-Augmented ELBO
  16. ...and 49 more sections

Key Result

Proposition 3.1

Let $T_\phi$ be invertible and set $\Delta W=T_\phi(U)$. With pushforwards $q_\phi := T_{\phi\#}q_\psi$ and $p_\phi := T_{\phi\#}p$, we have Hence, the ELBO written in $U$-space equals the ELBO written in $\Delta W$-space. Proof. See Appendix proposition_appendix.

Figures (5)

  • Figure 1: Bayesian-LoRA overview. Inducing variables $U$ are transformed by a flow $T_\phi$ and projected via conditional Gaussians to produce stochastic LoRA matrices $A$, $B$. The effective weight $W_{\text{eff}} = W_{\text{pre}} + \frac{\alpha}{r}BA$ is used for the forward pass; Monte Carlo samples capture epistemic uncertainty.
  • Figure 2: Ablation on inducing-point dimension $r = c$. Increasing the dimension improves calibration (lower ECE) with diminishing returns beyond $r{=}16$.
  • Figure 3: Relative changes in ACC, NLL, ECE and inference time with respect to $S=1$ for the ID ARC dataset and the OOD OBQA dataset at a fixed checkpoint.
  • Figure 4: Pareto analysis on the WinoGrande-M dataset (NLL vs ECE). Each point denotes a hyperparameter pair $(\mathrm{lr}, \mathrm{wd})$. Gray regions show dominated solutions, and the cross marks the “Best choice” near the Pareto front.
  • Figure 5: Pareto analysis on the ARC-Easy dataset (ACC vs NLL and ACC vs ECE). Each point denotes a hyperparameter pair $(\mathrm{lr}, \mathrm{wd})$. Gray regions show dominated solutions, and the cross marks the “Best choice” near the Pareto front.

Theorems & Definitions (5)

  • Proposition 3.1: KL invariance under $T_\phi$
  • Corollary 4.1: Deterministic limit
  • proof
  • Proposition 1.1: Details of $U$-space-independent KL
  • proof