---
title: 'FVAE-LoRA: Fine-Tuning via Latent Factorization'
url: https://www.emergentmind.com/topics/fvae-lora
type: topic
---

# FVAE-LoRA: Fine-Tuning via Latent Factorization

FVAE-LoRA is a parameter-efficient fine-tuning (PEFT) framework that augments the standard Low-Rank Adaptation (LoRA) method with an explicit latent-space factorization. By embedding a small variational autoencoder (VAE) at adaptation layers, FVAE-LoRA learns to disentangle task-salient from residual features in the low-rank update space, leading to increased robustness and downstream generalization, especially under spurious correlations and distribution shifts. Unlike standard LoRA variants, FVAE-LoRA utilizes a tailored Evidence Lower Bound (ELBO) objective to enforce this factorization without increasing inference-time parameter budgets [2510.19640].

## 1. Architectural Design and Latent Factorization

FVAE-LoRA operates by replacing the static LoRA low-rank update mechanism with an encoder-decoder (VAE) pipeline. Given a frozen linear transformation $W \in \mathbb{R}^{k \times d}$ and input activation $x \in \mathbb{R}^{d}$, the standard LoRA approach augments with $W + BA$ for rank $r$ factors $A \in \mathbb{R}^{r \times d}$ and $B \in \mathbb{R}^{k \times r}$. FVAE-LoRA introduces two parallel encoders $q_{\phi_1}, q_{\phi_2}$, each mapping $x$ to latent codes $z_1,z_2 \in \mathbb{R}^r$:

- $z_1$ parameterizes the low-rank update through $B z_1$, driving the downstream task.
- $z_2$ captures residual or non-task-salient information.

Each encoder is a two-layer MLP producing diagonal-Gaussian posteriors: $q_{\phi_i}(z_i|x) = \mathcal{N}(\mu_i(x), \mathrm{diag}(\sigma^2_i(x)))$. The decoder $p_\theta(x|z_1,z_2)$ is a shallow MLP reconstructing $\hat x$ from $[z_1;z_2] \in \mathbb{R}^{2r}$. The prior for $z_1$ is $p_1(z_1) = \mathcal{N}(0,I)$ (enforcing compact, task-relevant codes), and for $z_2$ it is $p_2(z_2) = \mathcal{N}(1.5 \cdot \mathbf{1}, I)$ (expanding residual space). A schematic (for one linear layer) is as follows:

```
   x
   │
   ├── q_{φ₁} → z₁ ─┐
   │                ├──⁺ B z₁
   ├── q_{φ₂} → z₂ ─┘
 + W x → output y
   └── p_θ([z₁;z₂]) →  x̂  (reconstruction)
```

Only $z_1$ is used for forming output activations; $z_2$ modulates the VAE loss and factorization but has no downstream effect during inference.

## 2. Training Objective and Information Factorization

The training objective couples the downstream task loss with a custom VAE-derived ELBO that explicitly encourages separation of task-driven and residual subspaces. For a single activation $x$ at an adapted layer, define

\[
\mathcal{L}^{\text{FVAE}}(x) = \alpha \, \mathbb{E}_{z_1,z_2} [\log p_\theta(x|z_1,z_2)] - \beta\, \mathrm{KL}(q_{\phi_1}(z_1|x)\,\|\,p_1(z_1)) + \delta \, \Gamma(q_{\phi_1}, q_{\phi_2}),
\]

with the repulsive cross-prior term

\[
\Gamma = \mathbb{E}_{z_1 \sim q_{\phi_1},z_2 \sim q_{\phi_2}} [ \log p_2(z_2) - \log p_1(z_1) ].
\]

The coefficients $\alpha, \beta, \delta$ tune (respectively) reconstruction-vs-adaptation, information bottleneck on $z_1$, and the repulsive force between the two latents. For a dataset $\{x^{(n)}\}$, the total loss is

