Latent Space Factorization in LoRA

AI summary

Overview

Research area: Parameter-efficient fine-tuning (PEFT) of large foundation models, specifically low-rank adaptation (LoRA), and latent-variable modeling with variational autoencoders (VAEs).

Technical level: Intermediate. The high-level idea is intuitive, but the paper includes a formal ELBO derivation, a Wasserstein-based bound on a discrepancy term, and KL-decomposition analysis in the appendices.

Scope: The paper proposes FVAE-LoRA, a LoRA variant that adds a factorized VAE with two latent spaces so that one latent space carries task-salient information and the other absorbs residual information, and evaluates it across image, text, and audio tasks.

What This Paper Is About

Standard LoRA adapts a frozen weight matrix by learning low-rank matrices A and B, but nothing in that mechanism explicitly forces the learned low-rank subspace to hold information that is actually relevant to the downstream task. The authors argue this implicit learning through gradient descent can leave irrelevant or even detrimental pre-training information mixed into the update, which may explain why LoRA sometimes underperforms full fine-tuning. Their goal is to add explicit control over what the low-rank update encodes.

Key Contributions

1. A new PEFT method (FVAE-LoRA): A VAE with two factorized latent spaces, z₁ and z₂, is integrated into the LoRA framework so that task-salient information (z₁) is disentangled from residual information (z₂).
2. A factorizing ELBO formulation: A novel Evidence Lower Bound is introduced that enforces factorization between the two latent spaces during training, including a repulsive regularizer term denoted Γ.
3. Strong empirical performance: Experiments on image, text, and audio benchmarks show FVAE-LoRA consistently outperforming standard LoRA, with comparisons also against full fine-tuning and other LoRA variants.
4. Empirical validation of robustness: Targeted spurious-correlation experiments on Waterbirds, CelebA, and Animals are used to argue that z₁ captures task-critical information, yielding better worst-group accuracy.

Main Findings

• Image classification: With a ViT-B/16 backbone on DTD, EuroSAT, GTSRB, RESISC45, SUN397, and SVHN, FVAE-LoRA reaches an average top-1 accuracy of 89.53%, above LoRA (87.39%), DoRA (88.32%), PiSSA (87.04%), OLoRA (86.63%), rsLoRA (85.98%), and slightly above full fine-tuning (89.38%). It uses 0.7311% of parameters versus 0.7240% for LoRA and 0.7451% for DoRA. Individual FVAE-LoRA scores are DTD 78.19, EuroSAT 97.78, GTSRB 97.98, RESISC45 93.57, SUN397 73.14, and SVHN 96.55; full fine-tuning leads on EuroSAT (98.30) and GTSRB (98.85).
• Commonsense reasoning: On Llama-3-8B across PIQA, SIQA, ARC-c, ARC-e, OBQA, HellaSwag, and WinoGrande, FVAE-LoRA averages 87.82% accuracy at 0.0850% parameters, ahead of HiRA (87.40% at 0.0848%), LoRA (77.82% at 0.0848%), ChatGPT (77.57%), Prompt Tuning (34.12%), and P-Tuning (12.13%).
• GLUE benchmark: On roberta-base over SST2, CoLA, QNLI, MRPC, RTE, STSB, and WNLI, FVAE-LoRA achieves the highest average (81.21), ahead of full fine-tuning (80.67) and LoRA (79.81), at 0.4759% parameters versus 0.4710% for LoRA. Its WNLI score of 64.33 stands out against LoRA's 57.73.
• Speech recognition: On TIMIT with Wav2Vec2-Large and CTC loss, FVAE-LoRA reports a phoneme error rate of 8.09, better than LoRA's 9.38 and close to full fine-tuning's 7.48, at 0.4999% parameters versus 0.4961% for LoRA.
• Spurious correlation robustness: FVAE-LoRA raises worst-group accuracy over LoRA on Animals (62.0 vs 54.79) and CelebA (43.33 vs 40.00), and slightly on Waterbirds (75.85 vs 75.49), while average accuracy remains comparable (Animals 89.55 vs 88.20; Waterbirds 90.99 vs 90.39; CelebA 95.77 vs 96.09). The reported accuracy disparity drops from 34.8 for LoRA to 31.71 for FVAE-LoRA.
• Ablation: Replacing the proposed objective with a two-latent-variable VAE (VAE2LAT) or with its β-VAE version (β = 10) reduces the six-dataset image average to 86.43 and 87.29 respectively, versus 89.53 for FVAE-LoRA.

