Source-linked AI summary
Structured Disentangled Representations
Babak Esmaeili, Hao Wu, Sarthak Jain, Alican Bozkurt, N. Siddharth, Brooks Paige, Dana H. Brooks, Jennifer Dy, Jan-Willem van de Meent
TL;DR
Disentanglement methods based on simple diagonal Gaussian priors do not reliably separate discrete factors. This paper introduces a two-level hierarchical objective derived from the ELBO, and experiments show improved discrete disentanglement, other-factor disentanglement, and generalization to unseen combinations.
Problem
Existing unsupervised disentanglement approaches generally assume simple diagonal Gaussian priors and do not reliably disentangle discrete factors of variation.
Method
HFVAEs use a two-level hierarchical objective that controls independence between variable blocks and individual variables within blocks through total-correlation decompositions.
Results
Experiments across datasets demonstrate that HFVAEs disentangle discrete variables and improve disentanglement of other variables, including in image and text settings.
Takeaways & Limitations
The framework provides structured control over where statistical independence is induced and supports disentanglement of higher-dimensional variable groups.
Takeaways & Limitations
Discrete-factor disentanglement depends on how separable the underlying classes or identities are, and the framework is less effective when visual distinctions are difficult to distinguish.
Abstract
from arXiv · showhide
Deep latent-variable models learn representations of high-dimensional data in an unsupervised manner. A number of recent efforts have focused on learning representations that disentangle statistically independent axes of variation by introducing modifications to the standard objective function. These approaches generally assume a simple diagonal Gaussian prior and as a result are not able to reliably disentangle discrete factors of variation. We propose a two-level hierarchical objective to control relative degree of statistical independence between blocks of variables and individual variables within blocks. We derive this objective as a generalization of the evidence lower bound, which allows us to explicitly represent the trade-offs between mutual information between data and representation, KL divergence between representation and prior, and coverage of the support of the empirical data distribution. Experiments on a variety of datasets demonstrate that our objective can not only disentangle discrete variables, but that doing so also improves disentanglement of other variables and, importantly, generalization even to unseen combinations of factors.
1 Introduction
Deep generative models use low-dimensional latent codes, but their dimensions are often not individually interpretable. HFVAEs introduce a two-level objective that controls independence among variable groups and within groups.
- Background: Deep generative models encode high-dimensional data in low-dimensional latent variables used for generation and downstream tasks.VAEs jointly learn probabilistic encoders and decoders by optimizing the evidence lower bound.
- Motivation: Disentangled representations aim for individual latent dimensions whose perturbations produce interpretable changes in the data.Standard latent dimensions do not necessarily correspond to semantically meaningful variation.
- Contribution: HFVAEs use a two-level hierarchical VAE decomposition to control statistical independence between variable groups and between variables within each group.At each level, total correlation induces statistical independence.
- Prior work: The objective builds on prior modifications of VAE terms, including fully unsupervised approaches such as β-VAE.These approaches alter, add, or remove terms in the VAE objective to induce disentanglement.
2 A Unified View of Generalized VAE Objectives
The VAE objective can be viewed as a KL divergence between generative and inference models, clarifying how its terms constrain data and latent marginals. Decomposing the objective exposes distinct roles for reconstruction identifiability, mutual information, data likelihood, and prior matching.
- VAE formulation: VAEs jointly learn a generative decoder distribution and an encoder that performs approximate inference over latent variables.Together, the encoder and decoder define a probabilistic autoencoder.
- KL formulation: The VAE objective is equivalently expressed as a KL divergence between the generative model pθ(x, z) and inference model qφ(z, x).This formulation makes the optimization targets and marginal-matching conditions explicit.
- Marginal consistency: Minimizing the KL requires the model and inference joint distributions to match, implying pθ(x)=q(x) and qφ(z)=p(z).The inference marginal qφ(z) is the encoder distribution averaged over the empirical data.
- Four-term decomposition: The four-term decomposition separates marginal consistency from conditional consistency between data and latent variables.Terms 3 and 4 constrain data and latent marginals, while terms 1 and 2 govern conditional identifiability and mutual information.
- Conditional terms: Combining terms 1 and 3 recovers the expected decoder log-likelihood, avoiding the intractability of evaluating pθ(x) pointwise.Term 2 regularizes term 1 by minimizing I(z; x) under the inference model.
- Removing terms: Removing terms 3 or 4 permits mismatches between model and empirical data marginals or between the inference marginal and prior.Removing term 1 can allow an expressive decoder to ignore the latent code.
3 Hierarchically Factorized VAEs (HFVAEs)
HFVAE decomposes the VAE’s KL regularization hierarchically, separately controlling independence between variable groups and within groups. This yields an objective that generalizes prior disentanglement objectives while enabling practical batch-based estimation.
- Hierarchical decomposition: The hierarchical decomposition separates total-correlation penalties between groups from KL penalties on constituent variables.The resulting hierarchy distinguishes group-level independence from within-group distribution matching.
- Hierarchical decomposition: Total correlation regularization encourages statistical independence between latent variables and can therefore induce disentangled representations.For grouped variables, the decomposition can continue recursively to impose independence at multiple levels.
- HFVAE objective: α controls mutual-information regularization, β controls total correlation between groups, and γ controls total correlation within groups.Setting β = γ recovers a non-hierarchical special case that includes related objectives.
- HFVAE objective: Unlike related objectives, HFVAE includes a component that need not assume the prior factorizes across variables.Matching inference and prior total correlations allows the prior structure to participate in the decomposition.
- Approximation: The objective is optimized by Monte Carlo estimates of qφ(z), qφ(zd), and qφ(zd,e) using the same training batch.This avoids requiring an exact full-dataset computation of the inference marginals during each objective evaluation.
4 Related Work
Related work modifies VAE objectives or uses alternative generative frameworks to encourage disentanglement. HFVAE generalizes several objective modifications while differing in its hierarchical decomposition and treatment of prior factorization.
- VAE objectives: HFVAE is related to VAE objectives that alter mutual-information or KL terms, including approaches based on covariance matching and correlation explanation.The paper positions its objective as a generalization of several recently proposed VAE modifications.
- Alternative frameworks: InfoGAN encourages disentanglement by maximizing mutual information between individual features and data under the generative model.This provides a GAN-based alternative to modifying the VAE objective.
- β-VAE trade-offs: Using β < 1 in β-VAE has been explored to improve reconstruction quality when disentanglement is not the primary goal.The related-work discussion connects this setting to the trade-off between likelihood and KL regularization.
- Objective comparison: Table 1 compares objectives in autoencoding deep generative models and marks objectives whose priors factorize.The asterisk denotes the factorization condition p(z) = Q d p(zd).
5 Experiments
Experiments across image and text datasets evaluate whether HFVAE learns interpretable, hierarchically structured representations, including discrete factors and controlled topic correlations. Results show discrete-variable disentanglement, pruning of unused dimensions, and generalization to unseen factor combinations, with effectiveness depending on class separability.
- Experimental Setup: HFVAE experiments cover CelebA, dSprites, MNIST, F-MNIST, and 20NewsGroups, comparing its objective with VAE and other disentanglement objectives.The experiments use Concrete variables for explicit classes and continuous variables for style where appropriate.
- Image Representations: HFVAE recovers seven interpretable MNIST features, including slant, width, height, openness, stroking, thickness, and roundness.The representation ignores remaining dimensions, whose mutual information with the data drops to zero.
- Text Representations: HFVAE organizes 20NewsGroups topics into internally uncorrelated groups while preserving sparse correlations between related groups.It learns 11 topics with significant mutual information, compared with nonzero mutual information for all 50 VAE topics.
- Scope Boundary: Discrete-factor disentanglement depends on how visually separable the underlying classes or identities are, working less effectively for difficult-to-distinguish dSprites shapes.F-MNIST classes are visually distinctive, while dSprites shape differences can be only a few pixels.
- Discrete Factors: HFVAE better separates digit identity from thickness than β-VAE, concentrating label information in the discrete variable rather than spreading it across continuous dimensions.The thickness manipulation preserves digit identity across the tested range for HFVAE, whereas β-VAE does not.
- Hyperparameter Analysis: Increasing β with fixed γ=3 increases the mutual information gap, while γ=β from 2 to 6 reduces total correlation by 5 nats and mutual information by 0.5 nats.The analysis reports a trade-off between total correlation and mutual information and finds that higher β helps the Concrete variable capture label information.
- Zero-shot Generalization: After combinations of factors are removed from training, HFVAE reconstructs unseen digits and correctly encodes their stroke and character-width features.The result supports evaluation of generalization to previously unseen combinations of factors.
6 Discussion
The discussion frames structured disentanglement as extending beyond independent scalar factors, while detailing estimators for entropy and mutual information used in the objective.
- Scope and motivation: The work targets correlations between higher-dimensional latent variables, whereas prior disentanglement research largely studied uncorrelated scalar factors.The authors describe this as a first step toward more structured disentangled representations.
- Scope and motivation: The objective can enforce or relax statistical independence between groups of variables, enabling disentanglement of higher-dimensional representations.
- Future work: The discussion identifies datasets for more rigorous testing of correlations between higher-dimensional variables as future work.
- Estimation procedure: Replacing excluded samples with other batch elements produces an estimator over C = B −1 samples, and the estimator is unbiased.
- Estimation procedure: The entropy estimator uses samples from the inference marginal and yields an upper bound on the marginal entropy.
B Model Architectures
The image experiments use four datasets with continuous and categorical latent variables, supported by dataset-specific encoder and decoder architectures.
- Datasets: The experiments cover dSprites, MNIST, F-MNIST, and CelebA, spanning shapes, handwritten digits, clothing, and celebrity faces.The datasets contain 737,280 dSprites images, 60,000 MNIST images, 60,000 F-MNIST images, and 202,599 CelebA images.
- Latent variables: Each dataset uses one normal hidden variable for continuous factors and one Concrete hidden variable for categorical factors.The implementation uses Adam with learning rate 1e-3 and default settings.
- Architectures: Tables 4–6 specify the encoder and decoder architectures for MNIST, F-MNIST, dSprites, and CelebA.
C Latent Traversals
The latent traversals visualize how changing one latent dimension affects generated examples while all other dimensions remain fixed.
- MNIST: In the MNIST traversals, one zd dimension varies from -3 to 3 while the remaining dimensions stay fixed at 0.The visualization uses 10% supervision to extract the digit more reliably and expose style features.
- F-MNIST: In the F-MNIST traversals, one zd dimension varies from -3 to 3 while the remaining dimensions stay fixed at 0.
- 20NewsGroups: The 20NewsGroups visualization compares learned topics under the HFVAE and VAE objectives.It reports frequent words, their mutual information with x, and mutual information between latent codes and document-category indicators.
D.1 Model Architectures
The topic-model experiments use 20NewsGroups documents with bag-of-words inputs and compare ProdLDA and NVDM architecture settings.
- Dataset: 20NewsGroups contains 11,314 documents partitioned into 20 categories and represented with a 2,000-word bag-of-words vocabulary.Stopwords are removed using the Mallet stopwords list.
- HFVAE: HFVAE uses two hidden variables, zc1 and zc2, each with 25 dimensions.
- Optimization: ProdLDA uses Adam with β1 = 0.99, β2 = 0.999, and learning rate 1e-3, while NVDM uses learning rate 5e-5.
D.2 Neural variational document model
The HFVAE improves disentanglement of document-category information compared with a standard NVDM, while individual features may capture multiple topics.
- The HFVAE uses two 25-dimensional latent variables, with β = 7 and γ = 4, allowing within-group but preventing across-group correlations.The baseline is a standard NVDM with a 50-dimensional latent variable trained using the normal VAE objective.
- HFVAE latent dimensions show higher disentanglement because category indicators generally correlate with only one latent feature.Figure 12 compares mutual information between latent dimensions and binary document-category indicators.
- A single latent feature can capture two distinct topics through negative and positive weights in the likelihood model.
D.3 Binary Indicator Variables for Document Category
The analysis groups related 20NewsGroups categories into binary topic indicators and measures their mutual information with individual latent dimensions.
- The 20NewsGroups analysis derives 10 binary variables indicating whether documents belong to specific grouped topics.Highly related newsgroups are treated as one topic before constructing the indicators.
- Mutual information is computed between each binary topic indicator b_l and individual latent dimension z_d.The resulting relationships are shown in Figure 12.
- Table 9 lists the topics formed by grouping highly related categories.