[Paper Review] Understanding disentangling in $\beta$-VAE

Christopher Burgess, Irina Higgins+5 more|arXiv (Cornell University)|Apr 10, 2018

Generative Adversarial Networks and Image Synthesis1 references280 citations

TL;DR

The paper analyzes why β-VAE learns disentangled representations through an information bottleneck view and proposes a capacity-increase training regime to improve disentanglement and reconstruction.

ABSTRACT

We present new intuitions and theoretical assessments of the emergence of disentangled representation in variational autoencoders. Taking a rate-distortion theory perspective, we show the circumstances under which representations aligned with the underlying generative factors of variation of data emerge when optimising the modified ELBO bound in $\\beta$-VAE, as training progresses. From these insights, we propose a modification to the training regime of $\\beta$-VAE, that progressively increases the information capacity of the latent code during training. This modification facilitates the robust learning of disentangled representations in $\\beta$-VAE, without the previous trade-off in reconstruction accuracy.

Motivation & Objective

• Investigate how β-VAE encourages disentangled representations under a rate-distortion perspective.
• Understand the role of the information bottleneck and posterior-prior capacity in shaping latent axes.
• Propose training modifications that progressively increase latent capacity to improve disentanglement without sacrificing reconstruction quality.

Proposed method

• Frame β-VAE as an information bottleneck where q(z|x) acts as a noisy channel transmitting data information to the decoder.
• Relate the β-weighted KL term to an upper bound on information transmitted through latent channels.
• Use a simplified model to show factor-wise information allocation under capacity constraints (KL terms per factor).
• Demonstrate a capacity-control objective that gradually increases the target KL (capacity) during training to encourage sequential disentanglement.
• Train on datasets like dSprites and chairs to visualize latent traversals and factor alignment.

Experimental results

Research questions

• RQ1 Why does β-VAE tend to align latent axes with underlying generative factors of variation?
• RQ2 How does increasing latent capacity during training affect disentanglement and reconstruction quality?
• RQ3 Can a controlled capacity increase yield robust disentangling across different datasets (e.g., dSprites, 3D Chairs)?
• RQ4 What factors contribute to the locality and axis-alignment of representations in β-VAE compared to standard VAE?

Key findings

• β-VAE disentangling emerges from pressures to preserve data locality while allocating capacity to factors contributing most to data likelihood.
• Increasing latent capacity gradually leads to axis-aligned latent factors corresponding to distinct generative factors (position, scale, shape, rotation, color).
• A capacity-increase training regime yields robust disentangling with improved reconstruction quality compared to fixed β-VAE objective.
• On coloured dSprites and 3D Chairs, latent traversals show factor-specific changes with high-quality reconstructions, indicating interpretable disentangled axes.
• The proposed capacity-control objective (gradually approaching target KL) enables disentangling without the typical reconstruction-disentanglement trade-off.