\[
\min_{\phi, \theta} \sum_n [ \mathcal{L}_{\text{task}}(x^{(n)}) + \lambda \sum_l \mathcal{L}^{\text{FVAE}}(x_l^{(n)}) ],
\]

where $\lambda \approx 10^{-3}$ (image tasks) to $10^{-4}$ (text/audio) balances FVAE regularization with task loss. 

The key innovation is the $\Gamma$ term, which, through choice of nonoverlapping priors, creates a geometric “repulsion” between $q_{\phi_1}$ and $q_{\phi_2}$, improving disentanglement. Theoretically, $\Gamma$ decomposes into a KL divergence difference, related to a lower bound on the Wasserstein distance $W_2(q_{\phi_1}, q_{\phi_2})$.

## 3. Optimization Protocol and Algorithmic Implementation

Training proceeds using AdamW with linear warm-up and decay. Hyperparameters include $r=16$ for all reported experiments; $\alpha=1$, $\delta=1$, and $\beta\in[1,10]$ to balance capacity and bottleneck; learning rates range from $5 \times 10^{-3}$ (ViT) to $3 \times 10^{-4}$ (Llama/FVAE-LoRA) and $1 \times 10^{-3}$ (LoRA baselines). For each batch, activations at adapted layers are encoded, sampled, decoded, and the FVAE loss is accumulated. The LoRA update is performed only via $z_1$. Notably, $\beta$-VAE annealing is unnecessary due to the effect of the $\delta$ term. 

Computation adds a ≈30% training-time overhead (due to the decoder path), but inference remains efficient, requiring only the $z_1$ encoder. Regularization is task-dependent, with weight decay (0.01 for vision, 0.0 for LLMs) applied as appropriate.

## 4. Comparative Analysis with LoRA Variants

Standard LoRA treats the low-rank update $A x$ as a direct, fully trainable projection with no explicit separation of task-relevant and nuisance factors. Other PEFT methods (e.g., AdaLoRA, DoRA, rsLoRA) focus on adapting rank, magnitude, or singular value structure but do not address semantic disentanglement within the update. FVAE-LoRA is unique in learning two complementary latent spaces and promoting explicit geometric separation using VAEs and nontrivial priors, with only $z_1$ driving downstream adaptation.

This semantic focus yields improved robustness to spurious correlations, as $z_2$ absorbs signals not essential for the core task. Unlike some adapters, FVAE-LoRA modules cannot be merged back into $W$, so the method remains a dynamic, input-dependent module at inference for enhanced robustness.

## 5. Empirical Performance Across Modalities

FVAE-LoRA was evaluated across vision, language, and audio domains. In all cases $r=16$ was used.

**Image Classification (ViT-B/16 backbone, q/k LoRA):**

| Method     | Params % | DTD   | EuroSAT | GTSRB | RESISC45 | SUN397 | SVHN  | Avg    |
|------------|----------|-------|---------|-------|----------|--------|-------|--------|
| Full-FT    | –        | 78.12 | 98.30   | 98.85 | 94.35    | 69.34  | 97.34 | 89.38  |
| LoRA       | 0.72     | 74.65 | 97.28   | 96.95 | 90.11    | 71.11  | 94.22 | 87.39  |
| FVAE-LoRA  | 0.73     | 78.19 | 97.78   | 97.98 | 93.57    | 73.14  | 96.55 | 89.53  |

**Language (Llama-3-8B, commonsense reasoning):**

| Method     | Params % | PIQA  | SIQA  | ARC-c | ARC-e | OBQA  | HellaSwag | WinoGrande | Avg    |
|------------|----------|-------|-------|-------|-------|-------|-----------|------------|--------|
| LoRA       | 0.085    | 80.74 | 75.59 | 67.58 | 82.11 | 75.20 | 85.73     | 77.82      | 77.82  |
| HiRA       | 0.085    | 88.63 | 80.40 | 81.66 | 93.56 | 87.20 | 94.48     | 85.87      | 87.40  |
| FVAE-LoRA  | 0.085    | 88.96 | 81.58 | 81.06 | 92.72 | 86.20 | 95.30     | 88.95      | 87.82  |

