[Paper Review] Disentangling Disentanglement in Variational Autoencoders

Émile Mathieu, Tom Rainforth+2 more|arXiv (Cornell University)|Dec 6, 2018

Adversarial Robustness in Machine Learning55 references52 citations

TL;DR

The paper generalizes disentanglement in VAEs as latent decomposition governed by two factors—overlap in latent space and aggregate encoding matching a structured prior—and shows how prior choice and a new objective with alpha/beta can yield richer, customizable representations beyond simple independence.

ABSTRACT

We develop a generalisation of disentanglement in VAEs---decomposition of the latent representation---characterising it as the fulfilment of two factors: a) the latent encodings of the data having an appropriate level of overlap, and b) the aggregate encoding of the data conforming to a desired structure, represented through the prior. Decomposition permits disentanglement, i.e. explicit independence between latents, as a special case, but also allows for a much richer class of properties to be imposed on the learnt representation, such as sparsity, clustering, independent subspaces, or even intricate hierarchical dependency relationships. We show that the $β$-VAE varies from the standard VAE predominantly in its control of latent overlap and that for the standard choice of an isotropic Gaussian prior, its objective is invariant to rotations of the latent representation. Viewed from the decomposition perspective, breaking this invariance with simple manipulations of the prior can yield better disentanglement with little or no detriment to reconstructions. We further demonstrate how other choices of prior can assist in producing different decompositions and introduce an alternative training objective that allows the control of both decomposition factors in a principled manner.

Motivation & Objective

• Motivate a general notion of disentanglement as latent-space decomposition with two factors: overlap and prior-structure alignment.
• Show limitations of standard disentanglement definitions for complex data and demonstrate a more flexible decomposition framework.
• Analyze beta-VAE to understand how it controls latent overlap and how prior choice affects disentanglement.
• Propose an alternative objective that explicitly regularizes both decomposition factors to enable structured representations (e.g., sparsity, clustering).

Proposed method

• Define decomposition in VAEs as fulfilling two factors: appropriate latent overlap and aggregate encoding matching the prior.
• Relate beta-VAE to an adjusted ELBO with annealed prior and a decoder-side reconstruction term, plus a maximum-entropy regularizer on the encoder.
• Provide theoretical results showing beta-VAE’s equivalence to a standard ELBO with latent-space scaling and prior annealing for Gaussian cases.
• Introduce an objective L_{alpha,beta} that adds a divergence term between q(z) and p(z) to control the second decomposition factor.
• Experiment with anisotropic and non-Gaussian priors to study axis-aligned disentanglement, clustering, and sparsity.
• Present a sparsity-focused prior and evaluate with sparsity metrics and reconstruction performance.

Experimental results

Research questions

• RQ1 How can disentanglement be generalized beyond independence to accommodate complex data generation processes?
• RQ2 What roles do latent overlap (I(x; z)) and the match between aggregate latent encodings and the prior play in achieving useful latent decompositions?
• RQ3 Can changing the prior structure and introducing an explicit regularizer on q(z) improve disentanglement and enable alternative decompositions like sparsity or clustering?
• RQ4 How does beta-VAE relate to the two-factor decomposition, and can we reformulate objectives to control both factors independently?
• RQ5 Do non-isotropic priors or priors designed for sparsity/ clustering yield tangible gains in disentanglement without sacrificing reconstruction?

Key findings

• beta-VAE mainly controls latent overlap via a maximum-entropy effect on the encoder, with its benefit limited by prior invariance to rotations under isotropic Gaussian priors.
• For Gaussian priors and encoders, beta-VAE is equivalent to a rescaled latent space optimization of the standard ELBO, up to constant terms.
• Rotational invariance of the isotropic Gaussian prior can hinder disentanglement, and breaking this invariance with structured priors improves disentanglement for a fixed reconstruction quality.
• Explicitly regularizing the aggregate posterior to match a structured prior (via alpha) can improve the alignment of q(z) with p(z) and enable alternative decompositions such as clustering or sparsity.
• Experiments show that non-isotropic priors (e.g., anisotropic Gaussians or Student-t mixtures) can yield better disentanglement scores at similar reconstruction performance; priors promoting sparsity or clustering achieve this in Fashion-MNIST and synthetic datasets.
• An objective incorporating both decomposition factors (beta for overlap and alpha for prior alignment) enables learning of sparse and clustered latent representations without severely hurting reconstructions.