Latent Space Factorization in LoRA

TL;DR

This work tackles the limitation of LoRA by proposing FVAE-LoRA, which embeds a factorized variational autoencoder into the LoRA update to explicitly separate task-salient information ($\mathbf{z}_1$) from residual information ($\mathbf{z}_2$) via a novel ELBO with a repulsion term $\Gamma$. During training, VR encodes inputs into two latent spaces, and only $\mathbf{z}_1$ is used to form the low-rank adaptation, enhancing task relevance and reducing vulnerability to spurious correlations. Empirical results across image, text, and audio tasks show FVAE-LoRA outperforms standard LoRA and approaches or exceeds full fine-tuning budgets, with notable robustness gains under distribution shifts. The work also provides ablations and spurious-correlation analyses, illustrating the practical impact of latent-space factorization for robust, parameter-efficient model adaptation.

Abstract

Low-rank adaptation (LoRA) is a widely used method for parameter-efficient finetuning. However, existing LoRA variants lack mechanisms to explicitly disambiguate task-relevant information within the learned low-rank subspace, potentially limiting downstream performance. We propose Factorized Variational Autoencoder LoRA (FVAE-LoRA), which leverages a VAE to learn two distinct latent spaces. Our novel Evidence Lower Bound formulation explicitly promotes factorization between the latent spaces, dedicating one latent space to task-salient features and the other to residual information. Extensive experiments on text, audio, and image tasks demonstrate that FVAE-LoRA consistently outperforms standard LoRA. Moreover, spurious correlation evaluations confirm that FVAE-LoRA better isolates task-relevant signals, leading to improved robustness under distribution shifts. Our code is publicly available at: https://github.com/idiap/FVAE-LoRA

Figures (2)

• Figure 1: Comparison between LoRA and the proposed FVAE-LoRA. During training, FVAE-LoRA factorizes the latent space into two components, $\mathbf{z}_1$ and $\mathbf{z}_2$, where only the task-salient latent factor $\mathbf{z}_1$ is propagated downstream. At inference, only the encoder corresponding to $\mathbf{z}_1$ is required.
• Figure 2: Random samples drawn from the train split of the Animals dataset, illustrating an original image ($\mathbf{x}$), its core object features ($\mathbf{x}_\text{core}$), and its spurious background features ($\mathbf{x}_\text{spurious}$).