Understanding Variational Autoencoders: A Guide for Machine Learning Students

Introduction to Variational Autoencoders

In the landscape of generative machine learning, the Variational Autoencoder (VAE) stands as a cornerstone architecture. Unlike standard autoencoders that map inputs to fixed points in a latent space, VAEs map inputs to probability distributions. This probabilistic approach allows the model to learn a continuous, structured latent space, making it a powerful tool for generating new, realistic data samples.

For undergraduate students, understanding VAEs is essential not only for generative modelling tasks but also for grasping the intersection of deep learning and Bayesian inference. This article breaks down the architecture, the mathematical objectives, and the practical implementation details you need to excel in your studies.

The Architecture: Encoder and Decoder

A VAE consists of two primary neural networks: the encoder and the decoder. The encoder takes an input $x$ and outputs the parameters of a distribution in the latent space, typically the mean $\mu$ and the variance $\sigma^2$. Instead of a single vector, we sample a latent vector $z$ from this distribution: $z \sim \mathcal{N}(\mu, \sigma^2)$.

The decoder then takes this sampled $z$ and attempts to reconstruct the original input $x$. By forcing the model to sample from a distribution, we ensure that the latent space is smooth and continuous, which is critical for generative tasks.

The Reparameterisation Trick

A significant challenge in training VAEs is that sampling is a stochastic process, which is non-differentiable. To allow backpropagation, we use the reparameterisation trick. Instead of sampling $z$ directly from $\mathcal{N}(\mu, \sigma^2)$, we express $z$ as:

$$z = \mu + \sigma \odot \epsilon$$

where $\epsilon \sim \mathcal{N}(0, 1)$. Here, $\epsilon$ is an auxiliary noise variable. Because $\mu$ and $\sigma$ are now deterministic outputs of the encoder, the gradient can flow through them during training.

Worked Example 1: Suppose an encoder outputs $\mu = [0.5, -0.2]$ and $\sigma = [0.1, 0.3]$. If we sample $\epsilon = [0.8, -0.5]$, calculate the latent vector $z$.

1. $z_1 = 0.5 + (0.1 \times 0.8) = 0.5 + 0.08 = 0.58$
2. $z_2 = -0.2 + (0.3 \times -0.5) = -0.2 - 0.15 = -0.35$ Result: $z = [0.58, -0.35]$.

The Loss Function: ELBO and KL Divergence

The VAE loss function consists of two parts: the reconstruction loss and the KL divergence. The reconstruction loss ensures the output resembles the input, while the KL divergence acts as a regulariser, forcing the latent distribution to be close to a standard normal distribution $\mathcal{N}(0, 1)$.

The KL divergence between the learned distribution $q(z|x)$ and the prior $p(z)$ is given by:

$$D_{KL} = -\frac{1}{2} \sum (1 + \log(\sigma^2) - \mu^2 - \sigma^2)$$

Worked Example 2: Given $\mu = 0$ and $\sigma^2 = 1$ (log variance = 0), calculate the KL divergence.

1. $D_{KL} = -0.5 \times (1 + 0 - 0^2 - 1) = -0.5 \times (0) = 0$. This confirms that when the distribution matches the prior, the penalty is zero.

Latent Space Structure

The latent space is the compressed representation of your data. In a well-trained VAE, similar inputs are mapped to nearby regions in the latent space. Because we use the KL divergence to pull the distributions towards $\mathcal{N}(0, 1)$, the latent space avoids 'holes' or disjointed clusters, allowing for smooth interpolation between data points.

Common Mistakes

1. Ignoring the Reparameterisation Trick: Attempting to backpropagate through a stochastic sampling node will result in a broken gradient flow.
2. Misinterpreting the KL Term: Forgetting that the KL divergence is a regulariser, not just a constraint. If the weight of the KL term is too high, the model may ignore the input and just output the prior.
3. Confusing AE and VAE: Remember that standard autoencoders learn a deterministic mapping, whereas VAEs learn a probabilistic mapping.

Frequently Asked Questions

What is the purpose of the latent space? It provides a compressed, lower-dimensional representation of the input data, capturing the essential features required for reconstruction.

Why do we use KL divergence? It prevents the encoder from assigning each input to a unique, isolated point, ensuring the latent space is continuous and generative.

Can VAEs generate new data? Yes. By sampling $z$ from the prior $\mathcal{N}(0, 1)$ and passing it through the decoder, the model generates entirely new samples.

Conclusion

Variational Autoencoders are a fascinating blend of probability theory and neural network architecture. By mastering the reparameterisation trick and the balance between reconstruction and regularisation, you are well on your way to understanding modern generative AI. To see these concepts in motion, visit MathInstructor AI to generate a free animated lesson on VAEs.