An Introduction to Discrete Variational Autoencoders

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.

• The paper introduces a rigorous treatment of discrete VAEs, shifting from Gaussian to categorical latent spaces for better symbolic representation.
• It details optimization challenges, using REINFORCE-based gradient estimators and discussing relaxations like Gumbel-Softmax to manage variance.
• The tutorial provides an end-to-end implementation recipe and explores implications for interpretable, structured generative modeling.

Discrete Variational Autoencoders: Foundations and Training

Background and Motivation

Variational Autoencoders (VAEs) have established themselves as a core probabilistic generative modeling approach, unifying deep neural networks with latent variable models under the evidence lower bound (ELBO) principle. Traditionally, VAEs employ Gaussian latent spaces since the continuous structure supports the reparameterization trick, enabling backpropagation through stochastic nodes and thus tractable optimization (Kingma et al., 2013, Doersch, 2016). However, many application domains, such as discrete symbolic reasoning or text, may be more naturally and efficiently modeled through discrete latent factors. The paper "An Introduction to Discrete Variational Autoencoders" (2505.10344) delivers a technically rigorous yet accessible treatment of discrete VAEs, where the latent space structure is categorical rather than Gaussian.

The tutorial articulates both the mechanistic and probabilistic perspectives, clarifies the required optimization techniques (including the necessity for specialized gradient estimators), and presents an end-to-end recipe for effective training. This exposition bridges the conceptual gap between standard VAEs and their discrete-latent variants and highlights the unique strengths and obstacles encountered when using discrete representations.

Theoretical Formulation of (Discrete) VAEs

A standard autoencoder is composed of an encoder network $f_{\bm{\theta}}$ that maps input data $\mathbf{x}$ to a latent representation $\mathbf{z}$, and a decoder $g_{\bm{\phi}}$ that attempts to reconstruct the original input from this compressed representation. In a vanilla autoencoder, both the encoding and decoding are deterministic.

Figure 1: An Autoencoder architecture with encoder $f_{\bm{\theta}}$ and decoder $g_{\bm{\phi}}$ forming a latent bottleneck $\mathbf{z}$.

In a VAE, the latent representation is modeled as a probability distribution—typically Gaussian—where the encoder predicts distribution parameters rather than single points. The model is trained by maximizing the ELBO, a lower bound on the marginal log-likelihood of the observed data, encapsulating both reconstruction fidelity and a regularization constraint on the latent distribution.

Figure 2: A Variational Autoencoder, where the encoder outputs parameters of a latent probability distribution; sampling is performed and the ELBO is maximized.

However, when the underlying semantics of the domain are inherently discrete, such as categorical symbols, structured events, or indices, a discrete latent space is more suitable. In the discrete VAE, the latent space comprises $D$ categorical latent variables, each allowing $K$ possible states. The encoder outputs the categorical probability parameters, while the decoder reconstructs the input conditioned on one-hot samples from these latent variables.

Figure 3: The Discrete VAE: input is mapped to the parameters of $D$ categorical distributions. Samples from these distributions are concatenated and decoded to reconstruct the input.

The discrete VAE can be formally described as follows:

• Prior: Each latent variable $\mathbf{z}^{(d)}$ is assigned a uniform categorical prior $$p(\mathbf{z}^{(d)}) = \text{Cat}(K^{-1}, ..., K^{-1})$$.
• Encoder/Posterior: The encoder produces categorical parameters $f_{\bm{\phi}}(\mathbf{x})$ for each latent dimension, forming $q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})$.
• Decoder/Likelihood: The decoder $g_{\bm{\theta}}$ generates a Bernoulli distribution per input dimension (typically used for binarized data like MNIST), modeling $p_{\bm{\theta}}(\mathbf{x}|\mathbf{z})$.

Optimization and Gradient Estimation

