Table of Contents
Fetching ...

An Introduction to Discrete Variational Autoencoders

Alan Jeffares, Liyuan Liu

TL;DR

The paper addresses learning with discrete latent spaces in VAEs, arguing that categorical latents can better capture structured data such as text. It presents a canonical discrete VAE with independent categorical latents, derives an ELBO-based objective, and shows how to compute gradients for both the encoder and decoder using the log-derivative trick. A concrete MNIST-style implementation with a Bernoulli decoder is described, together with a minimalist PyTorch training recipe and practical gradient expressions. The work provides a rigorous, self-contained guide to designing and training discrete latent representations in VAEs and sets a foundation for future discrete latent-variable methods.

Abstract

Variational Autoencoders (VAEs) are well-established as a principled approach to probabilistic unsupervised learning with neural networks. Typically, an encoder network defines the parameters of a Gaussian distributed latent space from which we can sample and pass realizations to a decoder network. This model is trained to reconstruct its inputs and is optimized through the evidence lower bound. In recent years, discrete latent spaces have grown in popularity, suggesting that they may be a natural choice for many data modalities (e.g. text). In this tutorial, we provide a rigorous, yet practical, introduction to discrete variational autoencoders -- specifically, VAEs in which the latent space is made up of latent variables that follow a categorical distribution. We assume only a basic mathematical background with which we carefully derive each step from first principles. From there, we develop a concrete training recipe and provide an example implementation, hosted at https://github.com/alanjeffares/discreteVAE.

An Introduction to Discrete Variational Autoencoders

TL;DR

The paper addresses learning with discrete latent spaces in VAEs, arguing that categorical latents can better capture structured data such as text. It presents a canonical discrete VAE with independent categorical latents, derives an ELBO-based objective, and shows how to compute gradients for both the encoder and decoder using the log-derivative trick. A concrete MNIST-style implementation with a Bernoulli decoder is described, together with a minimalist PyTorch training recipe and practical gradient expressions. The work provides a rigorous, self-contained guide to designing and training discrete latent representations in VAEs and sets a foundation for future discrete latent-variable methods.

Abstract

Variational Autoencoders (VAEs) are well-established as a principled approach to probabilistic unsupervised learning with neural networks. Typically, an encoder network defines the parameters of a Gaussian distributed latent space from which we can sample and pass realizations to a decoder network. This model is trained to reconstruct its inputs and is optimized through the evidence lower bound. In recent years, discrete latent spaces have grown in popularity, suggesting that they may be a natural choice for many data modalities (e.g. text). In this tutorial, we provide a rigorous, yet practical, introduction to discrete variational autoencoders -- specifically, VAEs in which the latent space is made up of latent variables that follow a categorical distribution. We assume only a basic mathematical background with which we carefully derive each step from first principles. From there, we develop a concrete training recipe and provide an example implementation, hosted at https://github.com/alanjeffares/discreteVAE.
Paper Structure (14 sections, 58 equations, 3 figures, 1 algorithm)

This paper contains 14 sections, 58 equations, 3 figures, 1 algorithm.

Figures (3)

  • Figure 1: An Autoencoder. The encoder network $f_{\bm{\theta}}$ maps to a latent representation $\mathbf{z}$ which is decoded by a second network, $g_{\bm{\phi}}$, attempting to reproduce the original input.
  • Figure 2: A Variational Autoencoder. Unlike the vanilla autoencoder, the encoder of this variational version outputs the parameters of a probability distribution from which we sample and decode. The parameters are learned by maximizing the ELBO which provides a lower bound on the log-likelihood.
  • Figure 3: The Discrete Variational Autoencoder. The input is encoded into the parameters of $D$ categorical distributions; then, we sample one of the $K$ categories from each. These samples are passed back through the decoder that attempts to reproduce the original input.