15 Variational Autoencoders

The previous chapters of this book developed two families of generative models. Gaussian mixture models (Chapter 7) introduced latent variable modelling and the expectation–maximisation algorithm for discrete latent spaces. Generative adversarial networks (Chapters 9–13) showed how to train implicit generators through adversarial games. In this chapter we combine neural networks with variational inference to build the Variational Autoencoder (VAE), a principled framework for learning deep latent variable models.

The journey proceeds in five steps:

1. We start with deterministic autoencoders and their connection to PCA.

2. We introduce deep latent variable models and explain why the posterior is intractable.

3. We derive the Evidence Lower Bound (ELBO) from two complementary perspectives.

4. We assemble the VAE: an amortised encoder, a neural decoder, and the reparameterisation trick for end-to-end training.

5. We discuss practical issues: posterior collapse, latent space structure, and generation quality.

Autoencoders

Before introducing the probabilistic machinery of VAEs, we review their deterministic predecessor: the autoencoder.

The encoder compresses each input $𝒙\in {\mathrm{ℝ}}^d$ into a latent code $𝒛=f_{\varphi }(𝒙)\in {\mathrm{ℝ}}^k$; the decoder attempts to reconstruct $𝒙$ from $𝒛$. Because $k<d$, the bottleneck forces the network to discover a compact representation that captures the most important structure of the data.

Autoencoders find applications in dimensionality reduction, denoising (by training on corrupted inputs), feature learning, and anomaly detection (by flagging inputs with high reconstruction error). However, they suffer from a fundamental limitation for generative modelling.

Linear Autoencoders and PCA

We now make precise the connection between autoencoders and the classical technique of Principal Component Analysis (PCA). Recall that PCA projects data onto the top $k$ eigenvectors of the covariance matrix $\mathrm{\Sigma }=\frac{1}{N}{\sum }_{i=1}^N({𝒙}_i-\overline{𝒙}){({𝒙}_i-\overline{𝒙})}^{\top }$, thereby minimising the reconstruction error among all linear projections.

Latent Variable Models

We now shift from deterministic compression to probabilistic generative modelling. The central idea is to posit a latent (unobserved) variable $𝒛$ that explains the observed data $𝒙$ through a generative process: first sample $𝒛$ from a simple prior, then generate $𝒙$ conditioned on $𝒛$.

This framework generalises the Gaussian mixture model studied in Exercise 14: in a GMM the latent variable is discrete (the cluster label) and the decoder is a simple Gaussian $𝒩({𝝁}_k,{\mathrm{\Sigma }}_k)$. The deep LVM replaces the discrete $𝒛$ with a continuous Gaussian and the simple decoder with a neural network, vastly increasing expressiveness.

The goal is to find parameters $\theta$ that maximise the log-likelihood of the observed data: (MLE)$$\begin{array}{ccc}& {\theta }^{\ast }={\mathrm{arg\hspace{0.17em}max}}_{\theta }{𝔼}_{p_{\mathrm{data}}(𝒙)}[\log p_{\theta }(𝒙)].& \text{<span class="tml-eqn"></span>}\end{array}$$ This is equivalent to minimising the KL divergence ${𝖣}_{\mathrm{KL}}(p_{\mathrm{data}}\Vert p_{\theta })$ (cf. Definition 3).

The Evidence Lower Bound

Since maximising $\log p_{\theta }(𝒙)$ directly is intractable, we introduce an approximate posterior $q(𝒛)$ and derive a tractable lower bound on the log-evidence. We present two derivations that reveal complementary insights.

Jensen Derivation

The ELBO decomposes naturally into two interpretable terms.

The reconstruction term encourages $q$ and $\theta$ to explain the data well; the KL term penalises $q$ for deviating from the prior. This tension between fidelity and regularity is the defining feature of variational methods.

Bayes Derivation

The second derivation reveals the gap between the ELBO and the true log-evidence.

Since ${𝖣}_{\mathrm{KL}}\ge 0$, we immediately recover the lower-bound property $\log p_{\theta }(𝒙)\ge ℒ$. More importantly, the gap identity tells us exactly when the bound is tight.

The VAE Architecture

Classical variational inference optimises a separate $q(𝒛)$ for each data point $𝒙$, which is computationally expensive for large datasets. The VAE introduces amortised inference: a single neural network, called the encoder or inference network, maps any $𝒙$ to the parameters of an approximate posterior $q_{\varphi }(𝒛|𝒙)$.

