Variational autoencoders

The big idea

A plain autoencoder compresses well, but its latent space is a scatter of islands. Pick a random point between them and the decoder draws a smudge, because nothing in training ever asked that point to mean anything. To generate new data you want the opposite: a latent space you can sample from blindly, where every point decodes to something plausible and nearby points decode to similar things.

A variational autoencoderAn autoencoder whose encoder outputs a probability distribution over codes, trained so that sampled codes reconstruct the input and the code distribution stays close to a simple prior, usually a standard normal. (VAE) gets there with two changes. The encoder no longer outputs one point; it outputs a small cloud, a mean and a spread, and the decoder must reconstruct the input from a random point in that cloud. And a penalty pulls every cloud toward the same standard normal distribution centred on zero. Think of a library where each book must be shelved somewhere inside one room, and each shelf label is a little fuzzy. Fuzzy labels force similar books onto neighbouring shelves (a wrong guess still lands on something similar), and the single room means there are no empty wings. Then you can walk to any shelf and find a sensible book.

How it works

Step 1: A probabilistic encoder

For an input $x$, the encoder outputs two vectors: a mean $\mu (x)$ and a standard deviation $\sigma (x)$, one pair per latent number. Together they define a Gaussian $q(z\mid x)=\mathcal{N}(\mu ,{\sigma }^2)$: "this input's code is probably around here". During training the decoder receives a sample $z$ from that Gaussian, not $\mu$ itself.

Step 2: The loss: reconstruction plus KL

Training balances two terms for each input:

The first term is the reconstruction loss from Autoencoders , now averaged over random codes from the cloud. For pixels in [0, 1] it is usually binary cross-entropy. The second term is the KL divergenceKullback-Leibler divergence: a measure of how much one probability distribution differs from another; zero only when they are identical. between the encoder's Gaussian and the standard normal prior. For one latent number it has a closed form:

Worked example: $\mu =1$, $\sigma =1$ gives $\frac{1}{2}(1+1-1)-0=0.5$. With $\mu =0$, $\sigma =0.1$ the code is in the right place but too certain: $\frac{1}{2}(0+0.01-1)-\log 0.1=-0.495+2.303=1.81$.

The two terms pull in opposite directions. Reconstruction wants each input's cloud to be tiny and far from the others, so codes are unambiguous. KL wants every cloud to be wide and centred on zero. The compromise packs the codes of all the data together around the origin, with overlapping clouds, so the space between training examples is filled with codes the decoder was actually trained on.

Step 3: The reparameterization trick

There is a technical snag. The loss involves sampling $z$, and you cannot backpropagate through "draw a random number". The fix is to move the randomness outside the network:

Now $z$ is an ordinary differentiable function of $\mu$ and $\sigma$, with $ϵ$ as a fixed input for this step, so gradients flow from the decoder back into the encoder. Kingma and Welling's 2013 paper, Auto-Encoding Variational Bayes, made this reparameterization of the bound the centre of the method; Rezende, Mohamed and Wierstra published a closely related approach, stochastic backpropagation, in 2014. In practice the encoder outputs $\log {\sigma }^2$ rather than $\sigma$, so the spread is always positive.

Step 4: The ELBO

Why this particular loss? A VAE is a probabilistic model: draw $z$ from the prior, then draw $x$ from the decoder's distribution $p(x\mid z)$. Training it properly would maximize $\log p(x)$ for the data, but that requires integrating over every possible $z$, which is intractable. The negative of the loss above is a lower bound on $\log p(x)$, the evidence lower boundEvidence lower bound: a quantity that is always at most log p(x) and can be estimated and optimized; maximizing it trains both the encoder and the decoder of a VAE. (ELBO):

The gap between the two sides is how far the encoder's $q(z\mid x)$ is from the true posterior over codes. So maximizing the ELBO does two useful things at once: it improves the model of the data, and it makes the encoder a better approximation of which codes could have produced each input. You do not need the derivation to use a VAE, but it explains the name: "variational" refers to fitting an approximate distribution $q$ by optimization.

Step 5: Sampling and interpolation

After training, throw the encoder away to generate: draw $z\sim \mathcal{N}(0,I)$ and decode it. Because training pulled all the codes toward that same normal, random draws land where the decoder has seen data. Walking a straight line between two codes decodes into a gradual change from one input to the other, since the space between them was trained too.

Try it

The same architecture as the autoencoder in the previous lesson (64, 64, 2, 64, 64 units), trained offline in numpy on the scikit-learn digits with the VAE loss; the page uses the real decoder's weights.

1. Compare the scatter with the autoencoder's. The codes now fill a round blob centred on zero, about one unit wide in each direction, matching the prior's circles; the plain autoencoder's codes were spread over a much larger, ragged area.
2. Press "Sample z from N(0, I)" several times. Almost every sample lands among training codes and decodes to a recognizable digit, which is exactly what random codes in the plain autoencoder's box failed to do.
3. Interpolate from 4 to 7, then do the same in the previous lesson's autoencoder. Here the strip changes from 4 to 7 directly; the plain autoencoder passes through something like a 1 on the way. Then try 1 to 0: with only two latent numbers, the 1s overlap the 7s, so the start already looks like a 7. A 2-D latent space is great for plotting and too small for 10 digit classes.

