Source-linked AI summary
Understanding disentangling in $β$-VAE
Christopher P. Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, Alexander Lerchner
TL;DR
The paper investigates how β-VAE can produce representations aligned with generative factors while addressing the tension between disentanglement and reconstruction. It develops a rate-distortion-based account and introduces controlled capacity increases during training, yielding robust disentangling with high-quality reconstructions.
Problem
Task-trained representations can overfit their training distribution, motivating better understanding of how disentangled representations support knowledge transfer and generalisation.
Method
The paper analyses β-VAE from a rate-distortion perspective and progressively increases the latent encoding capacity by controlling the target KL divergence during training.
Results
Controlled capacity increase produces robust disentangling of the generative factors alongside high-quality reconstructions.
Takeaways & Limitations
Gradually increasing latent capacity enables progressively more factors to be represented while retaining disentanglement in previously learned factors.
Takeaways & Limitations
The VAE formulation assumes Gaussian prior and posterior distributions with diagonal covariance, typically using an isotropic unit-Gaussian prior.
Abstract
from arXiv · showhide
We present new intuitions and theoretical assessments of the emergence of disentangled representation in variational autoencoders. Taking a rate-distortion theory perspective, we show the circumstances under which representations aligned with the underlying generative factors of variation of data emerge when optimising the modified ELBO bound in $β$-VAE, as training progresses. From these insights, we propose a modification to the training regime of $β$-VAE, that progressively increases the information capacity of the latent code during training. This modification facilitates the robust learning of disentangled representations in $β$-VAE, without the previous trade-off in reconstruction accuracy.
1 Introduction
Representation learning strongly influences machine-learning performance, yet task-trained representations can overfit and generalise poorly. The section motivates unsupervised disentangled representations, including β-VAE, as a potentially more transferable alternative.
- Representation learning is central to machine learning because algorithm performance depends critically on input representations.
- Modern deep-learning representations can overfit training tasks and lack properties needed for knowledge transfer and out-of-distribution generalisation.
- Task-agnostic unsupervised learning is proposed as a way to learn features useful across varied tasks, with disentangled representations identified as a promising approach.
- Disentangled representations assign separate latent units to individual generative factors while remaining relatively invariant to other factors.
- β-VAE is an unsupervised visual disentangled-representation model that constrains latent capacity and encourages factorisation.
- The proposed training extension relaxes the information bottleneck during training to improve disentangling robustness and reconstruction accuracy.
2 Variational Autoencoder (VAE)
The VAE learns a data distribution through encoder and decoder parameterisations, optimising a tractable lower bound. Gaussian diagonal-covariance assumptions and reparameterisation make gradient estimation practical.
- The VAE aims to learn the marginal likelihood of data generated from latent factors.
- The encoder and decoder distributions are parameterised by φ and θ, respectively.
- Maximising the variational objective is equivalent to maximising a lower bound on the true objective because the KL divergence is non-negative.
- The VAE commonly uses diagonal-covariance Gaussian prior and posterior distributions, with the prior set to an isotropic unit Gaussian N(0, 1).
- These Gaussian parameterisations enable the reparameterisation trick for estimating gradients with respect to encoder parameters.
3 β-VAE
β-VAE modifies the VAE objective with a tunable β that pressures the posterior toward a factorised unit-Gaussian prior. This can improve disentanglement, but stronger values commonly reduce reconstruction fidelity.
- β-VAE introduces an adjustable hyperparameter β into the original VAE objective.
- Values of β usually greater than one produce more disentangled latent representations, while β = 1 recovers the original VAE.
- The β-weighted pressure for the posterior to match a factorised unit-Gaussian prior constrains latent bottleneck capacity while retaining reconstruction information.
- Higher β values needed for disentangling often create a trade-off with reconstruction fidelity.
4 Understanding disentangling in β-VAE
β-VAE treats the latent posterior as an information bottleneck whose capacity and factorisation pressures can align latent dimensions with distinct generative factors. Progressively increasing capacity reveals factors in stages, improving disentanglement while recovering detailed reconstructions.
- Information bottleneck: The β-VAE objective jointly maximises reconstruction log likelihood and minimises a β-weighted KL divergence that limits information transmitted through the latent bottleneck.The KL term is zero when posterior channels match the unit Gaussian prior, corresponding to zero latent capacity.
- Information bottleneck: Reducing posterior spread or broadening posterior variance increases overlap, encouraging neighbouring data points to remain close in latent space and reducing reconstruction confusion.This locality pressure follows because overlapping posteriors can confuse samples from nearby data points.
- Factor alignment: Compared with a standard VAE, β-VAE represents Gaussian-blob position factors in two latent dimensions rather than four, illustrating more factorised traversals.Figure 2 compares β = 1 and β = 150 models using reconstructions, latent traversals, and position-tuning heatmaps.
- Factor alignment: β-VAE can align latent dimensions with generative factors because distinct factors make different contributions to reconstruction and diagonal covariance requires different capacities to occupy separate dimensions.A separate latent axis also provides a smooth way to preserve locality for each new factor without affecting other axes.
5 Improving disentangling in β-VAE with controlled capacity increase
The controlled-capacity objective gradually increases latent encoding capacity, producing robustly disentangled representations and high-quality reconstructions on dSprites and Chairs.
- The modified β-VAE objective controls latent encoding capacity through a gradually increased target KL divergence.Capacity C is increased from zero toward a value sufficient for good reconstructions.
- On coloured dSprites, controlled capacity increase demonstrates robust disentangling of all dataset factors with high-quality reconstructions.The model separates position, size, shape, rotation, and colour across latent traversals.
- Latent traversals isolate Y position, X position, scale, shape, rotation, and colour changes across successive latent dimensions.The final two dimensions have effectively zero KL divergence and produce no output effect.
- The dSprites reconstructions are difficult to distinguish by eye from their corresponding data samples.Traversal images are also reported as high quality, contrasting with earlier fixed β-modulated KL results.
- On the 3D Chairs dataset, reconstructions are high quality and traversals produce smooth changes with plausible chairs.Individual traversals appear isolated to qualitative features such as viewing angle, size, and leg or back styles.
6 Conclusion
The paper explains why β-VAE can learn axis-aligned disentangled representations and proposes controlled capacity increase as a training modification. It reports more robust disentangling together with better reconstruction fidelity than the original formulation.
- The paper identifies pressures that align representational axes with generative factors while preserving locality in the data.
- It proposes gradually increasing latent posterior encoding capacity by raising average KL divergence from zero during training.This replaces the fixed β-weighted KL term in the original β-VAE objective.
- The modified training regime promotes robust disentangled representations combined with better reconstruction fidelity than the original β-VAE formulation.
A.1 Model Architecture
The experiments use a shared convolutional and fully connected architecture, with Gaussian latent variables parameterized by encoder outputs.
- The encoder uses four convolutional layers followed by two fully connected layers.Each convolutional layer has 32 channels, 4x4 kernels, and stride 2; each fully connected layer has 256 units.
- The latent distribution uses a fully connected layer to parameterize the mean and log standard deviation of Gaussian random variables.The architecture uses 10 Gaussian variables, or 32 for the CelebA experiment.
A.2 Training Details
Training uses a large capacity-control parameter so the actual KL remains close to the target, with dataset-specific linear capacity schedules.
- The capacity-control parameter was set to 1000 to keep the actual KL close to the target KL C.
- For dSprites, C increased linearly from 0 to 25 nats over 100,000 training iterations.
- For CelebA, the target capacity increased to 50 nats.