The VAE is trained by maximising the ELBO summed (or averaged) over the dataset. Using the amortised posterior $q_{\varphi }$ in place of a generic $q$, the VAE objective becomes:

(Objective)$$\begin{array}{ccc}& ℒ(\varphi ,\theta )={𝔼}_{p_{\mathrm{data}}(𝒙)}[{𝔼}_{q_{\varphi }(𝒛|𝒙)}\left[\log p_{\theta }(𝒙|𝒛)\right]-{𝖣}_{\mathrm{KL}}\left(q_{\varphi }(𝒛|𝒙)\Vert p(𝒛)\right)].& \text{(VAE)}\end{array}$$

Gaussian VAE: Closed-Form KL Divergence

A major computational advantage of the Gaussian VAE is that the KL term in the ELBO can be evaluated in closed form. The general KL divergence between multivariate Gaussians was derived in Proposition 2; here we specialise to the diagonal-covariance case that arises in the standard VAE.

Substituting the closed-form KL and a single-sample Monte Carlo estimate for the reconstruction term, the practical VAE loss for a minibatch $\{{𝒙}_1,\dots ,{𝒙}_M\}$ becomes:

(FULL LOSS)$$\begin{array}{ccc}& \widehat{ℒ}(\varphi ,\theta )=\frac{1}{M}\sum _{m=1}^M[\log p_{\theta }({𝒙}_m|{𝒛}_m)-\frac{1}{2}\sum _{j=1}^d({\mu }_{\varphi ,j}^2({𝒙}_m)+{\sigma }_{\varphi ,j}^2({𝒙}_m)-\log {\sigma }_{\varphi ,j}^2({𝒙}_m)-1)],& \text{<span class="tml-eqn"></span>}\end{array}$$ where ${𝒛}_m={𝝁}_{\varphi }({𝒙}_m)+{\sigma }_{\varphi }({𝒙}_m)\odot {ϵ}_m$ with ${ϵ}_m\sim 𝒩(0,I)$.

The Reparameterisation Trick in VAEs

The ELBO involves the expectation ${𝔼}_{q_{\varphi }(𝒛|𝒙)}[\log p_{\theta }(𝒙|𝒛)]$. To maximise $ℒ$ with gradient-based optimisation, we need ${\nabla }_{\varphi }$ of this expectation. However, naïvely sampling $𝒛\sim q_{\varphi }(𝒛|𝒙)$ creates a stochastic node that blocks gradient flow: the sample depends on $\varphi$, but the sampling operation is not differentiable.

The solution is the reparameterisation trick, introduced in Definition 5 (Chapter 3). Here we show how it applies specifically to the VAE.

VAE Training Algorithm

We now consolidate the preceding results into complete training and sampling algorithms.

Posterior Collapse

One of the most common failure modes in VAE training occurs when the encoder learns to ignore the input, making the approximate posterior match the prior on some or all latent dimensions.

Posterior collapse arises from the tension between the two ELBO terms. The KL term ${𝖣}_{\mathrm{KL}}(q_{\varphi }(𝒛|𝒙)\Vert p(𝒛))$ is minimised when $q_{\varphi }=p$ (yielding ${𝖣}_{\mathrm{KL}}=0$), while the reconstruction term ${𝔼}_{q_{\varphi }}[\log p_{\theta }(𝒙|𝒛)]$ benefits from an informative posterior. If the decoder is powerful enough to model the data distribution without relying on $𝒛$ (for example, an autoregressive decoder such as PixelCNN), then the “easiest” solution is to set $q_{\varphi }\approx p$ and let the decoder do all the work.

Several mitigation strategies have been proposed:

• KL annealing (warm-up): During early training, weight the KL term by a factor ${\beta }_t$ that increases from $0$ to $1$ over the first several epochs. This gives the encoder time to learn useful representations before the full regularisation pressure is applied [4].

• Free bits: Replace ${𝖣}_{\mathrm{KL}}(q\Vert p)$ with $\max ({𝖣}_{\mathrm{KL}}(q\Vert p),\lambda )$ per dimension, ensuring that each latent dimension carries at least $\lambda$ nats of information [5].

• $\beta$-VAE weighting: Use $\beta {𝖣}_{\mathrm{KL}}$ with $\beta <1$ to reduce the regularisation pressure, at the cost of less structured latent spaces.

Latent Space Structure and Generation Quality

The KL regularisation term in the ELBO endows the VAE's latent space with useful geometric properties that deterministic autoencoders lack.

Exercises