---
title: 'Learnable-β VAE: Adaptive Disentangled Modeling'
url: https://www.emergentmind.com/topics/learnable-beta-vae-l-vae
type: topic
---

# Learnable-β VAE: Adaptive Disentangled Modeling

The Learnable-$\beta$ Variational Autoencoder (L-VAE) is a variational autoencoder framework that jointly learns both the disentangled representation of data and the hyperparameters of its cost function, thereby dynamically controlling the trade-off between disentanglement and reconstruction. L-VAE is an extension of $\beta$-VAE, with the key innovation being the replacement of the hand-tuned $\beta$ trade-off parameter by two learnable positive weights, allowing the model to discover an effective balance between reconstruction fidelity and the structured factorization of the latent space. Experimental results demonstrate that L-VAE achieves state-of-the-art or near–state-of-the-art performance across established disentanglement benchmarks, while preserving the simplicity and effectiveness characteristic of VAE-based objectives [2507.02619].

## 1. Mathematical Foundations

L-VAE modifies the standard $\beta$-VAE loss, which augments the Evidence Lower Bound (ELBO) by a fixed $\beta$ weighting parameter on the KL divergence, as follows:

\[
\mathcal{L}_{\beta\text{-VAE}}(\theta,\phi) =
    - \mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})} \left[\log p_\theta(\mathbf{x}|\mathbf{z})\right]
    + \beta \; D_\mathrm{KL}\left(q_\phi(\mathbf{z}|\mathbf{x}) \parallel p(\mathbf{z})\right)
\]

L-VAE introduces two learnable, positive scales $\{\sigma_0,\sigma_1\}$ and a quadratic regularizer, yielding the loss:

\[
\mathcal{L}_{\mathrm{L\text{-}VAE}}(\theta, \phi, \sigma_0, \sigma_1)
    =
    - \frac{1}{\sigma_0^2} \, \mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})} [\log p_\theta(\mathbf{x}|\mathbf{z})]
    + \frac{1}{\sigma_1^2} \, D_\mathrm{KL}(q_\phi(\mathbf{z}|\mathbf{x}) \,\|\; p(\mathbf{z}))
    + \big(\sigma_0^2 + \sigma_1^2\big)
\]

- The reconstruction loss and KL divergence are weighted by $1/\sigma_0^2$ and $1/\sigma_1^2$ respectively.
- $\sigma_0^2 + \sigma_1^2$ regularizes the weights, preventing unbounded growth and collapse.
- The effective $\beta$ can be expressed as $\widehat\beta = \sigma_0^2 / \sigma_1^2$.

This formulation can be interpreted as introducing adaptive, scale parameters that enable the model to self-tune the balance between reconstruction and disentanglement during training, circumventing the need for manual hyperparameter sweeps.

## 2. Learnable Weight Parameterization and Optimization

The scale parameters $\{\sigma_0,\sigma_1\}$ are parameterized as exponentials to ensure positivity:

\[
\sigma_i = \exp(s_i / 2), \quad s_i \in \mathbb{R},\ i = 0,1
\]

Optimization is performed jointly over the model parameters $(\theta,\phi)$ and the log-scales $(s_0,s_1)$, using Adam. The quadratic penalty $\sigma_0^2 + \sigma_1^2$ constrains both values towards zero, thereby regulating the weighting on both loss terms. Empirical results indicate convergence of $s_0$, $s_1$ (and hence the relative weighting $\widehat\beta$) within the initial 50,000 to 100,000 training iterations; after convergence, the ratio $\sigma_0^2/\sigma_1^2$ remains nearly constant.

The optimization workflow, including parameter updates and loss computation, is as follows:

```python
for minibatch in dataset:
    # Encode input
    mu, logv = encoder(minibatch)
    z = sample_gaussian(mu, logv)
    rec_loss = -sum(log_p_x_given_z(x, z) for x, z in zip(minibatch, z))
    kl_loss = sum(KL_div(mu, logv) for mu, logv in zip(mu, logv))
    sigma0 = exp(s0 / 2)
    sigma1 = exp(s1 / 2)
    total_loss = (1/sigma0**2)*rec_loss + (1/sigma1**2)*kl_loss + (sigma0**2 + sigma1**2)
    # Backpropagate w.r.t. encoder, decoder, and s0, s1
    optimizer.step()
```

## 3. Model Architectures

Two principal neural architectures were implemented, contingent upon the dataset's structure:

**A. Multi-Layer Perceptron (MLP):**  
- *Dataset:* dSprites (64$\times$64$\times$1), $L=5$ latent dimensions  
- *Encoder:* [1200 → 1200 → 2$\times$latent\_dim] (ReLU activations, sigmoid output)  
- *Decoder:* [1200 → 1200 → 64$\times$64] (ReLU activations, sigmoid output)

**B. Convolutional Neural Networks (CNN):**  
- *Datasets:* MPI3D, Falcor3D, Isaac3D (64$\times$64$\times$3), 
    - MPI3D, Falcor3D: $L=7$; Isaac3D: $L=9$
- *Encoder:*
    - Conv(3$\rightarrow$32, $4\times4$, stride=2) $\rightarrow$ ReLU
    - Conv(32$\rightarrow$32, $4\times4$, stride=2) $\rightarrow$ ReLU
    - Conv(32$\rightarrow$64, $4\times4$, stride=2) $\rightarrow$ ReLU
    - Conv(64$\rightarrow$64, $4\times4$, stride=2) $\rightarrow$ ReLU
    - Conv(64$\rightarrow$32, $4\times4$, stride=1) $\rightarrow$ ReLU
    - Flatten $\rightarrow$ FC $\rightarrow$ $2\times$latent\_dim
