Scalable Variational Bayesian Fine-Tuning of LLMs via Orthogonalized Low-Rank Adapters
Authors: Haotian Xiang, Bingcong Li, Qin Lu
Organizations: School of Electrical and Computer Engineering, University of Georgia, Athens, GA, USA · Department of Computer Science, ETH Zurich, Zurich, Switzerland
When deploying large language models (LLMs) to safety-critical applications, uncertainty quantification (UQ) is of utmost importance to self-assess the reliability of the LLM-based decisions. However, such decisions typically suffer from overconfidence, particularly after parameter-efficient fine-tuning (PEFT) for downstream domain-specific tasks with limited data. Existing methods to alleviate this issue either rely on Laplace approximation based post-hoc framework, which may yield suboptimal calibration depending on the training trajectory, or variational Bayesian training that requires multiple complete forward passes through the entire LLM backbone at inference time for Monte Carlo estimation, posing scalability challenges for deployment. To address these limitations, we build on the Bayesian last layer (BLL) model, where the LLM-based deterministic feature extractor is followed by random last layer parameters for uncertainty reasoning. Since existing low-rank adapters (LoRA) for PEFT have limited expressiveness due to rank collapse, we address this with Polar-decomposed Low-rank Adapter Representation (PoLAR), an orthogonalized parameterization paired with Riemannian optimization to enable more stable and expressive adaptation. Building on this PoLAR-BLL model, we leverage the variational (V) inference framework to put forth a scalable Bayesian fine-tuning approach which jointly seeks the PoLAR parameters and approximate posterior of the last layer parameters via alternating optimization. The resulting PoLAR-VBLL is a flexible framework that nicely integrates architecture-enhanced optimization with scalable Bayesian inference to endow LLMs with well-calibrated UQ. Our empirical results verify the effectiveness of PoLAR-VBLL in terms of generalization and uncertainty estimation on both in-distribution and out-of-distribution data for various common-sense reasoning tasks.
Figures & tables
In-Distribution Datasets
Out-of-Distribution Datasets (OBQA → X)
Small Shift
Large Shift
Metric
Method
WG-S
ARC-C
ARC-E
WG-M
OBQA
BoolQ
ARC-C
ARC-E
Chem
Phy
ACC ( ↑ )
MLE
77.92 ±0.62
81.05 ±1.62
90.66 ±0.10
82.80 ±0.96
88.30 ±0.36
87.86 ±0.50
79.33 ±0.64
85.66 ±0.50
48.00 ±2.00
43.33 ±1.53
LA
77.92 ±0.62
81.05 ±1.62
90.66 ±0.10
82.80 ±0.96
88.30 ±0.36
87.86 ±0.50
79.33 ±0.64
85.66 ±0.50
48.00 ±2.00
43.33 ±1.53
PoLAR-MLE
78.09 ±0.39
81.21 ±1.02
90.25 ±1.24
83.61 ±0.32
86.67 ±1.22
89.16 ±0.27
80.97 ±0.76
85.57 ±0.21
48.33 ±0.58
43.33 ±2.52
PoLAR-LA
78.09 ±0.39
81.21 ±1.02
90.25 ±1.24
83.61 ±0.32
86.67 ±1.22
89.16 ±0.27
80.97 ±0.76
85.57 ±0.21
48.33 ±0.58
43.33 ±2.52
Table 1: Performance of different methods on Llama-3.1-8B . ACC and ECE are reported in percentages. The evaluation is done across six in-distribution common-sense reasoning datasets with fine-tuning of 5000 steps. For OOD evaluation, models are trained on OBQA and tested on other datasets. Bold and underlined denote the best and second-best performance, respectively.
Figure 1: Ablation studies on the WG-S dataset using LLaMA 2-7B : (a) Performance of LoRA-VBLL using different ranks; (b) VBLL coupled with different adapters; and (c) ECE and NLL performances of PoLAR-VBLL with and without LA.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Metric
Method
Datasets
WG-S
ARC-C
ARC-E
OBQA
ACC (%)
MLE
68.99 ± 0.58
69.10 ± 2.84
85.65 ± 0.92
81.52 ± 0.25
MAP
68.62 ± 0.71
67.59 ± 0.40
86.55 ± 0.55
81.38 ± 0.65
MCD
69.46 ± 0.62
68.69 ± 1.30
86.21 ± 0.46
81.72 ± 0.10
ENS
69.57 ± 0.66
66.20 ± 2.01
84.40 ± 0.81
81.38 ± 0.91
BBB
66.54 ± 7.87
68.13 ± 1.27
86.86 ± 0.74
82.06 ± 0.59
Appendix
Table 2: Performances on ID datasets in terms of ACC, ECE, and NLL using LlaMA2-7B . Bold and underlined denote the best and the second-best performance, respectively. Here, we include PoLAR-VBLL with and without LA.
Method
Training Memory (MB)
Test Memory (MB)
Inference Time (s)
MLE
19,474
17,298
8
PoLAR-MLE
19,556
17,150
7
PoLAR-VBLL (Ours)
20,178
18,423
12
PoLAR-BLoB
19,738
19,830
80
LoRA-BLoB
20,972
18,762
89
PoLAR-LA-LL
19,556
18,074
30
Appendix
Table 3: GPU memory usage and inference time comparison across different uncertainty quantification methods. Configuration: LoRA rank = 8, batch size = 4, dataset = ARC-E, LoRA targets = [ q_proj , v_proj ].
Method
Variational Parameters per Layer
Calculation
LoRA-BLoB
131,072
r×d×2=16×4096×2
PoLAR-BLoB
512
r×r×2=16×16×2
ScalaBL
32
r×2=16×2
Appendix
Table 4: Variational parameters per layer for different methods.
In-Distribution Datasets
Out-of-Distribution Datasets (OBQA → X)
Small Shift
Large Shift
Metric
Method
WG-S
ARC-C
ARC-E
WG-M
OBQA
BoolQ
ARC-C
ARC-E
Chem
Phy
ACC ( ↑ )
MLE
77.92 ±0.62
81.05 ±1.62
90.66 ±0.10
82.80 ±0.96
88.30 ±0.36
87.86 ±0.50
79.33 ±0.64
85.66 ±0.50
48.00 ±2.00
43.33 ±1.53
LA
77.92 ±0.62
81.05 ±1.62
90.66 ±0.10
82.80 ±0.96
88.30 ±0.36
87.86 ±0.50
79.33 ±0.64
85.66 ±0.50
48.00 ±2.00
43.33 ±1.53
MAP
77.32 ±1.31
80.59 ±0.86
90.01 ±0.37
81.96 ±0.36
88.13 ±0.64
88.54 ±0.34
79.50 ±0.56
85.29 ±0.61
47.33 ±1.15
44.00 ±5.00
MCD
76.92 ±0.92
81.14 ±1.50
91.01 ±0.40
82.90 ±0.39
88.00 ±0.20
88.62 ±0.15
80.33 ±0.52
84.78 ±0.61
46.67 ±2.89
43.67 ±2.08
Appendix
Table 5: Performance of different methods on Llama-3.1-8B . ACC and ECE are reported in percentages. The evaluation is done across six in-distribution common-sense reasoning datasets with fine-tuning of 5000 steps. For OOD evaluation, models are trained on OBQA and tested on other datasets. Bold and underlined denote the best and second-best performance, respectively.
Method
ACC (%) ↑
ECE (%) ↓
NLL ↓
PoLAR-LA
70.33 ± 0.69
12.16 ± 2.58
0.69 ± 0.03
PoLAR-LA-LL
70.33 ± 0.69
14.63 ± 1.14
0.71 ± 0.05
PoLAR-BLoB
70.39 ± 0.26
12.06 ± 0.81
0.73 ± 0.04
PoLAR-VBLL (w/o LA)
71.62 ± 0.27
8.26 ± 0.60
0.66 ± 0.03
PoLAR-VBLL (Full)
71.62 ± 0.27
7.31 ± 0.32
0.60 ± 0.01
Appendix
Table 6: Ablation study comparing PoLAR-VBLL with PoLAR-based baselines on WG-S using Llama-2-7B . All methods use identical PoLAR adapters, isolating the contribution of different uncertainty quantification approaches.
Training steps
VBLL (Jensen)
VBLL (50-sample MC)
Absolute Gap
0
69.50
61.23
8.27
50
50.71
50.97
0.26
100
45.57
45.65
0.08
150
40.95
40.86
0.09
200
36.96
36.95
0.01
250
33.57
33.74
0.17
Appendix
Table 7: Comparison of training loss trajectories between the Jensen bound and 50-sample Monte Carlo estimation on WGS dataset using Llama-3.1-8B .
Prior Scale ( σθ )
ACC ( ↑ )
ECE ( ↓ )
NLL ( ↓ )
0.1
70.92±0.24
9.10±0.53
0.68±0.04
1.0 (Default)
71.62±0.27
8.26±0.60
0.66±0.03
10.0
70.70±0.50
10.59±1.17
0.69±0.07
Appendix
Table 8: Sensitivity analysis of prior scale σθ on WG-S under LLaMA 2 7B. Results are averaged over three random seeds with standard deviations reported. ACC: accuracy (%), ECE: expected calibration error (%), NLL: negative log-likelihood.
Figure 2: Stable rank comparison between PoLAR and LoRA across three datasets. PoLAR consistently achieves higher stable rank values, indicating better preservation of feature geometry and effective utilization of the allocated parameter space.
Large language models (LLMs) exhibit remarkable reasoning capabilities, but their task-specific fine-tuning is notoriously plagued by overconfidence, severely hindering trustworthy deployment. We propose Data-Adaptive Lower-Rank Adaptation (DALorRA), a simple and effective variational Bayesian sparse framework that shifts the paradigm of uncertainty quantification from the dense parameter space to the lightweight rank level of low-rank adaptation (LoRA). With the insight that LoRA essentially aggregates multiple rank-one components that may provide superfluous model capacity, DALorRA imposes stochastic masking on rank dimensions, enabling Bayesian regularization of model capacity during training and ensemble-like calibration during inference. Extensive experiments demonstrate DALorRA's excellent calibration of LLMs without compromising reasoning accuracy.
Jijie Zhang, Zhe Ren, Quan Zhang +1
School of Artificial Intelligence, Jilin University · Michigan State University
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 ≈1.2× 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.
Moule Lin, Shuhao Guan, Andrea Patane +2
School of Computer Science and Statistics, Trinity College Dublin, Dublin, Ireland and Lero the Research Ireland Centre for Software, Ireland · School of Computer Science, University College Dublin, Dublin, Ireland
Low-Rank Adaptation (LoRA) has become the standard for fine-tuning large pre-trained models at reduced computational cost. However, its low-rank point-estimate updates limit expressiveness, leave a persistent gap relative to full fine-tuning accuracy, and provide no built-in uncertainty quantification, limiting its applicability in settings where reliability matters as much as accuracy. We introduce BaLoRA, a Bayesian extension of LoRA with a novel input-adaptive Bayesian parameterization of LoRA matrices that adds minimal parameters and compute. Surprisingly, not only does the Bayesian extension yield well-calibrated uncertainty estimates, but the adaptive noise injection underlying our approach also significantly improves prediction accuracy, narrowing the gap with full fine-tuning across both natural language reasoning and vision tasks. When applied to band gap prediction in metal-organic frameworks, BaLoRA produces zero-shot test-time uncertainty estimates that correlate more strongly with model error than a trained ensemble of LoRA models, and improve monotonically with compute without sacrificing accuracy.