A central challenge in training discrete VAEs is the inability to naively apply the reparameterization trick due to the inherently non-differentiable nature of discrete sampling. The tutorial distills the standard ELBO objective and decomposes the gradients for encoder and decoder parameters:

• Decoder gradients ($\bm{\theta}$): The expectation in the ELBO over $q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})$ is independent of $\bm{\theta}$, so gradients can be propagated directly through Monte Carlo samples.
• Encoder gradients ($\bm{\phi}$): The posterior $q_{\bm{\phi}}$ depends on $\bm{\phi}$, preventing direct application of Monte Carlo. Instead, the paper deploys the log-derivative (score function/REINFORCE) trick:

$$\nabla_{\bm{\phi}} \mathbb{E}_{q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})} [f(\mathbf{z})] = \mathbb{E}_{q_{\bm{\phi}}} [f(\mathbf{z}) \nabla_{\bm{\phi}}\log q_{\bm{\phi}}(\mathbf{z}|\mathbf{x})]$$

This estimator is unbiased but can suffer from high variance. Advanced methods such as Gumbel-Softmax relaxation [jang2017categorical] or the use of control variates [grathwohl2018backpropagation] can further improve optimization, but the paper maintains focus on the conceptually clear REINFORCE-based approach while explicitly detailing the mathematical derivations leading to the final gradient forms.

The ELBO for a batch is efficiently estimated using:

$$\mathcal{L}^{\mathrm{ELBO}}_{\bm{\theta}, \bm{\phi}}(\mathbf{x}) \approx \mathrm{Entropy}(f_{\bm{\phi}}(\mathbf{x})) - D \log K - \mathrm{BCE}(g_{\bm{\theta}}(\mathbf{z}), \mathbf{x})$$

where entropy regularizes the posterior towards a high-entropy (less certain) regime, and BCE (binary cross entropy) quantifies reconstruction error.

Practical Implementation and Recipe

The tutorial provides a distilled recipe for discrete VAE training, suitable for direct translation into PyTorch or similar frameworks. The steps are as follows:

1. Encode input batch to obtain categorical parameters for $D$ latent variables.
2. Sample one-hot latent variables from these categorical distributions.
3. Decode the concatenated one-hot vectors to produce reconstructed outputs.
4. Compute gradients using Monte Carlo for decoder parameters and REINFORCE (score function) for encoder parameters.
5. Update model parameters via stochastic gradient ascent on the ELBO.

An efficient and minimal implementation is referenced to supplement the mathematical derivations.

Implications and Future Directions

This systematic treatment elucidates both the theoretical and practical landscape for learning discrete latent structures with VAEs. Discrete VAEs facilitate modeling of domains where symbolic structure, clustering, or room for interpretability is paramount. They have found broad application in generative text models, symbolic reasoning, and recent work on scalable sparse autoencoder analysis for interpretability of LLMs (Gao et al., 2024, Lieberum et al., 2024). Furthermore, discrete bottlenecks are critical for vector quantized approaches that have advanced state-of-the-art high-fidelity image synthesis ([van2017neural], [razavi2019generating]).

A salient practical implication is the ongoing challenge of high-variance gradient estimation for discrete variables. Future research will likely focus on improved relaxations (Gumbel-Softmax and extensions [jang2017categorical], [liu2023bridging]), tighter variational bounds, control variates, and hybrid discrete-continuous latent structures to exploit both the representational benefits of discreteness and the optimization advantages of continuous spaces. There is also significant scope for further analysis of the inductive biases imposed by discrete latent spaces, particularly for structured or symbolic data.

Conclusion

The paper provides a rigorous, technically complete introduction to discrete VAEs, synthesizing both foundational probability and practical methods for training with categorical latent variables. Through careful exposition of the ELBO, detailed gradient derivations, and clear mapping to practical training regimes, it sets a reference point for practitioners and researchers considering discrete latent generative modeling. This framework paves the way for innovation in interpretable, structured, and symbolic AI representations, with clear applications in language, vision, and large-scale model interpretability.