PaperlayerView exact revision on arXiv

Source-linked AI summary

Discrete Variational Autoencoders

Jason Tyler Rolfe

arXiv:1609.02200v2stat.MLcs.LG

TL;DR

Discrete latent variables fit datasets organized into classes but resist efficient backpropagation-based training. The paper introduces discrete VAEs, combining an undirected discrete component with hierarchical continuous layers and a hierarchical posterior, and reports state-of-the-art performance on three image benchmarks. The approach also has scope boundaries: standard reparameterization requires differentiable inverse CDFs, while the hierarchical model requires strong regularization and can be affected by RBM mixing.

  • Problem

    Discrete latent variables naturally model class-structured data, but efficient VAE-style backpropagation generally cannot pass through them.

  • Method

    Discrete VAEs combine an undirected binary latent-variable model with directed continuous latent layers and use a hierarchical posterior to accommodate strong correlations.

  • Results

    The resulting discrete variational autoencoder achieves state-of-the-art performance on permutation-invariant MNIST, Omniglot, and Caltech-101 Silhouettes.

  • Takeaways & Limitations

    The model learns both discrete object classes and their continuously deformable pixel-level realizations from unsupervised data.

  • Takeaways & Limitations

    Standard reparameterization requires an invertible, differentiable inverse CDF, while the hierarchical model can overfit without strong prior regularization and training quality depends on RBM sampling.

Abstract

from arXiv · show

Probabilistic models with discrete latent variables naturally capture datasets composed of discrete classes. However, they are difficult to train efficiently, since backpropagation through discrete variables is generally not possible. We present a novel method to train a class of probabilistic models with discrete latent variables using the variational autoencoder framework, including backpropagation through the discrete latent variables. The associated class of probabilistic models comprises an undirected discrete component and a directed hierarchical continuous component. The discrete component captures the distribution over the disconnected smooth manifolds induced by the continuous component. As a result, this class of models efficiently learns both the class of objects in an image, and their specific realization in pixels, from unsupervised data, and outperforms state-of-the-art methods on the permutation-invariant MNIST, Omniglot, and Caltech-101 Silhouettes datasets.

1 INTRODUCTION

The paper motivates combining discrete variables for disconnected object classes with continuous variables for within-class transformations, then introduces discrete VAEs to train this hybrid structure through the VAE framework. Standard VAEs and related reparameterization methods do not directly support discrete latent distributions because their inverse CDF is not differentiable.

  • Motivation: Image datasets can contain disconnected smooth manifolds, with discrete object types and continuous changes in pose, position, and lighting.The paper uses images as projections of real-world objects whose transformations remain smooth within a class but not across classes.
  • Motivation: Most probabilistic models use exclusively discrete or exclusively continuous latent variables, rather than representing class identity and within-class variation separately.The paper contrasts discrete models such as DBMs and NADEs with continuous models such as VAEs and GANs.
  • Contribution: Discrete VAEs combine an undirected binary latent-variable model with multiple directed continuous latent layers to capture object class and continuously deformable realization.The model is designed to represent both the discrete class of an object and its specific realization in pixels.
  • Contribution: Discrete VAEs enable efficient VAE training with backpropagation through binary latent variables while using a hierarchical posterior to keep the evidence lower bound tight under strong correlations.The hierarchical posterior is intended to model correlations that arise in the latent variables.
  • VAE limitation: The reparameterization trick works when samples can be generated by a differentiable deterministic function of inputs, parameters, and independent random variables.This provides low-variance stochastic gradient estimates for suitable continuous distributions and tractable approximating posteriors.
  • VAE limitation: For discrete distributions, the inverse CDF required by the reparameterization trick is undefined or nondifferentiable because the CDF is piecewise constant.The quantile derivative is either zero or infinite for a discrete distribution, preventing the usual reparameterization formulation.

2 BACKPROPAGATING THROUGH DISCRETE LATENT VARIABLES BY ADDING

The method smooths discrete latent variables with continuous auxiliaries so the VAE objective becomes differentiable while preserving the model’s discrete structure. A spike-and-exponential transformation provides an invertible, stochastic nonlinearity for backpropagation through binary variables.

  • Smoothing discrete variables: Continuous auxiliaries ζ smooth discrete variables z, making the conditional-marginal CDF invertible and differentiable for reparameterized gradients.The smoothing distribution preserves the model form and KL term while enabling direct use of the VAE gradient equations.
  • Smoothing discrete variables: The prior is augmented with ζ and requires p(x|ζ,z,θ) = p(x|ζ,θ), which insulates z from observed variables x during marginalization.This yields a deterministic, differentiable autoencoding network conditioned on independent uniform noise.
  • Autoencoding computation: The posterior computes q(z = 1|x, φ), combines it with ρ ∼U[0, 1], and applies F−1 to produce ζ before evaluating log p(x|ζ, θ).Conditioned on x and ρ, the autoencoder is deterministic and differentiable, allowing low-variance backpropagation.
  • Spike-and-exponential transformation: The transformation is invertible almost surely: ζi = 0 corresponds to zi = 0, while ζi > 0 corresponds to zi = 1.As β approaches infinity, ζi approaches zi almost surely, allowing smoothing variables to be removed after training.
  • Spike-and-exponential transformation: The spike-and-exponential inverse CDF rectifies q when q ≤ 1 −ρ and becomes increasing but concave-down when q > 1 −ρ.Its stochastic behavior is qualitatively similar to dropout and batch-normalization noise.

