Source-linked AI summary
Isolating Sources of Disentanglement in Variational Autoencoders
Ricky T. Q. Chen, Xuechen Li, Roger Grosse, David Duvenaud
TL;DR
The paper asks how ELBO terms contribute to unsupervised disentanglement and how such representations should be evaluated. It decomposes the ELBO, proposes β-TCVAE and the classifier-free MIG metric, and reports that total correlation is strongly correlated with disentanglement in models trained with its framework.
Problem
Unsupervised disentanglement is difficult, while existing classifier-based metrics are ad hoc and sensitive to hyperparameters.
Method
The paper decomposes the ELBO, trains β-TCVAE with weighted minibatch sampling, and introduces the classifier-free mutual information gap metric.
Results
Total correlation is strongly correlated with disentanglement in models trained with the proposed framework, while β-TCVAE consistently achieves higher MIG scores than β-VAE and InfoGAN and matches FactorVAE.
Takeaways & Limitations
β-TCVAE provides a no-extra-hyperparameter refinement of β-VAE, and MIG offers classifier-free evaluation applicable to varied latent distributions.
Takeaways & Limitations
The minibatch estimator is biased because its expectation is a lower bound, and MIG assumes the underlying factors and empirical generating process are known for evaluation.
Abstract
from arXiv · showhide
We decompose the evidence lower bound to show the existence of a term measuring the total correlation between latent variables. We use this to motivate our $β$-TCVAE (Total Correlation Variational Autoencoder), a refinement of the state-of-the-art $β$-VAE objective for learning disentangled representations, requiring no additional hyperparameters during training. We further propose a principled classifier-free measure of disentanglement called the mutual information gap (MIG). We perform extensive quantitative and qualitative experiments, in both restricted and non-restricted settings, and show a strong relation between total correlation and disentanglement, when the latent variables model is trained using our framework.
1 Introduction
The paper addresses unsupervised disentanglement by decomposing the ELBO, introducing β-TCVAE, and proposing a classifier-free disentanglement metric.
- Unsupervised learning of disentangled representations remains a difficult open problem.
- The paper decomposes the variational lower bound to explain β-VAE’s success in learning disentangled representations.
- β-TCVAE is a β-VAE plug-in replacement trained with weighted minibatches and no additional hyperparameters.
- Empirical evaluations suggest β-TCVAE learns more interpretable representations than existing methods and is fairly robust to random initialization.
- The paper introduces a classifier-free, information-theoretic disentanglement metric generalizable to arbitrarily distributed and non-scalar latent variables.
2 Background: Learning and Evaluating Disentangled Representations
The background frames disentanglement as learning interpretable latent factors and identifies limitations in both β-VAE explanations and classifier-based evaluation metrics.
- Learning Disentangled Representations: A VAE optimizes a tractable evidence lower bound instead of directly maximizing the intractable marginal log-likelihood.
- Learning Disentangled Representations: β-VAE uses a heavily penalized objective with β > 1 to learn disentangled representations, but its mechanism is not explicit.
- Learning Disentangled Representations: InfoGAN encourages interpretable latent representations by maximizing mutual information between observations and selected latent variables.
- Evaluating Disentangled Representations: Classifier-based metrics are ad hoc and hyperparameter-sensitive, motivating direct mutual-information measurement for arbitrary latent distributions.
3 Sources of Disentanglement in the ELBO
The ELBO decomposition isolates mutual-information, total-correlation, and dimension-wise KL terms, motivating β-TCVAE’s weighted minibatch training without extra hyperparameters.
- Sources of Disentanglement in the ELBO: The refined ELBO decomposition contains terms corresponding to mutual information and latent-variable dependence.
- Sources of Disentanglement in the ELBO: Total correlation measures dependence among latent variables, and penalizing it encourages statistically independent factors in the data distribution.
- Sources of Disentanglement in the ELBO: The dimension-wise KL term limits deviations of individual latent dimensions from their priors and acts as a complexity penalty on the aggregate posterior.
- Training with Minibatch-Weighted Sampling: Minibatch-weighted sampling enables scalable stochastic estimation of decomposition terms without additional hyperparameters or inner optimization loops.
- Training with Minibatch-Weighted Sampling: The minibatch estimator is biased because its expectation is a lower bound, although computing it requires no additional hyperparameters.
- Special case: β-TCVAE: β-TCVAE sets α = γ = 1 and tunes only β, while an equivalent objective in related work estimates total correlation with an auxiliary discriminator.
4 Measuring Disentanglement with the Mutual Information Gap
The mutual information gap (MIG) measures how much each latent variable captures a ground-truth factor while penalizing entangled, non-axis-aligned representations. It is classifier-free, bounded, and applicable across factorized latent distributions when mutual information can be efficiently estimated.
- Metric construction: MIG uses empirical mutual information between latent variables and known ground-truth factors to quantify disentanglement.The estimator is based on a joint distribution formed from the empirical data and encoder distributions.
- Metric construction: Normalized mutual information I(z_j; v_k)/H(v_k) is maximal when a latent variable has a deterministic, invertible relationship with a discrete factor.For discrete factors, the normalized quantity lies between 0 and 1.
- Metric construction: MIG measures the difference between the two latent variables with highest mutual information for each factor, enforcing axis-alignment.The gap also discourages redundant latent variables from encoding the same factor.
- Interpretation: The gap penalizes rotated representations in which several latent variables each contain information about multiple factors.It also favors compact representations in which one latent variable reliably models a factor.
- Properties: MIG is generally applicable to factorized latent distributions, including multimodal, categorical, and other structured distributions, provided mutual information can be estimated.The authors report that their simple estimator runs in reasonable time on the datasets studied.
5 Related Work
Prior work studies disentanglement through weak supervision, unsupervised generative modeling, factorial representations, and objective reweighting. The paper positions its analysis and metric as extending these lines of work toward stronger quantification and broader applicability.
- Prior approaches: Weakly supervised approaches invert generative processes with known disentangled factors while aiming to distill simpler representations rather than achieve perfect inversion.Unsupervised generative modeling frameworks have also explored disentanglement, and earlier methods succeeded in limited settings with few factors.
- Factorial representations: Factorial representations motivate disentanglement, but the paper presents a more general analysis of the underlying objective than earlier related treatments.Prior analyses considered total correlation under additional assumptions such as a flexible prior or q(z) = p(z).
- Objective-based methods: FactorVAE uses an objective equivalent to β-TCVAE but estimates total correlation with a density-ratio discriminator and an inner optimization loop.The paper instead emphasizes a refined ELBO decomposition and a training method for assigning arbitrary weights to its terms.
- Independent component analysis: Non-linear independent component analysis studies inversion with independent latent factors, whereas this paper targets maximizing mutual information between learned representations and ground-truth factors.The paper also states that simple priors can encourage interpretability by warping complex factors into simpler manifolds.
6 Experiments
Experiments evaluate β-TCVAE and MIG across controlled and biased factor distributions, quantitative comparisons, and qualitative image datasets. β-TCVAE generally improves the trade-off between density estimation and disentanglement, with total correlation strongly associated with disentanglement under the proposed method.
- Quantitative Comparisons: β-TCVAE consistently achieves higher MIG scores than β-VAE and InfoGAN, matches FactorVAE, and performs better when the density ratio trick is difficult to train.These results come from quantitative and qualitative experiments across multiple settings.
- ELBO vs. Disentanglement Trade-off: β-TCVAE provides a better trade-off between density estimation and disentanglement as β increases.Higher β overly penalizes mutual information in β-VAE, whereas β-TCVAE maintains higher disentanglement scores relative to β-VAE.
- Ablation Studies: Ablations show no significant performance difference after removing the index-code MI term or using a jointly trained factorized normalizing-flow prior.The results suggest that tuning the total-correlation weight is the most useful intervention for learning disentangled representations.
- Factorial vs. Disentangled Representations: For β-TCVAE, average total correlation and average MIG are strongly negatively correlated, whereas β-VAE shows a weaker correlation.At the same total correlation, β-TCVAE produces a better disentangled model; the comparison averages results over 40 random initializations on dSprites and faces.
- Correlated or Dependent Factors: Across sampling distributions with varying correlation and dependence, β-TCVAE is superior to β-VAE and InfoGAN, with substantially higher median scores in most configurations.The experiment uses a two-factor synthetic 3D-faces dataset and includes uniform, biased, dependent, and correlated settings.
- Qualitative Comparisons: Qualitative traversals show β-TCVAE learning additional chair properties and more meaningful extrapolations on CelebA than β-VAE.On chairs, the additional properties are material and leg rotation; on CelebA, β-TCVAE discovers 15 attributes and generates rare samples such as bald or mustached females.
7 Conclusion
The paper decomposes the ELBO to motivate β-TCVAE and introduces MIG as a classifier-free disentanglement metric. Experiments provide evidence that independence between latent variables can be strongly related to disentanglement in simple datasets with uniform factors.
- Method: The ELBO decomposition identifies a total-correlation penalty and motivates β-TCVAE, which trains with minibatch estimation and no additional hyperparameters.The method is designed as a simple plug-in replacement for β-VAE and can integrate into different frameworks.
- Metric: MIG is a classifier-free information-theoretic metric that evaluates disentanglement while enforcing compactness.The metric uses mutual information between latent variables and ground-truth factors.
- Conclusion: In simple datasets with uniform factors, independence between latent variables can be strongly related to disentanglement.This conclusion is presented as evidence within the paper’s supported experimental scope.
A.1 Qualitative Samples
The section describes qualitative samples and the stochastic estimation machinery used to estimate mutual information and ELBO-decomposition terms from minibatches.
- Real samples from the training dataset provide qualitative reference images.
- MIG evaluates mutual information between each latent variable z_j and each ground-truth factor v_k.The joint distribution is defined using the inference network and the empirical data-generating process.
- The estimator assumes that q(z_j|x) and p(n|v_k) can be sampled, with empirical distributions for p(v_k) and p(n|v_k).
- The mutual-information bound is tight when a latent variable and factor are invertibly related.Otherwise, a high conditional entropy H(v_k|z_j) indicates that the latent retains less information about the factor.
- Minibatch estimation uses sampled indices and weighted sampling to estimate q(z) and the decomposed ELBO terms.The estimator is exact when the minibatch contains the full empirical dataset, while its bias decreases as M increases.
- The stochastic estimator is a lower bound because of Jensen’s inequality, with equality at M = N.
C.2.2 Experiments
The experiments examine estimator choices, ablations, flexible priors, batch size, and metric sensitivity across β-TCVAE and related settings.
- MSS performs similarly to MWS in preliminary experiments.
- Removing the index-code MI penalty produces no significant difference between α = 0 and α = 1, and the preferred setting is dataset-dependent.
- A factorial normalizing flow prior yields no significant improvement over the factorial Gaussian prior in preliminary experiments.The flow is trained jointly with the β-TCVAE objective and is intended to model multimodal distributions.
- Lower batch sizes remain effective for β-TCVAE despite using a high batch size to address minibatch-estimation bias.This observation is reported on 3D faces with β = 6.
- Classifier-based metric values depend strongly on the hyperparameter L, which can alter model gaps and relative rankings.
F. MIG Traversal
MIG scores track increasing disentanglement and axis alignment, while revealing weaknesses associated with rotation, redundant factors, and latent uncertainty.
- MIG gives low scores when even two latent variables are not axis-aligned, but rotation is harder to score because some factors have redundant values.
- Traversal plots show only latent means, so similar-looking plots can conceal uncertainty changes that affect mutual information and MIG.
- Scores near 0.0 correspond to extremely entangled representations.
- Scores near 0.1–0.2 indicate reduced entanglement but persistent axis-alignment failures or incomplete structure.
- Scores near 0.3–0.4 correspond to increasingly disentangled and axis-aligned representations, with residual nuances, rotation entanglement, or imperfect position.
- Scores near 0.5 indicate axis-aligned, disentangled representations, although higher scores may partly reflect reduced latent entropy.
G. Disagreements Between Metrics
The comparison shows that MIG and the Higgins metric often agree on strong models but diverge when representations contain subtle entanglement or lack axis alignment.
- Both metrics generally agree on the most disentangled models, but MIG falls off more quickly than Higgins across trained β-VAE models.The comparison uses 200 models with varying β values and initializations.
- The Higgins metric tends to overestimate disentanglement because it misses subtle differences and axis-alignment failures.
- For the most controversial model, Higgins ranks it above 75% of models while MIG ranks it above only 26%.Only the scale factor appears disentangled; position is not axis-aligned and two rotation variables mirror each other.
- MIG penalizes models where multiple latent variables influence one factor, whereas Higgins does not enforce single-variable factor influence.
- The top five controversial models are identified by the largest difference between Higgins and MIG rankings.