**GLUE (RoBERTa-base):**

| Method     | Params % | SST2  | CoLA   | QNLI  | MRPC  | RTE   | STSB  | WNLI  | Avg    |
|------------|----------|-------|--------|-------|-------|-------|-------|-------|--------|
| Full-FT    | –        | 94.77 | 62.43  | 91.97 | 89.40 | 79.53 | 90.30 | 56.30 | 80.67  |
| LoRA       | 0.47     | 93.97 | 59.60  | 91.87 | 88.73 | 77.87 | 88.90 | 57.73 | 79.81  |
| FVAE-LoRA  | 0.48     | 94.10 | 60.37  | 91.63 | 89.53 | 79.90 | 88.60 | 64.33 | 81.21  |

**Audio (TIMIT, Wav2Vec2-Large):**

| Method     | Params % | PER ↓  |
|------------|----------|--------|
| Full-FT    | –        | 7.48   |
| LoRA       | 0.50     | 9.38   |
| FVAE-LoRA  | 0.50     | 8.09   |

In all domains, FVAE-LoRA matches or outperforms standard LoRA and, in some cases, full fine-tuning, despite training <1% of total parameters.

## 6. Robustness to Spurious Correlations and Worst-Group Generalization

Robustness was evaluated on benchmarks with train–test label/background decoupling: Animals (4 classes × 2 backgrounds), Waterbirds (land vs water), and CelebA (hair color vs gender). Metrics include worst-group (WG) accuracy, average (AVG) accuracy, and disparity $|\text{WG}-\text{AVG}|$.

| Method     | Params % | ANIM.WG | ANIM.AVG | WATR.WG | WATR.AVG | CEL.WG | CEL.AVG | Disparity |
|------------|----------|---------|----------|---------|----------|--------|---------|-----------|
| LoRA       | 0.72     | 54.8    | 88.2     | 75.5    | 90.4     | 40.0   | 96.1    | 34.8      |
| FVAE-LoRA  | 0.73     | 62.0    | 89.6     | 75.8    | 91.0     | 43.3   | 95.8    | 31.7      |

By channeling causal features to $z_1$ and spurious/environmental attributes to $z_2$, FVAE-LoRA improves worst-group performance (notably +7% on Animals) and reduces performance disparity.

## 7. Practical Considerations, Limitations, and Future Directions

Recommended settings include adopting $r=16$ for the low-rank subspace, specific Gaussian priors ($p_1 = \mathcal{N}(0,I)$, $p_2 = \mathcal{N}(1.5,I)$), $\alpha=1$, $\delta=1$, and tuning $\beta$ to trade between over/underfitting ($\beta \ll 1$ overfits; $\beta \gg 10$ underfits). The FVAE loss is weighted at $\lambda \approx 10^{-3}$ (image) or $10^{-4}$ (text/audio) relative to the primary task. Training learning rates and regularization mirror modality best practices; AdamW is recommended.

Limitations include current restriction to attention q/k projections. Extending factorized adaptation to feed-forward networks or value matrices may be beneficial. The FVAE-LoRA adapter cannot be statically merged with $W$, so inference modules remain dynamic and input-dependent. Training time increases by ≈30%, but inference cost remains similar to LoRA due to the omission of decoder usage. Open questions involve further examining generalization under broader distribution shifts and adaptation paradigms.

FVAE-LoRA offers a framework in which a minimal two-latent VAE regularizes low-rank adapters for explicit semantic disambiguation. The resulting models exhibit improved robustness, reduced sensitivity to spurious correlations, and strong empirical performance across heterogeneous tasks and modalities [2510.19640].

Source: https://www.emergentmind.com/topics/fvae-lora