3 ACCOMMODATING EXPLAINING-AWAY WITH A HIERARCHICAL

The paper replaces factorial discrete posteriors with a hierarchical approximating posterior to represent explaining-away correlations while retaining tractable reparameterized computation. Smoothed autoregressive dependencies allow the resulting autoencoder to remain deterministic and differentiable given stochastic inputs.

  • Motivation: Posterior inference can exhibit strong correlations from explaining-away, while factorial approximations impose independent discrete variables for tractability.The issue is especially relevant when the prior itself is an RBM with correlated latent variables.
  • Hierarchical approximating posterior: The hierarchical posterior divides RBM variables into disjoint groups and defines dependencies through a directed acyclic graphical model over those groups.This structure is introduced specifically to accommodate correlated posterior distributions while maintaining tractability.
  • Hierarchical approximating posterior: Each discrete group zj depends on preceding groups only through their smoothed variables ζi<j, enabling the autoregressive hierarchy to capture correlations and multiple modes.The smoothed mediation also supports differentiable computation through the discrete hierarchy.
  • Autoencoding computation: The hierarchical posterior increases the autoencoder’s depth: each layer computes q(zj = 1|ζi<j, x, φ), samples ζj with F−1, and recursively passes the results onward.After all groups are computed, the full ζ and x determine the reconstruction log-likelihood.
  • Gradient estimation: The remaining KL-term gradients are estimated stochastically, and Equation 12 is reported to have substantially lower variance than the naive approach.This provides the gradient estimator used alongside the hierarchical autoencoding computation.

4 MODELLING CONTINUOUS DEFORMATIONS WITH A HIERARCHY OF CONTINUOUS LATENT VARIABLES

The model adds hierarchical continuous latent layers beneath the discrete component to represent continuous deformations such as position and pose. Layer-wise autoregressive priors and posteriors retain tractable Gaussian KL terms and standard reparameterized gradients.

  • Continuous latent hierarchy: Additional continuous latent layers below the RBM increase the expressive power of both the generative model and the approximating posterior.The paper focuses on continuous layers because they have proven effective in generative adversarial networks and variational autoencoders.
  • Model structure: The model defines both posterior and prior as layer-wise fully autoregressive directed graphical models over the added continuous variables.The same autoregressive variable order is used for the approximating posterior and prior.
  • Continuous latent hierarchy: Continuous latent variables build manifolds for properties such as position and pose, conditioned on discrete variables representing object types.The discrete and continuous components therefore separate object identity from continuously deformable realization.
  • Model structure: The full RBM-associated latent set includes grouped discrete and smoothed variables, while each subsequent zm denotes a continuous latent layer.The resulting graphical model is depicted as a hierarchy beneath the RBM variables.
  • Optimization: Gaussian posterior and prior layers have a closed-form KL divergence when their covariance matrices are diagonal, with gradients passed through earlier layers using standard reparameterization.This makes the continuous hierarchy computationally efficient under the stated covariance condition.

5 RESULTS

The discrete VAE combines a smoothed RBM, hierarchical posterior, and continuous latent hierarchy, trained with persistent block-Gibbs sampling. Experiments show that sampling quality and posterior hierarchy affect performance, while moderate RBM size is sufficient on MNIST.

  • Model and training: Discrete VAEs combine a smoothed RBM, a hierarchical approximating posterior, and multiple directed continuous latent layers.The distributions are parameterized with neural networks, and the RBM prior is sampled using persistent block-Gibbs chains.
  • Evaluation: The models are evaluated on permutation-invariant MNIST, Omniglot, and Caltech-101 Silhouettes using importance-weighted test log-likelihood estimates.MNIST uses both static and dynamic binarization; discrete-VAE results are averaged over 16 runs.
  • Generative behavior: Well-separated RBM modes correspond to different digit types, while continuous latent variables generate variation within each shared RBM sample.On dynamically binarized MNIST, 100 block-Gibbs iterations separate successive sample rows, and modes generally align with individual digits.
  • Sampling: Thousands of single-temperature block-Gibbs iterations are required to mix between RBM modes, indicating that sample quality can constrain training.Performance improves when more block-Gibbs iterations are used per minibatch, motivating more effective samplers such as parallel tempering.
  • Architectural sensitivity: Log likelihood plateaus once the MNIST RBM has at least 64 units, whereas increasing posterior hierarchy depth helps until added parameters increase overfitting risk.The four-to-eight-layer posterior reduction may reflect the three encoder layers added per hierarchical posterior layer.

