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