PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Disentangled Joint Continuous and Discrete Representations

Emilien Dupont

arXiv:1804.00104v3stat.MLcs.LG

TL;DR

Disentanglement methods have primarily modeled continuous factors, although many datasets also contain discrete generative factors. JointVAE extends VAEs with relaxed discrete latents and controlled information capacity to learn both types unsupervised. It discovers such factors across several datasets, with a clear advantage when discrete factors are prominent, while remaining sensitive to hyperparameters and random seeds.

  • Problem

    Most disentanglement methods focus on continuous factors, making inherently discrete generative factors difficult to capture.

  • Method

    JointVAE augments a VAE with a relaxed discrete distribution and controls latent information capacity to learn jointly continuous and discrete representations unsupervised.

  • Results

    JointVAE disentangles continuous and discrete factors across MNIST, FashionMNIST, CelebA, and Chairs, with a clear advantage when discrete factors are prominent.

  • Takeaways & Limitations

    The learned representations can separate factors, infer unlabeled properties on test data, and support simple image editing.

  • Takeaways & Limitations

    Disentanglement quality is sensitive to hyperparameters and random seeds, and discrete factors are less clear in datasets such as CelebA and Chairs.

Abstract

from arXiv · show

We present a framework for learning disentangled and interpretable jointly continuous and discrete representations in an unsupervised manner. By augmenting the continuous latent distribution of variational autoencoders with a relaxed discrete distribution and controlling the amount of information encoded in each latent unit, we show how continuous and categorical factors of variation can be discovered automatically from data. Experiments show that the framework disentangles continuous and discrete generative factors on various datasets and outperforms current disentangling methods when a discrete generative factor is prominent.

1 Introduction

Disentangled representations align individual representation units with separate data factors while remaining invariant to others. The paper motivates combining continuous and discrete latent factors because many datasets contain both.

  • Motivation: Disentangled representations change one factor of variation while remaining invariant to others.Examples include separate units for object position, color, or scale.
  • Motivation: Most recent disentanglement methods focus on continuous factors, leaving inherently discrete factors difficult to capture.Distinct objects or entities may be discrete, while position and scale are continuous.
  • Related work: InfoGAN models discrete and continuous factors through mutual information but inherits GAN drawbacks including unstable training and reduced sample diversity.The passage notes that later GAN improvements mitigate some issues, while stable training remains challenging.
  • Related work: β-VAE is stable to train but can model only continuous latent variables.This motivates a VAE-based framework with additional flexibility for discrete generative factors.
  • Contribution: JointVAE learns disentangled continuous and discrete representations unsupervised while retaining VAE advantages such as stable training and large sample diversity.It discovers factors across MNIST, FashionMNIST, CelebA, and Chairs, including digit type separately from slant, width, and stroke thickness on MNIST.

2 Analysis of β-VAE

The β-VAE analysis explains how its KL term limits information transmitted by latent variables and how controlled capacity trades reconstruction quality against disentanglement. Gradually increasing the capacity offers a way to improve reconstruction without reducing disentanglement.

  • β-VAE objective: β-VAE combines a reconstruction likelihood with a KL penalty that encourages inferred latents to match a prior distribution.The encoder maps x to z, the decoder maps z back to x, and β weights the KL term.
  • Information control: The expected KL divergence is an upper bound on mutual information between the data and latent variables.A mini-batch estimate of mean KL therefore estimates an upper bound on the information z can transmit about x.
  • Information control: Penalizing mutual information improves disentanglement but increases reconstruction error.This establishes the central trade-off analyzed in the section.
  • Controlled capacity: The controlled-capacity objective forces the KL divergence toward a target capacity C using a weighting constant γ.The capacity is gradually increased during training to control how much information the model encodes.
  • Controlled capacity: Gradually increasing C improves reconstruction quality compared with the original objective without reducing disentanglement.This result is attributed to the controlled-capacity formulation.

