[Paper Review] Disentangling by Factorising

Hyunjik Kim, Andriy Mnih|arXiv (Cornell University)|Feb 16, 2018

Digital Media Forensic Detection49 references423 citations

TL;DR

FactorVAE introduces a total correlation penalty to encourage factorial latent code distributions in VAEs, yielding better disentanglement than beta-VAE for similar reconstruction quality. It also proposes a robust, discriminator-based metric for disentanglement and compares with InfoWGAN-GP.

ABSTRACT

We define and address the problem of unsupervised learning of disentangled representations on data generated from independent factors of variation. We propose FactorVAE, a method that disentangles by encouraging the distribution of representations to be factorial and hence independent across the dimensions. We show that it improves upon $β$-VAE by providing a better trade-off between disentanglement and reconstruction quality. Moreover, we highlight the problems of a commonly used disentanglement metric and introduce a new metric that does not suffer from them.

Motivation & Objective

• Motivate unsupervised learning of disentangled representations that align latent factors with independent data factors.
• Analyze the disentanglement-reconstruction trade-off in beta-VAE and propose a method to improve this balance.
• Introduce and validate a Total Correlation penalty to promote independence in the latent code.
• Propose a robust disentanglement metric that avoids weaknesses of prior metrics.
• Compare FactorVAE with beta-VAE and InfoWGAN-GP across multiple datasets with known and unknown factors.

Proposed method

• Augment the VAE objective with a Total Correlation term to encourage factorial latent distributions: maximize log-likelihood while penalizing KL(q(z)||bar{q}(z)) where bar{q}(z) is the product of marginals.
• Approximate the TC term using a discriminator that differentiates samples from q(z) and from bar{q}(z) via a density-ratio trick.
• Use a permutation-based sampling procedure (Alg. 1) to approximate bar{q}(z) without full data passes.
• Train the VAE and the TC-discriminator jointly, with TC gradient signals scaled by a gamma hyperparameter.
• Present pseudocode for FactorVAE (Alg. 2) and discuss stability considerations in latent-space divergences (as opposed to data-space).
• Introduce a new, hyperparameter-free disentanglement metric that measures which latent dimension corresponds to a fixed factor by evaluating the minimum empirical variance across dimensions after normalisation.

Experimental results

Research questions

• RQ1 Does a Total Correlation penalty on the latent code improve disentanglement without sacrificing reconstruction quality?
• RQ2 How does FactorVAE compare to beta-VAE and InfoWGAN-GP in terms of disentanglement and reconstruction across datasets with known and unknown factors?
• RQ3 What are the weaknesses of the Higgins et al. disentanglement metric, and can a more robust alternative be proposed?
• RQ4 Can a discriminator-based TC estimate provide a stable and effective optimization signal for disentangling?

Key findings

• FactorVAE achieves higher disentanglement scores than beta-VAE for similar reconstruction quality on 2D Shapes and 3D Shapes.
• The proposed TC-based penalty reduces Total Correlation in the latent code, yielding improved independence between latent dimensions.
• The new disentanglement metric is conceptually simpler, hyperparameter-free, and avoids the failure mode of prior metrics.
• InfoWGAN-GP generally underperforms compared to VAE-based methods on the tested shapes datasets and shows sensitivity to architecture.
• FactorVAE maintains competitive or better reconstruction quality while attaining superior disentanglement scores across several datasets, including those with unknown factors.
• Discriminator-based TC estimates tend to underestimate true TC, but decreasing TC correlates with better disentanglement during training.