- *Decoder:* Mirrors the encoder with transposed convolutions and applies sigmoid output.

For CelebA (128$\times$128), the architecture was further deepened, akin to a StyleGAN-inspired configuration, with the L-VAE loss and optimization scheme held constant.

## 4. Experimental Setup and Evaluation Methodology

Extensive experiments were performed on the following datasets:
- dSprites (binary shapes; 5 factors)
- MPI3D-complex (robot arm scenes; 7 factors)
- Falcor3D (living-room renderings; 7 factors)
- Isaac3D (robot arm + kitchen; 9 factors)
- CelebA (faces, 40 binary attributes; qualitative studies only)

Optimization employed Adam (parameters: $\beta_1=0.9$, $\beta_2=0.999$, $\epsilon=10^{-8}$), with a OneCycleLR learning rate schedule ($10^{-5} \rightarrow 10^{-4} \rightarrow 10^{-6}$) and batch sizes of 32, 64, 128, or 256. Training lengths ranged from 300,000 to 1,000,000 iterations, dataset-dependent.

Baselines included VAE ($\beta=1$), $\beta$-VAE ($\beta\in\{2,4\}$), ControlVAE, DynamicVAE, and $\sigma$-VAE, with hyperparameters selected as prescribed in their original works.

Evaluation deployed six disentanglement metrics via the Carbonneau et al. toolbox:
- Modularilty: $\beta$-VAE score, FactorVAE score
- Compactness: MIG, SAP
- Explicitness: Explicitness score
- Holistic: IRS (Interventional Robustness Score)

## 5. Quantitative Results and Comparative Analysis

L-VAE demonstrated consistently strong performance across all disentanglement metrics and baseline comparisons. Table 1 summarizes the reconstruction loss (MSE, lower is better) and $\beta$-VAE score (higher is better) on four key datasets:

| Dataset    | Model        | Recon Loss ↓ | $\beta$-VAE Score ↑ |
|------------|--------------|--------------|---------------------|
| dSprites   | VAE          | 12.46        | 0.91                |
|            | $\beta$-VAE (2) | 25.04    | 0.97                |
|            | ControlVAE   | 28.00        | 0.96                |
|            | DynamicVAE   | 31.57        | 0.93                |
|            | $\sigma$-VAE | 29.21        | 0.85                |
|            | **L-VAE**    | **21.14**    | **0.97**            |
| MPI3D      | VAE          | 11.21        | 0.70                |
|            | $\beta$-VAE (2) | 15.31    | 0.67                |
|            | ControlVAE   | 14.34        | 0.67                |
|            | DynamicVAE   | 15.81        | 0.54                |
|            | $\sigma$-VAE | 5.42         | 0.66                |
|            | **L-VAE**    | **10.79**    | **0.71**            |
| Falcor3D   | VAE          | 215.03       | 0.87                |
|            | $\beta$-VAE (4) | 105.17   | 0.92                |
|            | ControlVAE   | 187.41       | 0.78                |
|            | DynamicVAE   | 216.87       | 0.72                |
|            | $\sigma$-VAE | 78.82        | 0.89                |
|            | **L-VAE**    | **97.97**    | **0.88**            |
| Isaac3D    | VAE          | 13.37        | 0.75                |
|            | $\beta$-VAE (2) | 17.08    | 0.78                |
|            | ControlVAE   | 27.45        | 0.60                |
|            | DynamicVAE   | 35.55        | 0.49                |
|            | $\sigma$-VAE | 23.09        | 0.67                |
|            | **L-VAE**    | **12.97**    | **0.76**            |

Notable trends include:
- L-VAE reduces reconstruction loss substantially relative to ControlVAE and DynamicVAE.
- L-VAE achieves best or second-best disentanglement on every dataset.
- L-VAE ranks first or second across all six disentanglement metrics on all datasets evaluated [2507.02619].

## 6. Qualitative Analysis: Latent Traversals

Qualitative latent traversal experiments on CelebA revealed that sweeping individual latent coordinates in the L-VAE's latent space cleanly and independently manipulates semantically meaningful features. For example, moving one coordinate yielded changes in attributes such as “Smile,” “Bangs,” “Eyeglasses,” “Camera position,” “Receding hairline,” “Gender,” “Mustache,” “Make-up,” and “Apparent age,” while leaving other content unchanged. This result confirms the effectiveness of L-VAE for producing disentangled representations without compromising reconstruction quality.

## 7. Context, Implications, and Summary

L-VAE represents a methodological advance in disentangled generative modeling by incorporating learned weighting coefficients rather than fixed hyperparameters, thereby obviating extensive hyperparameter tuning. The end-to-end trainable structure, quadratic regularization, and empirical demonstration of stable convergence further enhance its applicability. Experimental results suggest that dynamic weighting in the loss offers robust trade-offs across heterogeneous domains, maintaining reconstruction quality while consistently attaining high disentanglement as measured by multiple metrics [2507.02619].

The clarity and stability of the learned $\widehat\beta$ and the empirical convergence of the weighting parameters imply an effective decomposition of disentanglement and fidelity in practice. A plausible implication is that self-tuning weighting in VAE objectives could serve as a general strategy for multi-objective generative modeling in future research directions.

Source: https://www.emergentmind.com/topics/learnable-beta-vae-l-vae