Variational Autoencoders

The previous two chapters developed the machinery for gradient-based VI: the ELBO, the reparameterisation trick, and Adam. A Variational Autoencoder (VAE) Kingma & Welling, 2014Rezende et al., 2014 puts these tools to work in a deep generative model. The core idea is simple:

Fit a nonlinear latent variable model by jointly learning a decoder (generative model) and an encoder (approximate posterior) using the reparameterised ELBO.

The result is a model that can both generate new data (by sampling from the prior and decoding) and embed data into a structured latent space (by encoding).

Topics covered:

• The VAE generative model: from PPCA to deep latent variable models

• The ELBO as reconstruction loss + KL regulariser

• Amortised inference: why we use an encoder network instead of per-datapoint variational parameters

• The complete training algorithm in a single clean loop

• Code: a VAE on 2D data with a 1D latent space

The Generative Model¶

From PPCA to VAEs¶

Recall Probabilistic PCA from Chapter 6:

The decoder $\mbz \mapsto \mathbf{W}\mbz + \mbmu$ is linear. A VAE replaces it with a neural network:

where $g : \mathbb{R}^H \to \mathbb{R}^D$ is a multilayer perceptron (MLP) with parameters $\mbtheta$, and $p(\mbx \mid g(\mbz_n; \mbtheta))$ is a suitable observation model (Gaussian for real data, Bernoulli for binary data, etc.).

The prior $p(\mbz_n) = \mathcal{N}(\mathbf{0}, \mathbf{I})$ acts as a regulariser: it encourages the latent space to be compact and well-organised.

Why Is the Posterior Intractable?¶

In PPCA the posterior $p(\mbz_n \mid \mbx_n)$ is Gaussian and has a closed form. In a VAE, the nonlinear decoder means

is no longer a standard distribution — the normalising constant requires integrating over $\mbz_n$, which is intractable. This is where variational inference comes in.

The ELBO: Reconstruction + KL¶

Introduce a variational approximation $q(\mbz_n; \mblambda_n)$ to the intractable posterior. The ELBO for data point $n$ is:

This decomposition has a clean interpretation:

Maximising the ELBO encourages good reconstructions while keeping the latent codes close to the prior — the same tension as a regularised autoencoder.

Gaussian Encoder and Decoder¶

Use a diagonal Gaussian encoder $q(\mbz_n; \mblambda_n) = \mathcal{N}(\mbmu_n, \mathrm{diag}(\mbsigma_n^2))$. The KL term then has a closed form:

For a Gaussian decoder $p(\mbx_n \mid \mbz_n) = \mathcal{N}(g(\mbz_n; \mbtheta), \mathbf{I})$, the reconstruction term is just

where $\mbepsilon_n \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ is the reparameterised noise.

Amortised Inference: The Encoder Network¶

The Per-Datapoint Bottleneck¶

With $N$ data points we have $N$ sets of variational parameters $\{\mblambda_n\}$. If we optimise each independently — running many gradient steps on $\mblambda_n$ for every update of $\mbtheta$ — the E-step dominates the computation and the method does not scale.

The Key Observation¶

The optimal variational parameters are a function of the data point:

Instead of learning $N$ independent $\mblambda_n$ vectors, we learn a single encoder network (also called a recognition network or inference network) $f(\mbx; \mbphi)$ that maps any data point to its variational parameters:

Amortisation means we amortise the cost of inference across the entire dataset: one forward pass through the encoder gives $\mblambda_n$ instantly, at the cost of a slightly suboptimal approximation.

The Amortisation Gap¶

The encoder can only represent functions in the parametric class defined by its architecture. This introduces an amortisation gap:

There is also an approximation gap from restricting $q$ to the diagonal-Gaussian family. Cremer et al., 2018 study both gaps empirically.

The VAE Training Algorithm¶

With an encoder, the ELBO for data point $n$ becomes a function of both $\mbtheta$ (decoder) and $\mbphi$ (encoder):

Because $\mbz_n$ is reparameterised through $\mbphi$ via $\mbmu_\phi$ and $\mbsigma_\phi$, the gradient $\nabla_{\mbphi} \mathcal{L}_n$ flows through the encoder by standard backpropagation.

The complete algorithm — no separate E and M steps:

repeat:
    sample mini-batch {x_n} ⊂ {x_1, …, x_N}
    for each x_n in the mini-batch:
        compute (μ_φ(x_n), σ_φ(x_n)) via encoder forward pass
        sample ε ~ N(0, I);  z_n = μ_φ(x_n) + σ_φ(x_n) ⊙ ε
        compute ELBO_n = log p(x_n | g(z_n; θ)) − KL term
    loss = −mean(ELBO_n)
    loss.backward()             # gradients wrt both θ and φ
    optimiser.step()            # update θ and φ together

The encoder and decoder are trained jointly in a single optimisation loop — there is no explicit E-step / M-step separation.

Conclusion¶

A VAE is a nonlinear latent variable model trained by amortised variational inference.

The ELBO has two interpretable terms:

Key takeaways:

• Amortisation replaces $N$ separate optimisation problems (one per data point) with a single shared encoder, enabling scalability.

• Reparameterisation makes the ELBO differentiable through both the decoder (via $\mbtheta$) and the encoder (via $\mbphi$), so both can be trained in a single gradient loop.

• The KL term with a standard normal prior has a closed form for diagonal Gaussian encoders, eliminating MC noise from that term.

• A linear VAE (linear encoder and decoder) recovers PPCA exactly, so VAEs are a strict generalisation of the linear LVMs from Chapter 6.

• Two sources of sub-optimality: the approximation gap (diagonal Gaussian may not match the true posterior shape) and the amortisation gap (the encoder may not perfectly represent the optimal $\mblambda_n$ for every $\mbx_n$).