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.