Methodology in Plain English

The researchers keep the frozen base weights of a model and, for each targeted linear layer (the query and key matrices in the transformer models), train a small VAE alongside the downstream task. Unlike a normal VAE with one latent code, this one has two: z₁ and z₂. The input to the layer is fed into the VAE and reconstructed, so the latent codes must retain enough information to rebuild that input. At the same time, only z₁ is passed through a learned matrix B and added to the frozen layer's output, giving Wx + Bz₁. Because the downstream task loss only ever sees z₁, that latent space is pushed to hold task-relevant content, while z₂ takes on whatever else is needed for accurate reconstruction.

To keep the two spaces from collapsing into the same information, the objective adds a repulsive regularizer (Γ) that penalizes similarity between the second encoder's output and the first latent space's prior. The authors decompose Γ into a "mismatch" term, which is a difference of KL divergences and encourages the second encoder to align with its own prior rather than the first one's, and a "discrepancy" term, which they bound using the 2-Wasserstein distance between the two posteriors. Practically, both encoders are diagonal Gaussians parameterized by neural networks, the prior p₁ is a standard normal, and p₂ is centered at 1.5 so the two priors occupy distinct locations. Balancing coefficients α, β, and δ weight the reconstruction, KL, and repulsion terms. At inference, only the encoder for z₁ is needed. The rank r is fixed at 16 for all LoRA-based methods and also serves as the dimensionality of z₁.

Why This Matters

Impact on research: The paper reframes LoRA improvement as a question of what information the low-rank update carries, rather than only how the update is structured or optimized — distinguishing it from variants like AdaLoRA, DoRA, LoRA+, PiSSA, and rsLoRA. It also connects PEFT to the VAE disentanglement literature (β-VAE, FHVAE, FactorVAE, TCVAE, DIP-VAE) with a different goal: task-conditional separation rather than disentangling underlying data factors. Code is publicly released at https://github.com/idiap/FVAE-LoRA.

Real-world applications:

• Adapting vision backbones for specialized classification such as satellite imagery, traffic signs, remote sensing scenes, and texture recognition.
• Instruction or domain adaptation of large language models for reasoning-heavy assistants under tight memory budgets.
• Adapting speech recognition models to new languages, accents, or phonetic corpora with limited compute.
• Deployments where spurious cues in training data (background, demographics) could otherwise cause failures on minority groups.

Industry relevance: The parameter budgets reported (roughly 0.73% for image models, about 0.085% for Llama-3-8B, roughly 0.48% for roberta-base, and about 0.50% for Wav2Vec2-Large) mean FVAE-LoRA sits close to LoRA's memory and storage footprint, which matters for serving many adapted models on shared hardware. Its worst-group accuracy improvements are directly relevant to fairness-sensitive deployments.

Future Directions

• The paper notes the method is generic and could give latent-space control to any explicit LoRA method; combining it with variants such as DoRA, PiSSA, or quantized LoRA is an open avenue.
• The work gives a practical guide for selecting β and δ in an appendix, but how these coefficients should be chosen across new modalities and architectures remains an empirical question.
• Extending the evaluation to further domains, larger backbones, and higher ranks than the fixed r = 16 used here is a natural next step.
• The spurious-correlation evidence covers Waterbirds, CelebA, and Animals; whether the robustness gains hold under other distribution shifts, and whether the Γ-based repulsion can be tightened theoretically beyond the Wasserstein bound, are open questions.

Target Audience

Researchers and engineers working on parameter-efficient fine-tuning of foundation models, particularly those interested in LoRA variants, multi-modal adaptation, and robustness to spurious correlations. It is also relevant to practitioners in computer vision, NLP, and speech who need to adapt large pre-trained models under constrained memory and compute budgets, and to readers interested in applying VAE-style latent factorization to a non-generative goal.

Authors’ 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