Learnable-β VAE: Adaptive Disentangled Modeling

• The paper introduces a novel self-tuning VAE that replaces the fixed β parameter with learnable weights to balance reconstruction and disentanglement.
• It employs an adaptive loss function where reconstruction loss and KL divergence are weighted by positive scale parameters optimized via Adam with quadratic regularization for stability.
• Experimental results on benchmarks like dSprites, MPI3D, and CelebA demonstrate that L-VAE achieves state-of-the-art performance in disentangled representation learning.

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 (Ozcan et al., 3 Jul 2025).

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.
• $\beta$0 regularizes the weights, preventing unbounded growth and collapse.
• The effective $\beta$1 can be expressed as $\beta$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 $\beta$3 are parameterized as exponentials to ensure positivity:

$\beta$4

Optimization is performed jointly over the model parameters $\beta$5 and the log-scales $\beta$6, using Adam. The quadratic penalty $\beta$7 constrains both values towards zero, thereby regulating the weighting on both loss terms. Empirical results indicate convergence of $\beta$8, $\beta$9 (and hence the relative weighting $\beta$0) within the initial 50,000 to 100,000 training iterations; after convergence, the ratio $\beta$1 remains nearly constant.

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

$$\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)$$0

3. Model Architectures

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

A. Multi-Layer Perceptron (MLP):

• Dataset: dSprites (64$\beta$264$\beta$31), $\beta$4 latent dimensions
• Encoder: 1200 → 1200 → 2$\beta$5latent_dim
• Decoder: 1200 → 1200 → 64$\beta$664

B. Convolutional Neural Networks (CNN):

• Datasets: MPI3D, Falcor3D, Isaac3D (64$\beta$764$\beta$83),
  • MPI3D, Falcor3D: $\beta$9; Isaac3D: $\beta$0
• Encoder:
  • Conv(3$\beta$132, $\beta$2, stride=2) $\beta$3 ReLU
  • Conv(32$\beta$432, $\beta$5, stride=2) $\beta$6 ReLU
  • Conv(32$\beta$764, $\beta$8, stride=2) $\beta$9 ReLU
  • Conv(64$\beta$064, $\beta$1, stride=2) $\beta$2 ReLU
  • Conv(64$\beta$332, $\beta$4, stride=1) $\beta$5 ReLU
  • Flatten $\beta$6 FC $\beta$7 $\beta$8latent_dim
• Decoder: Mirrors the encoder with transposed convolutions and applies sigmoid output.

For CelebA (128$\beta$9128), 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: $$\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)$$0, $$\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)$$1, $$\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)$$2), with a OneCycleLR learning rate schedule ($$\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)$$3) 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 ($$\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)$$4), $$\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)$$5-VAE ($$\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)$$6), ControlVAE, DynamicVAE, and $$\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)$$7-VAE, with hyperparameters selected as prescribed in their original works.

Evaluation deployed six disentanglement metrics via the Carbonneau et al. toolbox:

• Modularilty: $$\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)$$8-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 $$\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)$$9-VAE score (higher is better) on four key datasets:

Dataset	Model	Recon Loss ↓	$\{\sigma_0,\sigma_1\}$0-VAE Score ↑
dSprites	VAE	12.46	0.91
	$\{\sigma_0,\sigma_1\}$1-VAE (2)	25.04	0.97
	ControlVAE	28.00	0.96
	DynamicVAE	31.57	0.93
	$\{\sigma_0,\sigma_1\}$2-VAE	29.21	0.85
	L-VAE	21.14	0.97
MPI3D	VAE	11.21	0.70
	$\{\sigma_0,\sigma_1\}$3-VAE (2)	15.31	0.67
	ControlVAE	14.34	0.67
	DynamicVAE	15.81	0.54
	$\{\sigma_0,\sigma_1\}$4-VAE	5.42	0.66
	L-VAE	10.79	0.71
Falcor3D	VAE	215.03	0.87
	$\{\sigma_0,\sigma_1\}$5-VAE (4)	105.17	0.92
	ControlVAE	187.41	0.78
	DynamicVAE	216.87	0.72
	$\{\sigma_0,\sigma_1\}$6-VAE	78.82	0.89
	L-VAE	97.97	0.88
Isaac3D	VAE	13.37	0.75
	$\{\sigma_0,\sigma_1\}$7-VAE (2)	17.08	0.78
	ControlVAE	27.45	0.60
	DynamicVAE	35.55	0.49
	$\{\sigma_0,\sigma_1\}$8-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 (Ozcan et al., 3 Jul 2025).

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 (Ozcan et al., 3 Jul 2025).

The clarity and stability of the learned $\{\sigma_0,\sigma_1\}$9 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.