Bayesian LoRA via Mean-Field VI

• The paper's main contribution is introducing a Bayesian framework that applies mean-field variational inference to Low-Rank Adaptation (LoRA) for uncertainty-aware fine-tuning.
• It employs a low-rank plus diagonal covariance structure with SWAG-inspired updates to achieve efficient calibration and reduced computational overhead.
• Experimental results on GLUE tasks demonstrate that the method delivers strong predictive performance with energy savings while robustly quantifying uncertainty.

Bayesian LoRA by Mean-Field Variational Inference is an approach that unifies parameter-efficient fine-tuning and uncertainty quantification in LLMs by applying Bayesian principles to Low-Rank Adaptation (LoRA) using mean-field variational inference (VI). This methodology emphasizes low-dimensional parameterization and storage efficiency, enabling strong predictive performance, calibration, and energy savings in downstream adaptation scenarios.

1. Probabilistic Model Formulation

The foundation of Bayesian LoRA is the adaptation of a frozen pre-trained weight matrix $W^0 \in \mathbb{R}^{m\times n}$ through a low-rank update: $\Delta W = A R B,$ where $A \in \mathbb{R}^{m \times r}$ and $B \in \mathbb{R}^{r \times n}$ are fixed projection matrices constructed from a truncated SVD of $W^0$: $W^0 \approx U_r S_r V_r^\top$ with $A = U_r S_r$ and $B = V_r^\top$. The core trainable (and stochastic) component is the small inner matrix $R \in \mathbb{R}^{r \times r}$ for each LoRA module.

The parameters of interest $\theta$ are defined as the vectorization of all $R$ matrices: $$\theta = \left[\mathrm{vec}(R_1); \mathrm{vec}(R_2); \dots \right].$$ A zero-mean isotropic Gaussian prior is imposed: $p(\theta) = \mathcal{N}(\theta | 0, \tau^2 I).$ The likelihood over data $\mathcal{D} = \{(x_i, y_i)\}_i$ is given by the standard model output, e.g., softmax loss for classification tasks: $$p(\mathcal{D} | \theta) = \prod_i \text{SoftmaxLoss}(y_i, f(x_i; W^0 + \Delta W(\theta))),$$ where $f(\cdot; W)$ denotes the network with weight $W$ at each layer. The stochasticity in the weight matrix resides strictly in $R$, confining all learning and uncertainty modeling to a small low-dimensional subspace (Marszałek et al., 17 Feb 2025).

2. Variational Family and Posterior Structure

Mean-field VI is employed to approximate the intractable posterior $p(\theta|\mathcal{D})$. In the B-LoRA-XS framework, the variational distribution $q(\theta)$ is taken as a Gaussian with a low-rank plus diagonal covariance structure: $$q(\theta) = \mathcal{N}(\mu, \Sigma), \qquad \Sigma = UU^\top + \mathrm{diag}(\nu),$$ with $U \in \mathbb{R}^{n \times k}$ ($k \ll n$) capturing a low-dimensional subspace of correlations and $\nu \in \mathbb{R}_+^n$ providing per-parameter variance. Unlike standard mean-field where all correlations are ignored, this construction retains aggregate global dependencies in a rank-$k$ subspace while maintaining linear storage cost in $n$.

The mean $\mu$ is tracked as a running average of $\theta$ during post–burn-in epochs, and $U$ is constructed from the latest $k$ deviation vectors (denoted $D = [d_1, ..., d_k]$ with $d_i = \theta_i - \mu$), following the SWAG moment-matching prescription. Per-parameter variance $\sigma^2 = \mathbb{E}[(\theta - \mu)^2]$ is estimated in parallel (Marszałek et al., 17 Feb 2025).

This approach can be conceptually reconciled with a mean-field variational family whose marginal posterior for each $R$ parameter is (up to low-rank cross-covariances) an independent Gaussian, while cross-layer correlations are captured in $U$.

3. ELBO Objective and Low-Dimensional Projections

Although the empirical procedure in B-LoRA-XS utilizes SWAG, the underlying optimization closely corresponds to implicit maximization of the evidence lower bound (ELBO): $$\operatorname{ELBO}(\mu, U, \nu) = \mathbb{E}_{\epsilon \sim \mathcal{N}(0, I_k), \eta \sim \mathcal{N}(0, I_n)}\left[\sum_{i} \log p(y_i|x_i; \mu + U\epsilon + \mathrm{diag}(\sqrt{\nu})\eta)\right] - \mathrm{KL}\left[\mathcal{N}(\mu, UU^\top+\mathrm{diag}(\nu)) \,||\, \mathcal{N}(0, \tau^2 I)\right].$$ In practice, no explicit backpropagation is performed through the KL; instead, the second moments are tracked in SWAG fashion. The fixed SVD projection ($A$, $B$) ensures that all Bayesian inference occurs in the subspace spanned by $R$, and the posterior covariance is constrained to be low rank plus diagonal at scale, minimizing parameter and computational overhead.

The low-rank structure appears both in the weight update and in the construction of $\Sigma$. These two factors—explicit subspace parameterization of $\Delta W$ and low intrinsic dimensionality of posterior uncertainty—enable the approach to outperform or match alternatives at a fraction of the memory and compute footprint (Marszałek et al., 17 Feb 2025).

4. Training Algorithm and Gradient Estimation

