- The paper presents a novel framework, PoLAR-VBLL, that advances uncertainty quantification in LLMs by integrating orthogonal low-rank adapters with variational Bayesian last-layer training.
- It employs polar decomposition with Stiefel manifold constraints to preserve stable rank and maintain robust feature geometry, enhancing in-distribution and out-of-distribution calibration.
- Experiments demonstrate superior accuracy and lower calibration error with a tenfold reduction in inference time compared to existing Bayesian fine-tuning methods.
Scalable Variational Bayesian Fine-Tuning of LLMs via Orthogonalized Low-Rank Adapters
Context and Motivation
The deployment of LLMs in safety-critical scenarios necessitates robust uncertainty quantification (UQ) for self-assessment of prediction reliability. Traditional PEFT approaches such as LoRA, while computationally feasible, frequently induce overconfidence and suboptimal calibration, particularly with limited domain-specific data. Prior UQ techniques suffer from either post-hoc approximations that depend heavily on deterministic training trajectories (e.g., Laplace) or Bayesian fine-tuning schemes that demand prohibitive inference costs due to repeated backbone forward passes (e.g., BLoB). The geometrical expressiveness of low-rank adaptation layers is central to effective UQ, yet standard LoRA suffers from extreme rank collapse, diminishing the dimensionality of learned feature spaces and undermining distance-aware uncertainty discriminability.
Methodological Contributions
The paper proposes the PoLAR-VBLL framework, integrating polar decomposed, orthogonality-constrained low-rank adapters (PoLAR) with variational Bayesian last-layer training (VBLL):
- PoLAR Adapter Architecture: The adaptation uses polar decomposition with Stiefel manifold constraints, preventing collapse of stable rank and thus enabling multi-directional feature geometry preservation. This property is theoretically and empirically validated to improve distance awareness, which is crucial for separating in-distribution (ID) from out-of-distribution (OOD) samples in UQ.
- VBLL Uncertainty Quantification: The model implements variational inference solely at the classification layer (last layer), maximizing the ELBO via a closed-form Jensen-tightened bound. The posterior is parameterized with per-class Gaussian distributions, capturing nontrivial covariance structure.
- Riemannian Optimization Techniques: Optimization on orthogonality-constrained parameters leverages infeasible landing field methods rather than standard retraction-based Riemannian algorithms, offering significant speedup for large-scale models.
- Inference and Laplace Refinement: Predictive uncertainty is evaluated by Monte Carlo sampling over the posterior of last-layer parameters, requiring only one backbone pass per test input. Optionally, post-hoc Laplace approximation is applied to locally refine the variational covariance, using the VBLL-discovered mode as initialization—empirically improving calibration without compromising accuracy.
Empirical Evaluation
Extensive experiments are conducted on Llama-3.1-8B and Llama-2-7B backbones, encompassing six common-sense reasoning benchmarks and additional OOD science tasks:
- Accuracy and Calibration: PoLAR-VBLL attains top or second-best accuracy on all ID and OOD benchmarks. It achieves lower expected calibration error (ECE) and negative log-likelihood (NLL) than all other methods, including baseline PEFT (MLE, MAP), deep ensembles, Monte Carlo dropout, Laplace-LoRA, BLoB, C-LoRA, and TFB variants.
- Stable Rank Preservation: Empirical analysis shows that PoLAR maintains average stable rank significantly above standard LoRA (2.86 vs. ~1.5), ensuring robust geometric diversity in feature extraction.
- Computational Efficiency: The framework delivers approximately a one order-of-magnitude reduction in inference time compared to BLoB-based Bayesian adapters, owing to its sampling restriction to the classification layer. Memory footprint is competitive and often lower than full-model Laplace approaches.
- Ablations and Sensitivity: Ablation analyses confirm that VBLL is the primary driver for calibration gain, with optional Laplace refinement providing incremental improvements. The Jensen tightness analysis demonstrates rapid convergence and stability of the bound throughout training. Sensitivity studies suggest robustness to initialization and prior scale.
Theoretical and Practical Implications
The PoLAR-VBLL design substantiates the critical role of geometric properties in PEFT adapters for downstream Bayesian inference quality. By combining orthogonality-preserving adaptation with scalable variational last-layer learning, the framework resolves longstanding UQ bottlenecks—namely, overconfidence after fine-tuning and inefficient sampling during downstream deployment. The practical pathway for rapid, memory-efficient, and trustworthy LLM adaptation for domain-specific tasks is established, facilitating deployment in settings where reliable uncertainty assessment is paramount (e.g., medical assistance, legal analysis, autonomous systems).
From a theoretical perspective, the synergy of PoLAR and VBLL demonstrates that the expressiveness of adapted feature spaces remains the limiting factor for Bayesian last-layer calibration and OOD detection, opening further investigation into geometric regularization of adapter architectures and integration with advanced Bayesian learning modalities for fine-tuned LLMs.
Future Directions
Continuous improvement in adapter architecture can be pursued, with focus on further mitigating rank collapse and optimizing geometric dispersion in feature embedding. Exploration of hierarchical Bayesian schemes, deep kernel learning integration, and efficient posterior refinement mechanisms (beyond Laplace) could enhance both the scalability and the quality of UQ in LLM adaptation. Additional work is warranted in characterizing uncertainty under extreme domain shifts and adversarial settings, and in generalizing the framework to multitask or multi-modal adaptation.
Conclusion
PoLAR-VBLL presents a unified and scalable paradigm for uncertainty-aware fine-tuning of LLMs, leveraging orthogonalized low-rank adapters and variational Bayesian classification layers. It offers superior calibration, accuracy, and inference efficiency relative to prior art, validating its applicability for trustworthy LLM adaptation in safety-critical deployments (2604.03388). The work sets a foundation for principled, architecture-aware advancements in Bayesian PEFT for future AI systems.
