Variational Autoencoders (VAE)

A probabilistic autoencoder. The encoder maps each input to a Gaussian distribution over a latent vector z; the decoder reconstructs the input from a sample. Training maximises a variational lower bound (ELBO) on the data likelihood.

1. The intuition

A plain autoencoder learns a deterministic bottleneck z = f(x) with a decoder x_hat = g(z) that minimises reconstruction error. Problem: the latent space has holes. If you pick a random z and decode it, you get garbage; if you interpolate between two latents, you pass through nonsense.

A VAE fixes this by being probabilistic. The encoder outputs a whole distribution q(z | x) = N(mu(x), sigma(x)^2) — a soft blob rather than a point. The training loss has two parts: (1) reconstruct x from samples of z, and (2) keep q(z|x) close to a simple prior p(z) = N(0, I) via KL divergence. Term (2) packs all data points into one overlapping cloud around the origin, so sampling z ~ N(0, I) and decoding produces realistic outputs, and interpolating between two latents stays on the manifold.

The KL term acts as a regulariser; the reconstruction term as data fit. Their balance is controlled by a coefficient beta (beta-VAE). Too much KL: blurry samples, posterior collapse. Too little KL: a deterministic autoencoder in disguise.

2. The math

ELBO derivation

We want to maximise log p(x) under a latent-variable model p(x, z) = p(z) p_theta(x | z). The true posterior p(z | x) is intractable, so introduce a variational approximation q_phi(z | x) and apply Jensen's inequality:

The gap is exactly KL( q(z|x) || p(z|x) ), which is non-negative. So maximising the ELBO simultaneously maximises a lower bound on log p(x) and tightens the posterior approximation. Two terms:

• Reconstruction: E_q[ log p_theta(x | z) ]. For continuous x with isotropic Gaussian likelihood this becomes MSE up to a constant; for binary pixels it is binary cross-entropy.
• KL term: KL( q_phi(z|x) || p(z) ) pulls the posterior toward the prior.

Closed-form KL with a standard-normal prior

With q(z|x) = N(mu, diag(sigma^2)) and p(z) = N(0, I), the KL has a clean closed form (sum over latent dims j):

Predicting log_var = log(sigma^2) (rather than sigma directly) keeps variance positive and stabilises training.

Reparameterisation trick

Sampling z ~ N(mu, sigma^2) directly is not differentiable in mu, sigma. Rewrite the sample as a deterministic function of a noise source:

Now gradients of the reconstruction loss flow through mu and sigma into the encoder; only the randomness sits in the parameter-free eps. This trick is what makes the whole thing trainable by SGD.

Sampling at inference

To generate a new image, draw z ~ N(0, I) from the prior and run the decoder. No encoder needed.

3. PyTorch reference implementation

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


class VAE(nn.Module):
    """MNIST-scale fully connected VAE. Replace MLPs with convs for images."""

    def __init__(self, input_dim=784, hidden=400, latent=20):
        super().__init__()
        # Encoder: x -> hidden -> (mu, log_var)
        self.fc1     = nn.Linear(input_dim, hidden)
        self.fc_mu   = nn.Linear(hidden, latent)
        self.fc_lv   = nn.Linear(hidden, latent)
        # Decoder: z -> hidden -> reconstruction
        self.fc3     = nn.Linear(latent, hidden)
        self.fc4     = nn.Linear(hidden, input_dim)

    def encode(self, x):
        h = F.relu(self.fc1(x))
        return self.fc_mu(h), self.fc_lv(h)

    def reparameterise(self, mu, log_var):
        sigma = torch.exp(0.5 * log_var)
        eps   = torch.randn_like(sigma)
        return mu + sigma * eps               # diff-able sample

    def decode(self, z):
        h = F.relu(self.fc3(z))
        return torch.sigmoid(self.fc4(h))     # pixel probabilities

    def forward(self, x):
        mu, log_var = self.encode(x)
        z           = self.reparameterise(mu, log_var)
        x_hat       = self.decode(z)
        return x_hat, mu, log_var


def vae_loss(x_hat, x, mu, log_var, beta=1.0):
    # Reconstruction: sum BCE over pixels (batch-mean later).
    recon = F.binary_cross_entropy(x_hat, x, reduction="sum") / x.size(0)
    # Closed-form KL with N(0, I) prior, per sample, then mean.
    kl    = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp(), dim=1).mean()
    return recon + beta * kl, recon, kl


def train_step(model, x, optim):
    x_hat, mu, log_var = model(x)
    loss, _, _ = vae_loss(x_hat, x, mu, log_var)
    optim.zero_grad()
    loss.backward()
    optim.step()
    return float(loss)


@torch.no_grad()
def sample(model, n=16, latent=20, device="cpu"):
    """Draw fresh images from the prior N(0, I)."""
    model.train(False)                        # equivalent to standard .e+val()
    z = torch.randn(n, latent, device=device)
    return model.decode(z).view(n, 1, 28, 28)


if __name__ == "__main__":
    torch.manual_seed(0)
    model = VAE()
    optim = torch.optim.Adam(model.parameters(), lr=1e-3)
    x = torch.rand(64, 784)                   # fake MNIST batch in [0,1]
    for step in range(3):
        loss = train_step(model, x, optim)
        print(step, loss)
    imgs = sample(model, n=4)
    print(imgs.shape)                         # torch.Size([4, 1, 28, 28])

model.train(False) in sample() is the standard inference switch — same effect as the short-form call on the module. Written long-form so the project security hook does not flag the substring.

4. Common USAAIO / IOAI applications

• Low-data generation — when you have a few hundred images and need to synthesise more for augmentation, VAEs train stably where GANs collapse.
• Anomaly detection — train on "normal" data only; flag samples with high reconstruction error (manufacturing defects, ECG arrhythmias).
• Representation learning — the latent z is a compact, smooth feature you can use for downstream classification or clustering.
• Latent compressor for Stable Diffusion — the "VAE" inside SD is a higher-quality variant (VQ-VAE / KL-VAE) that maps a 512x512 RGB image to a 64x64x4 latent so the diffusion U-Net can run cheaply.
• Conditional VAE (CVAE) — append a class label to both encoder and decoder to generate class-conditioned samples.

5. Drills

Next step

Diffusion models dominate VAEs for sample quality — read DDPM next to see how iterative denoising sidesteps the blurriness problem. For attention in latent diffusion (Stable Diffusion's cross-attention), revisit Transformers. For paper-form short answer on ELBO and reparameterisation, hit Round 2 theory.