Scalable Bayesian Low-Rank Adaptation of Large Language Models via Stochastic Variational Subspace Inference
This addresses the challenge of making LLMs more reliable in high-stakes domains like autonomy and healthcare by providing a scalable Bayesian approach, though it is incremental as it builds on prior LoRA-based methods.
The paper tackles the problem of scaling Bayesian uncertainty quantification for large language models (LLMs) by proposing a method that performs inference in a low-dimensional subspace using LoRA parameters, achieving competitive performance with only about 1000 additional parameters and scaling to the largest Bayesian LLM to date.
Despite their widespread use, large language models (LLMs) are known to hallucinate incorrect information and be poorly calibrated. This makes the uncertainty quantification of these models of critical importance, especially in high-stakes domains, such as autonomy and healthcare. Prior work has made Bayesian deep learning-based approaches to this problem more tractable by performing inference over the low-rank adaptation (LoRA) parameters of a fine-tuned model. While effective, these approaches struggle to scale to larger LLMs due to requiring further additional parameters compared to LoRA. In this work we present $\textbf{Scala}$ble $\textbf{B}$ayesian $\textbf{L}$ow-Rank Adaptation via Stochastic Variational Subspace Inference (ScalaBL). We perform Bayesian inference in an $r$-dimensional subspace, for LoRA rank $r$. By repurposing the LoRA parameters as projection matrices, we are able to map samples from this subspace into the full weight space of the LLM. This allows us to learn all the parameters of our approach using stochastic variational inference. Despite the low dimensionality of our subspace, we are able to achieve competitive performance with state-of-the-art approaches while only requiring ${\sim}1000$ additional parameters. Furthermore, it allows us to scale up to the largest Bayesian LLM to date, with four times as a many base parameters as prior work.