6 CONCLUSION

Discrete latent variables fit datasets organized into disconnected classes but obstruct efficient VAE backpropagation. The discrete VAE addresses this by projecting posterior and prior symmetrically into continuous space while retaining the original discrete KL divergence.

  • 6 CONCLUSION: Discrete latent variables naturally represent datasets composed of discrete classes, but gradients based on VAE backpropagation generally cannot pass through them.This motivates a method that preserves discrete representations while enabling efficient gradient approximation.
  • 6 CONCLUSION: The discrete VAE symmetrically projects the approximating posterior and prior into continuous space and evaluates the autoencoding term there.The original discrete representation is marginalized for the autoencoding term, while the KL divergence remains evaluated in discrete space.
  • 6 CONCLUSION: The model uses a hierarchical posterior over discrete latent variables and a hierarchy of continuous latent variables to increase representational power.This architecture combines discrete class structure with continuous variation in the generated observations.
  • 6 CONCLUSION: The resulting discrete VAE achieves state-of-the-art performance on permutation-invariant MNIST, Omniglot, and Caltech-101 Silhouettes.The method is designed to retain discrete latent variables while supporting VAE-style training.
  • Multivariate reparameterization: The multivariate CDF is generally non-invertible, so the method replaces it with conditional-marginal CDFs whose inverses can be applied sequentially under positivity conditions.Each conditional-marginal CDF maps to [0, 1], and inversion proceeds in variable order when conditional-marginal probabilities are everywhere nonzero.

B THE DIFFICULTY OF ESTIMATING GRADIENTS OF THE ELBO WITH REINFORCE

REINFORCE provides a stochastic gradient estimate for models with discrete and continuous latent variables, but its high variance and limited use of decoder gradients can make training inefficient.

  • Estimator limitations: A stochastic ELBO-gradient estimator can handle discrete and continuous latent variables using computationally tractable samples, but its variance can make training slow and perform poorly.A baseline can reduce variance without changing the gradient.
  • Estimator limitations: REINFORCE estimates the effect of encoder changes without directly differentiating the conditional log likelihood through the decoder.Its estimate behaves like a finite-difference approximation rather than using the decoder gradient in the VAE reparameterization estimator.
  • Estimator limitations: REINFORCE captures less information about the decoder per sample because a change in one latent direction affects the estimate only when sampled in that direction.In a D-dimensional latent space, at least D samples are needed to span all directions.
  • Estimator limitations: REINFORCE can be much less efficient than direct-gradient estimation when the latent representation contains dozens of variables.The paper reports that hundreds of latent variables can nevertheless be used effectively when gradients are calculated efficiently.

C AUGMENTING DISCRETE LATENT VARIABLES WITH CONTINUOUS LATENT

Discrete VAE reparameterization augments discrete latent variables with continuous randomness so probability packets move smoothly, allowing decoder gradients to estimate changes in the ELBO.

  • Continuous reparameterization: VAEs partition probability into equal-mass packets, making high-probability regions more likely to be selected.The inverse CDF maps parameter-independent random variables to points in the posterior distribution.
  • Continuous augmentation: A packet’s mass stays constant while its location changes smoothly with encoder parameters, enabling the decoder gradient to estimate loss changes.This requires a differentiable inverse transformation and overlapping auxiliary-space regions for discrete states, apart from zero-or-one probabilities.
  • Limitation: The resulting gradient is low variance only when most packet movements have similar effects or packets are tightly clustered.Low-variance behavior is therefore conditional rather than guaranteed for arbitrary posterior geometries.
  • The discrete-variable problem: Discrete encoders prevent smooth packet motion because changing their parameters moves probability mass between separated discrete values.Most packets remain unchanged under small parameter changes, so decoder gradients cannot accurately capture the resulting loss change.
  • Continuous augmentation: The method augments discrete latent representations with continuous random variables, making the conditional-marginal CDF invertible and differentiable.The generative conditional depends only on the new continuous latent space, while the model retains the original discrete latent structure.

D.1 MIXTURE OF RAMPS

