BaLoRA: Bayesian Low-Rank Adaptation of Large Scale Models

TL;DR

BaLoRA is a Bayesian extension of LoRA that enhances uncertainty quantification and prediction accuracy in large-scale model fine-tuning with minimal additional parameters.

Contribution

It introduces a novel Bayesian parameterization of LoRA matrices that improves calibration, accuracy, and uncertainty estimation across NLP, vision, and materials science tasks.

Findings

BaLoRA yields well-calibrated uncertainty estimates.

It narrows the accuracy gap with full fine-tuning.

Uncertainty correlates strongly with model error in band gap prediction.

Abstract

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.