Module 3 — Variational autoencoders

The autoencoder of module 2 compresses well but generates poorly, because it never commits to a known distribution over its latent codes. A variational autoencoder — VAE — fixes that with two changes: the encoder produces a distribution rather than a point, and the loss adds a term that pulls that distribution towards a standard Gaussian. That is enough to turn the network into a generator that can draw new digits from noise. This module builds one on MNIST.

Encoder as a distribution, not a point

Instead of returning a single code $z = f(x)$, the VAE encoder returns two vectors: a mean $\mu(x)$ and a log-variance $\log \sigma^2(x)$. The latent variable is then sampled from

$z \sim \mathcal{N}(\mu(x), \sigma(x)^2 I).$

Two consequences follow immediately. First, the same $x$ can map to slightly different codes across calls: the encoder is now stochastic. Second, we can pull the encoder's distribution towards a prior we choose in advance, typically $\mathcal{N}(0, I)$, so that sampling from that prior at generation time makes sense.

import torch
import torch.nn as nn
import torch.nn.functional as F

class Encoder(nn.Module):
    def __init__(self, latent_dim=16):
        super().__init__()
        self.body = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 256), nn.ReLU(),
            nn.Linear(256, 128), nn.ReLU(),
        )
        self.mu = nn.Linear(128, latent_dim)
        self.log_var = nn.Linear(128, latent_dim)

    def forward(self, x):
        h = self.body(x)
        return self.mu(h), self.log_var(h)

The two linear heads share the trunk and produce two vectors of the same size. Predicting $\log \sigma^2$ rather than $\sigma$ is a numerical trick: taking exp(0.5 * log_var) yields a positive $\sigma$ without a constraint on the head, and gradients stay stable.

The reparameterisation trick

Now for the subtlety that made VAEs possible. During training we need gradients to flow through the sampling $z \sim \mathcal{N}(\mu, \sigma^2)$, so that reconstruction error can update $\mu$ and $\sigma$. But sampling is a stochastic operation and, on its own, has no derivative with respect to its parameters.

The reparameterisation trick rewrites the sample as a deterministic function of a fixed-shape noise:

$z = \mu + \sigma \cdot \varepsilon, \quad \varepsilon \sim \mathcal{N}(0, I).$

The randomness now sits in $\varepsilon$, which does not depend on the model. The path from $\mu$ to $z$ is a plain arithmetic operation, differentiable end to end. In code:

def reparameterise(mu, log_var):
    sigma = torch.exp(0.5 * log_var)
    eps = torch.randn_like(sigma)
    return mu + sigma * eps

This one line is what unlocked stochastic gradients for latent variable models. Without it we would have to fall back on high-variance score-function estimators, and VAEs would not train on anything larger than a toy example.

The trick does not sample from the prior

reparameterise(mu, log_var) samples from the posterior $q(z \mid x)$, not from the prior $p(z)$. At generation time we bypass the encoder entirely and draw $z \sim \mathcal{N}(0, I)$; at training time we must go through the encoder or the reconstruction term makes no sense.

Two terms in the loss: reconstruction and KL

The VAE loss combines a reconstruction term — the same pixel error we used in module 2 — with a KL divergence term that pulls $q(z \mid x)$ towards the prior:

$\mathcal{L}(x) = \underbrace{\| x - g(z) \|^2}_{\text{reconstruction}} + \beta \cdot \underbrace{D_{\text{KL}}\big( q(z \mid x) \,\|\, \mathcal{N}(0, I) \big)}_{\text{KL}}$

The negative of this quantity is a lower bound on $\log p(x)$, known as the ELBO — the Evidence Lower BOund. Maximising the ELBO therefore maximises a bound on the data likelihood.

For a Gaussian $q(z \mid x)$ against a standard Gaussian prior, the KL has a closed form:

$D_{\text{KL}} = -\frac{1}{2} \sum_{i=1}^{d} \big( 1 + \log \sigma_i^2 - \mu_i^2 - \sigma_i^2 \big)$

which translates directly:

def kl_divergence(mu, log_var):
    return -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp(), dim=1)

The two terms pull in opposite directions. Reconstruction wants the encoder to place each $x$ at a very specific code — the one the decoder reconstructs best. KL wants every code to be close to zero with unit variance, which erases individuality. The balance $\beta$ controls that tension. $\beta = 1$ is the standard VAE; $\beta > 1$ produces a disentangled latent space at the cost of reconstruction, and is known as the $\beta$-VAE.

The full model on MNIST

Wiring everything together yields a training loop that runs in a few minutes on a laptop:

class VAE(nn.Module):
    def __init__(self, latent_dim=16):
        super().__init__()
        self.encoder = Encoder(latent_dim)
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 128), nn.ReLU(),
            nn.Linear(128, 256), nn.ReLU(),
            nn.Linear(256, 28 * 28), nn.Tanh(),
        )

    def forward(self, x):
        mu, log_var = self.encoder(x)
        z = reparameterise(mu, log_var)
        x_hat = self.decoder(z).view(-1, 1, 28, 28)
        return x_hat, mu, log_var

model = VAE().to("cuda" if torch.cuda.is_available() else "cpu")
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(20):
    for x, _ in loader:  # loader from module 2
        x = x.to(next(model.parameters()).device)
        x_hat, mu, log_var = model(x)
        recon = F.mse_loss(x_hat, x, reduction="sum") / x.size(0)
        kl = kl_divergence(mu, log_var).mean()
        loss = recon + 1.0 * kl
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

To generate, we ignore the encoder and pass random codes through the decoder:

with torch.no_grad():
    z = torch.randn(64, 16, device=next(model.parameters()).device)
    samples = model.decoder(z).view(-1, 1, 28, 28)

The samples are recognisable digits, but they look soft. Every stroke is a little blurred, and shapes that could be a "9" or a "4" often end up somewhere in between.

Why VAE reconstructions blur

The blur is not a training bug: it is a direct consequence of the objective. The decoder receives a stochastic $z$ that scatters around $\mu(x)$ during training, and it minimises average squared error over that neighbourhood. The average of several sharp reconstructions is a blurry reconstruction — the same mechanism that makes squared-error regression blur images in general.

Diffusion models of module 6 avoid this by training on a very different target — noise prediction — that has no such averaging effect. GANs of module 4 avoid it by throwing likelihood away entirely and training on a discriminator's judgement of realism.

Latent traversal reveals what each dimension controls

Fix all coordinates of $z$ except one, vary that one across a range and decode. Each dimension often controls a coherent factor: slant, thickness, digit identity. On a plain autoencoder from module 2 these traversals rarely mean anything; on a VAE they do, because the KL term regularised the space.

In summary

• A VAE encoder returns $(\mu, \log \sigma^2)$; sampling uses the reparameterisation trick $z = \mu + \sigma \varepsilon$ with $\varepsilon \sim \mathcal{N}(0, I)$, which makes the operation differentiable.
• The loss is reconstruction plus KL divergence against the prior $\mathcal{N}(0, I)$; together they form the ELBO, a lower bound on $\log p(x)$.
• To generate, draw $z \sim \mathcal{N}(0, I)$ and pass it through the decoder, bypassing the encoder — the prior and the aggregate posterior now agree, so the samples are plausible.
• Reconstructions are blurry because the decoder minimises squared error over a stochastic $z$; the diffusion models of module 6 sidestep this by predicting noise rather than pixels.

Next module: GANs, which abandon likelihood for an adversarial game and produce sharper samples — at the cost of a training procedure that is much harder to keep alive.