Going deeper

Code: the KL formula and why reparameterize

import numpy as np

rng = np.random.default_rng(0)

# 1. The KL term for one latent number: q = N(mu, sigma^2) against the prior p = N(0, 1)
def kl(mu, sigma):                       # closed form used in every VAE
    return 0.5 * (mu**2 + sigma**2 - 1) - np.log(sigma)

def log_normal(z, mu, sigma):
    return -0.5 * np.log(2 * np.pi * sigma**2) - (z - mu) ** 2 / (2 * sigma**2)

for mu, sigma in [(0.0, 1.0), (1.0, 1.0), (2.0, 0.5), (0.0, 0.1)]:
    z = rng.normal(mu, sigma, 200_000)   # Monte Carlo estimate of E_q[log q - log p]
    mc = np.mean(log_normal(z, mu, sigma) - log_normal(z, 0, 1))
    print(f"mu={mu:3.1f} sigma={sigma:3.1f}   KL formula {kl(mu, sigma):6.3f}   sampled {mc:6.3f}")

# 2. Why the reparameterization trick: gradient of E[f(z)] with respect to mu,
#    for f(z) = (z - 3)^2 and z ~ N(mu, sigma^2). The true answer is 2 * (mu - 3).
mu, sigma, n = 1.0, 0.8, 100
f = lambda z: (z - 3) ** 2
reparam, score = [], []
for trial in range(2000):
    eps = rng.normal(size=n)
    z = mu + sigma * eps                         # z is a function of mu: gradients pass through
    reparam.append(np.mean(2 * (z - 3)))         # d f(mu + sigma * eps) / d mu
    score.append(np.mean(f(z) * eps / sigma))    # score-function (REINFORCE) estimator
print(f"\ntrue gradient {2 * (mu - 3):.2f}")
for name, g in [("reparameterized", reparam), ("score function ", score)]:
    print(f"{name}: mean {np.mean(g):6.2f}   spread (std) {np.std(g):5.2f}")

mu=0.0 sigma=1.0   KL formula  0.000   sampled  0.000
mu=1.0 sigma=1.0   KL formula  0.500   sampled  0.500
mu=2.0 sigma=0.5   KL formula  2.318   sampled  2.322
mu=0.0 sigma=0.1   KL formula  1.808   sampled  1.809

true gradient -4.00
reparameterized: mean  -4.00   spread (std)  0.16
score function : mean  -4.00   spread (std)  0.96

Two checks of the pieces above, in plain numpy.

• The closed-form KL is right. For four Gaussians, the formula and a Monte Carlo average of $\log q-\log p$ over 200,000 samples agree: 0.000, 0.500, 2.318 against 2.322, and 1.808 against 1.809. The last case, $\mu =0$ and $\sigma =0.1$, shows that being too certain is penalized even when the mean is perfect.
• Reparameterized gradients are far less noisy. For $f(z)=(z-3{)}^2$ with $z\sim \mathcal{N}(1,0.8^2)$, the true gradient of $\mathbb{E}[f(z)]$ with respect to $\mu$ is -4. Both estimators average -4.00 over 2,000 trials of 100 samples, but the reparameterized one has a spread of 0.16 and the score-function estimator (the same REINFORCE idea used for policy gradients) 0.96, six times larger. Lower-variance gradients are why VAEs train with ordinary backpropagation.

Tuning the balance

Weighting the KL term by a factor $\beta$ trades the two goals off explicitly. Larger $\beta$ gives a smoother, more regular latent space and blurrier reconstructions; smaller $\beta$ gives sharper reconstructions and more holes. If the decoder is very powerful, it can learn to ignore $z$ entirely and the KL term drives every code to the prior: the codes carry no information. Watching the KL term during training catches this.

Two properties of the standard VAE matter in practice. The decoder is trained with a per-pixel loss on a single sample, so where several outputs are plausible it hedges toward their average, which tends to look soft. And the latent prior is simple by design. Later lessons in this module trade those properties for sharper samples, and latent diffusion systems use an autoencoder, sometimes with a slight VAE-style KL penalty, only for compression, leaving generation to a diffusion model.

Trade-offs

Common mistakes

• Sampling from the latent space of a plain autoencoder and expecting realistic outputs.
• Forgetting the reparameterization trick and sampling inside the graph, so the encoder gets no gradient.
• Predicting $\sigma$ directly with no positivity constraint; predict $\log {\sigma }^2$ instead.
• Reporting only the reconstruction loss. Track the KL term too: near zero means the latent codes are being ignored.
• Using the encoder's random sample $z$ when you want a deterministic code for a downstream task; use $\mu$.

Where you will see this

VAEs are used where a smooth, sampleable latent space matters more than the sharpest samples: generating and interpolating molecules or designs, anomaly detection with a likelihood-style score, and learning compact representations of images, audio or sensor data. Their most visible role today is as the compression stage of latent diffusion image generators: an autoencoder (in one of the two published variants, trained with a slight VAE-style KL penalty) maps images to a smaller latent grid, and the diffusion model works there, as the diffusion lesson in this module shows.

Check yourself

Recap