🤖 machine learning

Bayesian Fine-tuning in Projected Subspaces

This paper proposes a novel framework for parameter-efficient Bayesian fine-tuning that achieves effective uncertainty quantification and improved model calibration by projecting weight space uncertainty into very low-dimensional subspaces, thereby overcoming the high parameter costs and training instability of existing Bayesian LoRA methods.

Original authors: Viktar Dubovik, Patryk Marszałek, Jacek Tabor, Tomasz Kuśmierczyk

Published 2026-05-11

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Viktar Dubovik, Patryk Marszałek, Jacek Tabor, Tomasz Kuśmierczyk

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Overconfident" AI

Imagine you have a giant, incredibly smart encyclopedia (a Large Language Model) that knows almost everything. It's great, but it's also huge and expensive to update.

To teach it new things, we usually use a method called LoRA (Low-Rank Adaptation). Think of LoRA like adding a small, sticky-note layer to the encyclopedia. Instead of rewriting the whole book, you just write a few new notes on the pages. This is cheap and fast.

However, there's a catch: Standard LoRA is like a student who memorizes the notes perfectly but has no idea when they might be wrong. It becomes overconfident. If you ask it a question it doesn't know, it will guess with 100% certainty, even if it's wrong. In the real world, we need AI to know when it's unsure (uncertainty quantification).

To fix this, scientists tried Bayesian LoRA. This is like giving the student a "confidence meter" so they can say, "I'm 80% sure of this." But, there's a new problem: To build this confidence meter, the student needs to carry around a massive backpack of extra data. This makes the method slow and expensive again, defeating the purpose of the cheap "sticky notes."

The Solution: The "Tiny Room" Strategy

The authors of this paper propose a clever trick. They ask: Do we really need the whole backpack to measure confidence? Or can we do it in a tiny, organized room?

They introduce a new framework where they don't just add notes to the whole book. Instead, they:

Project the problem: They take the massive, complex task of updating the AI and shrink it down into a very small, low-dimensional "room" (a subspace).
Freeze the walls: They use the original encyclopedia's structure to build the walls of this room. These walls are fixed and don't move.
Furnish the room: They only learn how to arrange the furniture (the core parameters) inside this tiny room.

