Bayesian Low-Rank LeArning (Bella): A Practical Approach to Bayesian Neural Networks

TL;DR

The paper addresses the computational intractability of Bayesian neural networks by introducing Bella, a low-rank perturbation framework that leverages a fixed pre-trained base model plus small, trainable adapters. By performing SVGD updates on low-rank factors, Bella achieves scalable posterior estimation with dramatically reduced parameter counts while maintaining competitive accuracy and uncertainty quantification across large-scale vision tasks. Key contributions include the concrete low-rank parameterization, kernelized SVGD updates on adapter factors, and extensive empirical validation showing favorable cost/benefit trade-offs on datasets like ImageNet, CAMELYON17, DomainNet, and VQA with LLaVA. The method enables practical Bayesian deep learning in real-world applications, offering strong out-of-distribution robustness and calibrated uncertainty with substantial resource savings.

Abstract

Computational complexity of Bayesian learning is impeding its adoption in practical, large-scale tasks. Despite demonstrations of significant merits such as improved robustness and resilience to unseen or out-of-distribution inputs over their non- Bayesian counterparts, their practical use has faded to near insignificance. In this study, we introduce an innovative framework to mitigate the computational burden of Bayesian neural networks (BNNs). Our approach follows the principle of Bayesian techniques based on deep ensembles, but significantly reduces their cost via multiple low-rank perturbations of parameters arising from a pre-trained neural network. Both vanilla version of ensembles as well as more sophisticated schemes such as Bayesian learning with Stein Variational Gradient Descent (SVGD), previously deemed impractical for large models, can be seamlessly implemented within the proposed framework, called Bayesian Low-Rank LeArning (Bella). In a nutshell, i) Bella achieves a dramatic reduction in the number of trainable parameters required to approximate a Bayesian posterior; and ii) it not only maintains, but in some instances, surpasses the performance of conventional Bayesian learning methods and non-Bayesian baselines. Our results with large-scale tasks such as ImageNet, CAMELYON17, DomainNet, VQA with CLIP, LLaVA demonstrate the effectiveness and versatility of Bella in building highly scalable and practical Bayesian deep models for real-world applications.

Figures (26)

• Figure 1: Train and Test Error (%) landscapes of CAMELYON17. Training landscape demonstrates the learned 3 particle approximation of the modes of the posterior from a pre-trained model $\boldsymbol\theta$---modification of parameters in the constrained region ($\Delta\boldsymbol\theta$) leads to approaching posterior modes in inference. Here, we observe a key benefit of our Bella approximation compared to a point estimate---while a single parameter particle (e.g. $\boldsymbol\theta +\Delta\boldsymbol\theta_1$) do not generalize well, Bayesian prediction with \ref{['eqn: bma']}, effectively an average over multiple parameter settings, leads to better performance.
• Figure 2: Evaluation of uncertainty estimations using Mutual Information on CAMELYON17 and CIFAR-10-C datasets. $\uparrow$ MI for Misclassified Examples is better---denoted by the distribution shifting $\rightarrow$. In contrast, $\downarrow$ MI for Correctly Classified Examples is better---denoted by the distribution shifting $\leftarrow$.
• Figure 3: Comparison of model robustness to $L_{\infty}$ FGSM adversarial attacks across varied attack budgets on the CIFAR-10 dataset.
• Figure 4: The impact of ranks on CAMELYON17 performance (left), as well as the impact of the number of parameter particles on CIFAR-10 (right) on Bella performance.
• Figure 5: Bella achieves similar Accuracy with full SVGD (Base) only with a fraction of cost.