Disentangling by Factorising

Citation

• Authors: Hyunjik Kim and Andriy Mnih
• Year: 2018
• Venue: ICML 2018 (PMLR 80)
• arXiv: arXiv:1802.05983
• URL: https://arxiv.org/abs/1802.05983

One-Line Takeaway

FactorVAE argues that beta-VAE pays for disentanglement by over-penalizing latent information, and replaces that blunt pressure with a separate total-correlation penalty on the aggregate posterior.

Technical Core

• Let $$q(z) = \mathbb{E}_{p_{\mathrm{data}}(x)}[q(z \mid x)]$$ denote the aggregate posterior.
• The paper starts from the beta-VAE decomposition $$\mathbb{E}_{p_{\mathrm{data}}(x)}\!\left[\mathrm{KL}(q(z \mid x)\,\|\,p(z))\right] = I(x; z) + \mathrm{KL}(q(z)\,\|\,p(z)),$$ and argues that increasing $\beta$ penalizes both the dependence-breaking term and the information-carrying term.
• FactorVAE keeps the standard VAE KL weight and instead optimizes $$\mathcal{L}_{\mathrm{FactorVAE}} = \mathbb{E}_{q(z \mid x)}[\log p(x \mid z)] - \mathrm{KL}(q(z \mid x)\,\|\,p(z)) - \gamma\, \mathrm{KL}(q(z)\,\|\,\bar q(z)),$$ where $$\bar q(z) = \prod_{j=1}^d q(z_j).$$
• The added term is the total correlation of the aggregate posterior, so the method explicitly encourages independence across latent coordinates without directly increasing the mutual-information penalty beyond the baseline VAE.

Method

• Train a standard VAE encoder and decoder with a diagonal-Gaussian posterior and a factorized Gaussian prior.
• Approximate samples from $\bar q(z)$ by taking a batch of latent codes sampled from $q(z)$ and randomly permuting each latent dimension across the batch.
• Estimate the density ratio $$\log \frac{q(z)}{\bar q(z)}$$ with a discriminator $D(z)$, giving the practical approximation $$\mathrm{TC}(z) = \mathrm{KL}(q(z)\,\|\,\bar q(z)) \approx \mathbb{E}_{q(z)}\!\left[\log \frac{D(z)}{1 - D(z)}\right].$$
• Update the VAE and discriminator jointly so the VAE reduces total correlation while preserving reconstruction quality.
• Replace the Higgins et al. disentanglement metric with a lower-variance majority-vote metric based on which latent dimension has the smallest normalized variance when one generative factor is held fixed.

Empirical Evaluation

• The main quantitative data sets are 2D Shapes and 3D Shapes; the paper also reports qualitative latent traversals on 3D Chairs, 3D Faces, and CelebA where ground-truth factors are unavailable.
• Baselines include a standard VAE, beta-VAE, InfoGAN, and InfoWGAN-GP.
• On 2D Shapes, FactorVAE achieves a better disentanglement-reconstruction trade-off than beta-VAE. The paper reports the best mean score around $0.82$ for $\gamma = 40$ versus about $0.73$ for beta-VAE with $\beta = 4$, with both at reconstruction error around $45$.
• On 3D Shapes, the top disentanglement scores of FactorVAE and beta-VAE are similar, but FactorVAE reaches them at lower reconstruction error: about $3515$ for $\gamma = 36$ versus $3570$ for beta-VAE with $\beta = 24$.
• InfoWGAN-GP is more stable than InfoGAN, but on 2D Shapes it still learns weak disentanglement and often allocates meaningful information to the wrong latent type or to regions of the prior that decode poorly.
• On the unknown-factor data sets, the paper is limited to latent-traversal inspection, but reports smaller reconstruction error and visually sensible factor directions for FactorVAE relative to beta-VAE.

Assumptions And Limits

The paper assumes data are generated from independent factors of variation, which makes the target notion of disentanglement most natural on synthetic data and much less clean on realistic data. It also stresses that low total correlation is necessary but not sufficient for disentanglement: a representation can have $\mathrm{TC}(q(z)) = 0$ while still collapsing most latents and carrying little useful information. The proposed quantitative metric also requires access to controlled factor-fixed sampling, so realistic data sets without known factors still fall back to qualitative traversals.

Relevance To This Project

This paper is directly relevant to the repo's recurring question of what the KL term should be doing. Its concrete message is: if the goal is cleaner latent structure, do not automatically increase the whole per-example KL penalty. Instead, separate "keep $z$ informative" from "shape the global latent geometry" and regularize the distribution-level property you actually want. In the current project's language, FactorVAE sits close to InfoVAE connections and Router marginal KL, but with a different target: it pushes the aggregate latent distribution toward factorization rather than merely toward a reference prior. That makes it a useful comparison point for any future objective that tries to make continuous reasoning latents less entangled without reintroducing a strong information bottleneck.

Comparison Notes

Higgins et al. (2017) increases the whole KL term, while this paper argues that only the dependence part should be targeted aggressively. Hoffman and Johnson (2016) provides the decomposition that makes this argument legible, and InfoVAE pushes the same separation in a different direction by decoupling mutual-information pressure from aggregate-prior matching. Burgess et al. (2018) is complementary: it keeps the beta-VAE framing but replaces a fixed bottleneck with a capacity schedule.

Related Notes

• beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework
• Understanding disentangling in beta-VAE
• InfoVAE: Balancing Learning and Inference in Variational Autoencoders
• InfoVAE connections
• Router marginal KL