The Analogy:
Imagine you are trying to rearrange a massive, chaotic warehouse (the AI's brain).

Standard LoRA is like hiring a crew to move a few boxes around. It's fast, but they don't know if the boxes are fragile.
Old Bayesian LoRA is like hiring a crew with a full team of safety inspectors, blueprints, and sensors for every single box. It's safe, but it's too expensive and slow.
This Paper's Method is like building a small, specialized "rearrangement pod" inside the warehouse. You lock the rest of the warehouse in place. You only move the furniture inside this small pod. Because the pod is so small and organized, you can easily measure exactly how much the furniture might wobble (uncertainty) without needing a massive team of inspectors.

How They Built the "Room" (Projections)

The paper tests four different ways to build the walls of this tiny room to see which one works best:

SVD (The "Principal" Method): They look at the original encyclopedia and find the most important directions (like the main aisles) and build the room along those lines. This works very well.
Whitened SVD (The "Data-Aware" Method): They look at the specific questions they want to answer and build the room based on those. This is like customizing the room for a specific type of customer. It often works the best.
DCT (The "Pattern" Method): They use a mathematical pattern (like a wave) to organize the room. It's a bit like using a standard grid. It works okay but isn't as good as the custom methods.
Random Projections (The "Shotgun" Method): They just pick random directions to build the room. This usually fails because the room ends up in a weird spot where the furniture doesn't fit.

The Finding: The "Principal" and "Data-Aware" methods (SVD variants) are the winners. They create a room where the AI can learn efficiently while still knowing when it's unsure.

The Results: Small Size, Big Confidence

The authors tested this on two types of AI: a medium-sized one (RoBERTa) and a giant one (LLaMA-7B).

Calibration: Their method makes the AI much better at knowing when it is unsure. It stops guessing confidently when it's wrong.
Efficiency: They achieved this high level of "confidence awareness" using 5 to 15 times fewer parameters (less memory and storage) than previous Bayesian methods.
Out-of-Distribution Detection: When the AI is asked a question it has never seen before (like a trick question), this method is much better at saying, "I don't know," compared to standard methods that just guess confidently.

The Takeaway

The paper proves that you don't need a massive, expensive system to make an AI humble and aware of its own uncertainty. By shrinking the learning process into a smartly designed, tiny "subspace" (a low-dimensional room), you can get the best of both worlds: fast, cheap training and smart, cautious decision-making.

They showed that even with a very small "rank" (a tiny room), the AI can still capture the complex relationships between its knowledge, provided the room is built in the right direction.

Technical Summary: Bayesian Fine-tuning in Projected Subspaces

Problem Statement

Large Language Models (LLMs) and other deep neural networks often suffer from overconfidence and poor calibration when fine-tuned, particularly on limited data. While Low-Rank Adaptation (LoRA) offers a parameter-efficient solution for fine-tuning by decomposing weight updates into low-rank matrices, standard LoRA lacks mechanisms for uncertainty quantification.

Bayesian variants of LoRA (Bayesian LoRA) have been proposed to address this by learning posterior distributions over parameters. However, these methods face a critical trade-off: modeling the posterior covariance matrix (necessary for epistemic uncertainty) typically requires a number of parameters that scales quadratically with the number of trainable weights. This significantly increases storage and computational costs, partially negating the efficiency gains of LoRA. Furthermore, training these high-dimensional Bayesian models can be unstable and difficult to converge.

Methodology

The authors propose a novel framework that combines parameter-efficient fine-tuning with Bayesian inference within a projected low-dimensional subspace. The core idea is to restrict the learning of Bayesian posteriors to a compressed manifold rather than the full parameter space.

1. Projected Subspace Parameterization

Instead of learning full low-rank updates $\Delta W = AB$ where both $A$ and $B$ are trainable (as in standard LoRA), the proposed method freezes the basis matrices $A$ and $B$ and learns only a small, square "core" matrix $R$ . The weight update is defined as:
$\Delta W = ARB$
Here, $A \in \mathbb{R}^{n \times r}$ and $B \in \mathbb{R}^{r \times m}$ are precomputed projection matrices derived from the pretrained weights $W_0$ , and $R \in \mathbb{R}^{r \times r}$ is the learnable parameter. This reduces the number of trainable parameters from $r(n+m)$ to $r^2$ , creating a highly compressed parameter space.

2. Projection Strategies

The paper evaluates four distinct methods for computing the frozen projection matrices $A$ and $B$ :

Truncated SVD: Aligns projections with the principal directions of the original weight matrix $W_0$ .
Truncated Whitened SVD (W-SVD): Accounts for input data statistics (covariance $\Sigma_{xx}$ ) to minimize activation errors rather than weight reconstruction errors.
2D Discrete Cosine Transform (DCT): Uses a fixed, universal basis of cosine waves, preceded by a deterministic permutation of weights to induce spatial correlation.
Random Projections: Uses Haar-distributed Gaussian orthogonal matrices, making no structural assumptions about the weights.

3. Bayesian Inference in the Subspace

The framework performs Bayesian inference on the core matrices $R$ rather than the full weight updates. Two approximate Bayesian methods are employed:

SWAG (Stochastic Weight Averaging-Gaussian): Exploits the trajectory of Stochastic Gradient Descent (SGD) to construct a Gaussian posterior with a low-rank plus diagonal covariance.
Laplace Approximation: Approximates the posterior around a Maximum A Posteriori (MAP) solution. The authors implement both Diagonal (DIAG) and Kronecker-factored (KRON) approximations for the Hessian/Fisher information matrix to balance expressiveness and memory efficiency.

Crucially, the full covariance of the weight updates $\Delta W$ is never materialized. Instead, the covariance is implicitly defined via the Kronecker structure of the projected subspace, allowing for efficient sampling and uncertainty estimation.

Key Contributions

Parameter-Efficient Bayesian Framework: The paper introduces a method that achieves effective uncertainty quantification in very low-dimensional parameter spaces ( $r^2$ parameters per layer) without sacrificing the efficiency benefits of LoRA.
Generalized Projection Strategies: Unlike prior work that primarily relied on SVD, this work introduces and evaluates a generalized framework including data-aware (Whitened SVD) and structure-agnostic (DCT, Random) projections.
Empirical Validation of Low-Rank Covariances: The authors demonstrate that weight covariances in these projected subspaces exhibit low-rank properties, meaning that even very low-rank covariance approximations (e.g., rank $k=2$ ) are often sufficient to model uncertainty effectively.
Scalability: The approach is validated on both RoBERTa-large and LLaMA2-7B, showing that the method scales to larger models while maintaining computational feasibility.

Experimental Results

The experiments were conducted on GLUE tasks (RTE, MRPC, CoLA, SST-2) using RoBERTa-large and commonsense reasoning tasks (OBQA, ARC-E, ARC-C) using LLaMA2-7B.

Calibration and Uncertainty: The proposed methods (B-LoRA-XS and L-LoRA-XS) consistently outperform standard LoRA and LoRA-XS in Expected Calibration Error (ECE) and Negative Log-Likelihood (NLL). They achieve calibration comparable to or better than full Bayesian LoRA baselines (like LoRA-SWAG) while using 5 to 15 times fewer parameters.
Projection Performance: Whitened SVD (W-SVD) and standard SVD projections consistently yielded the best performance. Data-agnostic methods like Random Projections and DCT performed significantly worse, particularly at lower ranks, highlighting the importance of aligning the subspace with the weight structure.
Covariance Rank: The models remained robust even with very low-rank covariance approximations. For instance, B-LoRA-XS achieved optimal calibration at covariance ranks as low as $k=2$ or $k=5$ , whereas standard Bayesian LoRA methods degraded significantly as rank decreased.
OOD Detection: The Bayesian variants demonstrated superior ability to distinguish In-Distribution (ID) from Out-of-Distribution (OOD) data compared to deterministic baselines, with the proposed subspace methods showing strong separation in predictive entropy.
LLaMA2-7B Results: On the larger model, the compressed variants (e.g., L-LoRA-XS with $r=64$ ) achieved significant calibration improvements with less than one-third of the parameter budget required by the Laplace-LoRA baseline, while maintaining competitive accuracy.

Significance and Claims

The paper claims that effective uncertainty quantification does not require high-dimensional parameter spaces. By projecting weight updates onto a carefully chosen low-dimensional manifold, the authors demonstrate that:

Uncertainty can be modeled efficiently: The correlations between weights can be captured effectively in low-dimensional subspaces, allowing for the use of low-rank covariance matrices that drastically reduce memory and compute costs.
Stability is improved: Operating in a reduced parameter space stabilizes optimization, leading to better convergence and calibration under limited computational budgets.
Synergistic Trade-off: The method offers a synergistic trade-off where predictive performance is maintained or improved, and calibration is significantly enhanced, all while using a fraction of the parameters required by existing Bayesian LoRA methods.

The authors conclude that their framework represents a conceptual innovation in applying Bayesian inference to compressed subspaces, providing a practical solution for uncertainty-aware fine-tuning of large models in resource-constrained settings.