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.
• $\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:

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
• Decoder: 1200 → 1200 → 64$\times$64

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 (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 $\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.