The training protocol, as aligned with mean-field VI, supports unbiased gradient estimation for $q$'s parameters via the reparameterization trick: $$\theta(\epsilon, \eta) = \mu + U\epsilon + \mathrm{diag}(\sqrt{\nu})\eta, \qquad \epsilon \sim \mathcal{N}(0, I_k), \; \eta \sim \mathcal{N}(0, I_n).$$ Gradients with respect to $\mu, U, \nu$ are in principle estimated by: $$\nabla_{\mu,U,\nu} \, \mathbb{E}_{\theta \sim q} [\log p(\mathcal{D}|\theta)] = \mathbb{E}_{\epsilon, \eta}\left[\nabla_{\mu,U,\nu} \log p(\mathcal{D}|\theta(\epsilon, \eta))\right].$$ However, in B-LoRA-XS, SGD is used to update $\theta$ directly, and after a burn-in period, low-rank covariances for posterior approximation are accumulated according to SWAG: second moments are constructed from the empirical deviations as model parameters traverse the solution landscape with a fixed learning rate. These statistics define the variational distribution used for Bayesian predictive inference (Marszałek et al., 17 Feb 2025).

A high-level algorithmic loop is as follows:

for epoch in training_epochs: for batch in data: if past_burn_in: # Sample theta from low-rank+diag Gaussian theta_sample = mu + U @ eps + sqrt(sq_avg - mu**2) * eta else: theta_sample = current_theta # Compute ΔW_l = A_l R_l B_l, forward/loss/backprop/update # SWAG update: accumulate mu, sq_avg, buffer last k deviations

5. Covariance Structure and Parameter Efficiency

For a concatenated $\theta$ vector of dimension $N = \sum_l r_l^2$, a dense covariance would cost $\mathcal{O}(N^2)$ in storage. B-LoRA-XS' low-rank plus diagonal approach requires only $N(k + 2)$ parameters ($k \ll N$), achieving linear scaling. The induced covariance of the actual weight update, with $\mathrm{vec}(\Delta W)$, is

$$\mathrm{Cov}(\mathrm{vec}(\Delta W)) = (B^\top \otimes A)\ \Sigma_R\ (B^\top \otimes A)^\top,$$

where $\Sigma_R$ is the covariance in $R$–space. Given $R$ is $r \times r$ and $\Sigma_R$ is rank $k$, $\Delta W$'s covariance is low rank with at most $k r$ nonvanishing directions. This property ensures not only memory efficiency but also that almost all uncertainty mass is confined to a tractable, easily sampled region of parameter space (Marszałek et al., 17 Feb 2025).

6. Practical Settings, Empirical Results, and Comparative Context

Empirical validation in (Marszałek et al., 17 Feb 2025) is conducted on four GLUE tasks (RTE, MRPC, CoLA, SST-2) using RoBERTa-large as the backbone, with B-LoRA-XS adapters injected into all major transformer projection submodules (Query, Value, Attention-Output, Output-FC). Typical parameter settings include adaptation ranks $r \in \{2,8,16,25\}$, covariance rank $k \in \{2,5,10,\ldots\}$, AdamW optimization, 10–25 burn-in epochs, and inference via $S=15$ posterior samples.

Key observed outcomes:

• B-LoRA-XS matches or exceeds SWAG-LoRA on calibration (ECE) and NLL, with 5–15× fewer Bayesian parameters.
• The SVD-based low-rank projection effectively disentangles adapted directions, enabling strong uncertainty quantification with $k$ as low as 2–5.
• Training is more stable and less sensitive to random seed than SWAG-LoRA.

A related approach, Bayesian-LoRA (Meo et al., 2024), incorporates mean-field VI for both rank and quantization selection via discrete Bayesian gates with Gumbel-sigmoid relaxations, thereby offering flexible per-layer adaptation and further reducing bit-operations required. However, in contrast to B-LoRA-XS, (Meo et al., 2024) does not focus on uncertainty quantification metrics such as ECE or NLL, but rather on adaptive rank/quantization and compute efficiency.

7. Extensions to Quantization and Rank Adaptation

The Bayesian-LoRA framework (Meo et al., 2024) extends mean-field VI in LoRA to include not only Gaussian uncertainty over continuous low-rank factors but also discrete latent variables governing quantization levels and effective ranks. This is achieved by:

• Factorizing the variational posterior over low-rank matrices ($A$, $B$), quantization gates ($z$), and rank-selection gates ($g$).
• Employing continuous relaxations (straight-through Gumbel-sigmoid) for the Bernoulli gates, enabling end-to-end training with reparameterized gradients.
• Penalizing KL divergence of quantization/rank gates in the objective, balancing compression and accuracy.

Empirical results in (Meo et al., 2024) show that Bayesian-LoRA achieves competitive or superior task performance to baseline LoRA variants at markedly lower memory and operation counts, with learned adaptations tailored per layer and module. A plausible implication is that the Bayesian VI perspective enables simultaneous adaptation, selective precision allocation, and energy-aware deployment in large-scale LLM fine-tuning.

---

For technical readers seeking full Bayesian fine-tuning recipes, the B-LoRA-XS framework rigorously specifies the probabilistic model, variational approximation, training regimen, and inference algorithm, realizing both uncertainty quantification and efficiency through judicious low-dimensional projection and mean-field modeling (Marszałek et al., 17 Feb 2025).