3 JointVAE Model

JointVAE extends β-VAE with jointly continuous and discrete latent variables, using relaxed categorical distributions and separate capacity control to encourage information in both channels. Its encoder samples both latent types, concatenates them, and passes the combined representation to the decoder.

  • JointVAE models continuous latents z and categorical latents c with a joint posterior, prior, and likelihood.
  • Assuming conditional independence between z and c, the joint KL divergence separates into continuous and discrete terms.
  • Separate capacities Cz and Cc are gradually increased to prevent the model from ignoring discrete variables or assigning all capacity to continuous channels.
  • Discrete posterior variables use Gumbel Softmax distributions, whose temperature-controlled relaxation makes categorical sampling differentiable.
  • The encoder outputs Gaussian parameters μ and σ2 for continuous latents and class probabilities α(i) for discrete latents, which are sampled, concatenated, and decoded.

4 Experiments

Experiments show JointVAE can disentangle continuous and discrete factors across image datasets, with especially clear advantages when discrete factors are prominent. The model also supports unsupervised factor inference and latent-space image editing, but remains sensitive to hyperparameters and random seeds.

  • MNIST: JointVAE discovers digit type as a discrete factor and stroke thickness, angle, and width as continuous factors on MNIST.The model used 10 continuous latent variables and one 10-dimensional discrete latent variable, generating realistic samples across many latent settings.
  • MNIST: JointVAE separates digit type from continuous factors on MNIST, whereas β-VAE and CCβ-VAE entangle digit type with continuous variation.This advantage is attributed to JointVAE’s ability to represent the inherently discrete digit-type factor.
  • FashionMNIST: JointVAE discovers interpretable FashionMNIST factors including sleeve length, heel size, and shirt color, although similar classes are not all disentangled.Discovered classes include dress, t-shirt, trousers, sneakers, bag, and ankle boot.
  • CelebA: On CelebA, JointVAE discovers azimuth, age, and background color continuously, while discrete settings correspond to different facial identities.The model used 32 continuous latent variables and one 10-dimensional discrete latent variable; samples retain facial details but are less sharp than entangled-model samples.
  • Chairs: On Chairs, JointVAE discovers rotation, width, and leg style, while discrete settings represent chair types and colors.The model used 32 continuous latent variables and three binary discrete latent variables.
  • dSprites: The dSprites evaluation reports that only JointVAE encodes all four continuous and one discrete factors into four continuous and one discrete latent variables.Scores were averaged over 10 random seeds using the best-hyperparameter model, excluding collapsed-to-mean outliers.
  • Detecting disentanglement in latent distributions: Increasing latent channel capacity progressively activates units that encode factors, while near-zero KL units have little or no effect on outputs.For MNIST, digit type is discovered before angle and width, consistent with its larger reconstruction-error reduction.
  • The inference network: JointVAE’s inference network identifies MNIST digit type with 88.7% unsupervised accuracy and infers chair rotation without labeled data.The same latent structure also enables editing celebrity faces by changing azimuth, background color, or hairstyle coordinates.

5 Conclusion

JointVAE learns disentangled continuous and discrete representations unsupervised while retaining VAE advantages such as stable training and sample diversity. It also supports inference of unlabeled quantities and image editing.

  • JointVAE learns disentangled continuous and discrete representations in an unsupervised manner.
  • The framework models complex jointly continuous and discrete generative factors while retaining stable training and large sample diversity.
  • JointVAE disentangles factors of variation across several datasets while producing realistic samples.
  • Its inference network can infer unlabeled quantities on test data and edit or manipulate images.
  • Future work includes combining JointVAE with FactorVAE or β-TCVAE and studying how latent channel capacities affect disentanglement stability.

A Proofs