The mixture-of-ramps transformation smooths binary latent variables through an invertible, differentiable inverse CDF. Its shape preserves useful gradients, while overly separated component distributions produce high-variance estimates.

  • The transformation uses linear conditional distributions for ζ_i given z_i=0 and z_i=1.
  • The CDF of q(ζ|x, φ) is formed by mixing the conditional CDFs according to q(z=1|x, φ).
  • Figure 7 plots the inverse CDF as a function of q(z=1|x, φ) for ρ values 0.2, 0.5, and 0.8.
  • The inverse CDF remains non-flat across mixture weights, unlike the sigmoid probability of z, so it does not eliminate gradients.
  • When conditional distributions overlap too little, the smoothing transformation becomes too sharp and gradient estimates have large variance.

E TRANSFORMATIONS FROM DISCRETE TO CONTINUOUS LATENT

The method augments discrete latent variables with continuous variables so the posterior can use differentiable inverse-CDF reparameterization. It derives tractable gradient calculations for hierarchical posteriors while addressing additional KL terms.

  • The input-independent expansion r(ζ|z) is insufficient when discrete latents capture much less information than the input entropy.
  • The resulting ELBO includes the joint divergence between q(z,ζ|x,φ) and p(z,ζ|x,θ).
  • For factorial distributions, the expectation in the additional gradient term can be evaluated analytically without injecting sampling variance.
  • The Gaussian conditional construction yields closed-form KL terms, including a weighted sum over the two binary states.
  • Nonfactorial posteriors complicate gradients, especially because cross-entropy terms involve dependencies between variables at different hierarchical levels.

F.3 NAIVE APPROACH

The naive cross-entropy gradient treats discontinuous discrete samples with REINFORCE-like estimators. Its variance grows poorly with RBM size, motivating a chain-rule decomposition that reduces the scaling.

  • The naive approach differentiates expectations using log-probability gradients over all variables, producing a REINFORCE estimator.
  • The estimator’s variance is proportional to the number of contributing terms, which grows quadratically with the number of RBM units.
  • The chain-rule formulation changes the discrete sample into a function of q(z=1) and ρ, allowing integration over the transition region.
  • The resulting stochastic estimate uses importance weighting and ignores samples inconsistent with the selected transition, exploiting z_i∈{0,1}.
  • The improved gradient scales linearly with RBM units, compared with quadratic scaling for the naive method.

H ARCHITECTURE

The architecture combines a stochastic posterior pass with a prior pass over hierarchical continuous latents and uses regularization to control overfitting. Performance is evaluated against several established models across four datasets.

  • H ARCHITECTURE: The stochastic ELBO approximation samples successive continuous latent layers in a posterior pass, then traverses the prior conditioned on those samples.
  • H ARCHITECTURE: The hierarchical architecture overfits when the prior is overparameterized, whereas increasing posterior capacity generally does not reduce performance in the examined range.
  • H ARCHITECTURE: Figure 9 varies hidden units and deterministic hidden layers separately in the prior and approximating posterior while fixing 10 continuous latent layers without parameter sharing.
  • H ARCHITECTURE: Table 2 states that smaller datasets require more regularization and achieve optimal performance with a smaller prior.
  • H ARCHITECTURE: The model therefore limits decoder and prior network complexity, controls hidden-layer sizes, and shares parameters across hierarchical layers.
  • H ARCHITECTURE: Log-likelihood comparisons cover dynamically and statically binarized MNIST, Omniglot, and Caltech-101 Silhouettes against multiple recent probabilistic models.

J SUPPLEMENTARY RESULTS

Simplifying the discrete VAE’s prior or approximating posterior worsens performance, while visualizations assess whether the RBM learns separated modes across datasets. The results show distinct modes for most MNIST digits, but weaker or unclear modes for Omniglot and Caltech-101 Silhouettes.

  • Simplified models: −92.7 negative log-likelihood with 128 RBM units and −88.8 with 200 RBM units after restricting the observation network to an SBN-like form.The restriction uses a linear transformation followed by a pointwise logistic nonlinearity.
  • Simplified models: −97.0 log-likelihood results when lateral RBM connections are removed, reducing the model to independent binary latent variables.The resulting architecture is described as a noisy sigmoid belief network.
  • Simplified models: −102.9 log-likelihood after replacing the hierarchical approximating posterior with a factorial posterior alongside the simplified prior.This result uses 200 binary latent variables.
  • Mode evolution: Consistency of object classes across successive rows indicates that the RBM prior has well-separated modes.Each sub-row holds the RBM constant while continuous latent variables generate within-class variation.
  • Mode evolution: MNIST retains distinct modes for most digit types, whereas Omniglot lacks obvious modes and Caltech-101 may show a mode for large roughly convex blobs.The Omniglot observation is plausibly attributed to insufficient capacity from an RBM with 128 units.
Loading 1609.02200v2…