The proof section defines aggregate encoding distributions and decomposes the latent KL divergence using mutual information. It assumes conditional independence between continuous and discrete latent variables.

  • The joint distribution of data and encodings is defined as q(z, x) = p(x)qφ(z|x).
  • The marginal latent distribution is defined by averaging qφ(z|x) over the data distribution p(x).
  • The KL divergence between posterior and prior is rewritten using mutual information under the joint data-encoding distribution.
  • The joint KL decomposition assumes continuous and discrete latents are conditionally independent in both posterior and prior distributions.

B Model architecture

The model uses ReLU nonlinearities throughout the encoder and decoder, except for a sigmoid decoder output. The 32-by-32 image configuration removes one convolutional layer from each side.

  • The encoder and decoder use ReLU nonlinearities, with a sigmoid output layer in the decoder.
  • The architecture shown in the table is used for 64 by 64 Chairs, CelebA, and dSprites images.
  • For 32 by 32 MNIST and FashionMNIST images, the encoder’s last convolutional layer and decoder’s first convolutional layer are removed.

C Training details

Training details vary latent dimensionality, optimizer learning rate, and the schedules for continuous and discrete latent capacities across model configurations. Capacity parameters are increased linearly over training iterations.

  • The listed configurations use either 10 continuous and one 10-dimensional discrete latent variable, or 32 continuous and three binary discrete variables.
  • Adam is used with learning rates of 5e-4 or 1e-4 across the configurations.
  • For one configuration, Cz and Cc increase linearly from 0 to 5 over 25000 iterations.
  • For another configuration, Cz increases from 0 to 5 and Cc from 0 to 10 over 50000 iterations.
  • For the 32-continuous, three-binary-discrete configuration, Cz increases from 0 to 30 and Cc from 0 to 5 over 100000 iterations.
  • For the 32-continuous, one-10-dimensional-discrete configuration, Cz increases from 0 to 50 and Cc from 0 to 10 over 100000 iterations.

C.5 dSprites

The dSprites configuration used six continuous latent variables and one three-dimensional discrete variable, with capacities increased during training and discrete capacity clipped at its categorical maximum.

  • Configuration: The latent distribution contained 6 continuous variables and 1 3-dimensional discrete variable.Training used Adam at learning rate 5e-4, batch size 64, and 30 epochs.
  • Capacity scheduling: Cz increased linearly from 0 to 40 over 300000 iterations.
  • Capacity scheduling: Cc increased linearly from 0 to 1.1 over 300000 iterations.
  • Capacity clipping: The discrete capacity was clipped when it exceeded the maximum capacity log n of the categorical variable.This bound follows from the KL divergence between a categorical variable and a uniform categorical variable.

D Things that didn’t work

Several attempted modifications failed to reliably disentangle joint continuous and discrete representations, while comparisons with InfoGAN highlighted entanglement and training-stability issues. CelebA also required an arbitrary choice of what discrete factors represent, and latent traversals used standardized ranges for visualization.

  • Things that didn’t work: Adding Gaussian and Gumbel-Softmax latents without changing the β-VAE loss generally caused the model to ignore discrete codes.
  • Things that didn’t work: Reweighting continuous and discrete KL terms either caused discrete codes to be ignored or made continuous factors enter the discrete latents.These outcomes occurred for large combinations of the two β values.
  • Things that didn’t work: Approximating the Gumbel-Softmax KL divergence produced high-variance estimates that impeded learning, with no closed-form expression available.
  • CelebA scope: On CelebA, the discrete factor choice was somewhat arbitrary: 10 dimensions encoded 10 facial identities, while more dimensions encoded more identities.This behavior was generally robust except when the number of discrete dimensions exceeded 100, when the mode became affected.
  • Comparison with InfoGAN: InfoGAN disentangled digit type from angle and width, but width and stroke thickness remained entangled, and its models were typically less stable to train.
  • Visualization: Continuous latent traversals ran from Φ−1(0.05) to Φ−1(0.95), whereas discrete traversals ran from 1 to the variable’s number of dimensions.
Loading 1804.00104v3…