Variational autoencoders

Meera had a marketing problem. The Christmas campaign needed posters of people enjoying plum cake, and the photographer’s quote was ₹1.8 lakh. “Can AI not make up some faces?” she asked. “Not real customers. I don’t want anyone’s photo used without asking.”

“It can,” said Priya, “and that instinct about consent is exactly right. But let’s be careful. The datasets people learn face generation on have their own rules. Arjun, the Pune GenAI Jam is this weekend. Build a face generator there as a learning project, and we’ll talk about what Crumb could actually use afterwards.”

Arjun had a head start, or thought he did. Back in Module 27 he had built an autoencoder for handwritten digits: a network that squeezes an image down to a handful of numbers and rebuilds it. “Squeeze faces down, then feed in new numbers to get new faces,” he said. “How hard can it be?”

Arjun’s autoencoder makes monsters

An autoencoder has two halves. The encoder squeezes an image into a short list of numbers called a latent code (latent means hidden: these numbers are never shown to anyone, the network invents them). The decoder rebuilds the image from that code. Training rewards one thing: the rebuilt image should match the original, pixel by pixel.

At the hackathon, Arjun swapped his digit autoencoder for a version that handled 64 × 64 colour images and trained it on CelebA, a public research dataset of about 200,000 celebrity face photos. After an hour on a rented GPU, the reconstructions looked decent: blurry, but clearly the same person, same hair, same smile.

Then the real test. He skipped the encoder, made up a random latent code (128 random numbers from a bell curve), and asked the decoder for a face.

He got a brown smear with two dark blotches where eyes might go. He tried again: a beige rectangle with noise. Ten tries, ten monsters.

“The decoder works,” he said. “I’ve watched it rebuild a thousand faces. Why does it fall apart on my numbers?”

Priya’s map of Pune

Priya opened a map of Pune on her laptop. “Imagine every face in the dataset gets an address on this map. That is the latent code: a location.”

“OK.”

“Your autoencoder was only ever rewarded for one thing: rebuilding faces from their own addresses. Did anything tell it where to put the addresses?”

“No.”

“So it put them wherever was convenient. A cluster here in Kothrud, a long thin line along the river, a few scattered in Hinjewadi. And between them?”

“Nothing. Empty land.”

“Now you pick a random point on the map.”

Arjun saw it. “I land in the middle of Pashan Lake. The decoder has never been asked to draw anything from there. It has no idea what lives at that address, so it makes something up.”

“That’s your monster. An autoencoder is a perfect filing system and a terrible map. To generate, you need a map with no empty land, where every location is close to some real faces, so a random point decodes into something sensible.”

Clouds, not pins

“So how do you fill the empty land?” Priya asked.

Arjun thought. “Make the addresses spread out more?”

“Two changes. First: each face no longer gets a pin, an exact point. It gets a cloud: a fuzzy neighbourhood. During training, we pick a random spot inside the cloud and ask the decoder to rebuild the face from there. Now the decoder learns that a whole area means ‘this face’, not one exact point.”

“And second?”

“Second, we charge rent. Every cloud is pulled towards the centre of the map and encouraged to be a standard, generous size. Clouds that run off to the edge, or shrink back into tiny pins, pay a penalty.”

“So two forces fight.”

“Exactly. Rebuilding faces well wants each cloud small and separate, so faces don’t get confused. The rent wants every cloud big, round and near the centre. The compromise is a map where neighbouring clouds overlap, the gaps fill in, and the whole middle of the map is covered with faces. Then you sample a random point near the centre, and you’re always somewhere sensible, often between two faces, which gives you a blend.”

That is a variational autoencoder, or VAE, proposed in 2013. The clouds are probability distributions, the rent is a quantity called the KL divergence, and the trick that makes training work is called the reparameterisation trick. We will meet each precisely, after following one face through the machine.

Where the map analogy breaks: the real latent space is not two-dimensional; Arjun’s has 128 dimensions. Our intuitions about “nearby” and “empty land” hold up surprisingly well, but you cannot look at the whole thing at once.

Following one face through the VAE

To keep the numbers readable, imagine a VAE with a 2-number latent space (real ones use 64 to 512). One CelebA photo, a smiling woman with dark hair, goes in.

Step 1: The encoder describes a cloud, not a point. Instead of outputting one code, the encoder outputs two lists: a centre $\mu$ (mu, the average) and a spread. For maths reasons it outputs the log-variance, $\log {\sigma }^2$, rather than the spread $\sigma$ (sigma) itself. Say it outputs:

• $\mu =[0.5,-1.0]$
• $\log {\sigma }^2=[-0.69,0.0]$, which means ${\sigma }^2=[0.5,1.0]$, so $\sigma =[0.707,1.0]$

This face’s cloud is centred at (0.5, -1.0), a bit narrower in the first direction than the second.

Step 2: Pick a random spot in the cloud. Draw two random numbers $ϵ$ (epsilon) from a standard bell curve, say $ϵ=[0.2,-1.5]$. Then:

$$z=\mu +\sigma \cdot ϵ=[0.5+0.707\times 0.2,-1.0+1.0\times (-1.5)]=[0.641,-2.5]$$

Step 3: Decode. The decoder turns $z=[0.641,-2.5]$ into a 64 × 64 image.

Step 4: Score the rebuild. The reconstruction loss compares the output with the original photo, pixel by pixel. Using squared error summed over all $64\times 64\times 3=12,288$ pixel values, a decent reconstruction might score around 60 (each pixel off by about 0.07 on a 0-to-1 scale).

Step 5: Charge the rent. The KL term measures how far this cloud is from the standard cloud at the centre of the map (centre 0, spread 1). Per dimension it is:

$$\text{KL}=-\frac{1}{2}\left(1+\log {\sigma }^2-{\mu }^2-{\sigma }^2\right)$$

• Dimension 1: $-\frac{1}{2}(1-0.69-0.25-0.5)=0.22$
• Dimension 2: $-\frac{1}{2}(1+0-1.0-1.0)=0.50$

Total rent: 0.72. Dimension 2 is perfectly sized but sits one unit off-centre, and pays 0.5 for it. Dimension 1 is nearer the centre but a bit too narrow.

Step 6: Add them up. Loss = reconstruction + KL ≈ 60 + 0.72 = 60.72. Training adjusts both networks to shrink this total, averaged over thousands of faces.

Now watch the rent fight back against cheating. Suppose the encoder tries to turn the cloud back into a pin to get a sharper rebuild, shrinking $\log {\sigma }^2$ in dimension 1 to -10. That dimension’s rent becomes $-\frac{1}{2}(1-10-0.25-0.00005)\approx 4.6$, twenty times more. Or suppose it shoves the face far away to $\mu =5$ to keep it separate: the rent for that dimension becomes $-\frac{1}{2}(1+0-25-1)=12.5$. Pins and remote addresses are expensive. Overlapping clouds near the centre are cheap.

Opening the hood: what the maths is really saying

Now the precise version.

The pieces

• The encoder outputs a probability distribution over codes for each image $x$: a Gaussian (bell curve) with centre $\mu (x)$ and spread $\sigma (x)$. This is written $q(z\mid x)=\mathcal{N}(\mu ,{\sigma }^2)$: “the cloud of likely codes for this image”.
• The prior, $p(z)=\mathcal{N}(0,I)$, is the standard cloud at the centre of the map: centre 0, spread 1 in every direction. It is where we will sample from when generating.
• The decoder turns a code $z$ back into an image.

The loss

The VAE loss for one image has the two forces you met:

$$\mathcal{L}(x)=\underset{\text{reconstruction}}{\underbrace{\Vert x-\text{decoder}(z){\Vert }^2}}+\beta \cdot \underset{\text{rent}}{\underbrace{\text{KL}(q(z\mid x)\Vert p(z))}}$$

KL divergence (Kullback-Leibler divergence) measures how different one probability distribution is from another. It is zero when they are identical and grows as they drift apart. For two Gaussians it has a neat closed form, summed over the $d$ latent dimensions:

$$\text{KL}=-\frac{1}{2}\sum _{j=1}^d\left(1+\log {\sigma }_j^2-{\mu }_j^2-{\sigma }_j^2\right)$$

That is exactly what you computed by hand. $\beta$ (beta) is a dial for how much rent to charge; the original VAE uses $\beta =1$.

The reparameterisation trick

There is a snag in Step 2. Training works by backpropagation: sending correction signals (gradients) backwards from the loss through every operation to every weight, as you saw in Module 27. But “pick a random point in the cloud” is a dice roll. You cannot ask a dice roll how its result would change if $\mu$ were a little bigger. The gradient path to the encoder is cut.

The trick is to move the dice roll outside the network. Roll $ϵ$ from a fixed standard bell curve first, as an ingredient handed in from outside. Then compute $z=\mu +\sigma \cdot ϵ$, which is ordinary arithmetic. Now if $\mu$ goes up by 0.1, $z$ goes up by exactly 0.1; if $\sigma$ doubles, $z$’s offset doubles. Gradients flow straight through to the encoder.

Priya’s version: “A chef who rolls dice in the middle of the recipe can’t tell you how the cake would change if she used more sugar, because the dice keep changing everything. So roll the dice before you start, write the number on the order slip, and follow a fixed recipe. Same randomness, but now every step is traceable.”

Why output $\log {\sigma }^2$ instead of $\sigma$? A spread must be positive, and a network’s raw outputs can be any number. Treating the output as a log and exponentiating it guarantees positivity and keeps training stable.

Generating new faces

After training, throw the encoder away. Sample $z$ from the prior, $\mathcal{N}(0,I)$, which is plain torch.randn, and decode. Because the rent packed every face’s cloud around the centre, random samples land among real faces’ clouds, and you get plausible new ones.

Building a face VAE on CelebA

You will train this on a free Colab GPU (Google Colab gives you a notebook with a GPU in the browser; choose Runtime → Change runtime type → T4 GPU). Locally, use uv add torch torchvision; in Colab, PyTorch is preinstalled.

Step 1: get the faces. CelebA has 202,599 aligned face photos of about 10,000 celebrities. Its licence allows non-commercial research use only, so it is fine for learning and not for Crumb’s posters. torchvision can download it, but the download is hosted on Google Drive and often fails with a quota error. If it does, download the aligned images (img_align_celeba.zip) from the dataset’s official page or a mirror such as Kaggle, unzip it into data/celeba/, and use ImageFolder, which treats the single subfolder as one class.

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

tfm = transforms.Compose([
    transforms.CenterCrop(148),   # crop to the face (images are 178 x 218)
    transforms.Resize(64),        # shrink to 64 x 64 so training is fast
    transforms.ToTensor(),        # pixels become numbers between 0 and 1
])

try:
    faces = datasets.CelebA("data", split="train", download=True, transform=tfm)
except Exception as err:          # Google Drive quota errors are common
    print("Download failed, using local folder instead:", err)
    faces = datasets.ImageFolder("data/celeba", transform=tfm)

loader = DataLoader(faces, batch_size=128, shuffle=True, num_workers=2, drop_last=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
print(len(faces), "faces")

You should see 162,770 faces for the official training split, or 202,599 if you used the whole folder.

Step 2: the model. The encoder is a convolutional network (the image-scanning layers from Module 28): each layer halves the image size, from 64 × 64 down to 4 × 4, then two small layers produce $\mu$ and $\log {\sigma }^2$. The decoder is its mirror, using transposed convolutions that double the size at each step.

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

LATENT = 128   # the number of dimensions on our map

class VAE(nn.Module):
    def __init__(self, latent=LATENT):
        super().__init__()
        # Encoder: 3x64x64 image -> 256x4x4 features, halving the size each layer
        self.enc = nn.Sequential(
            nn.Conv2d(3, 32, 4, stride=2, padding=1), nn.BatchNorm2d(32), nn.LeakyReLU(0.2),     # 32x32
            nn.Conv2d(32, 64, 4, stride=2, padding=1), nn.BatchNorm2d(64), nn.LeakyReLU(0.2),    # 16x16
            nn.Conv2d(64, 128, 4, stride=2, padding=1), nn.BatchNorm2d(128), nn.LeakyReLU(0.2),  # 8x8
            nn.Conv2d(128, 256, 4, stride=2, padding=1), nn.BatchNorm2d(256), nn.LeakyReLU(0.2), # 4x4
            nn.Flatten(),
        )
        self.to_mu = nn.Linear(256 * 4 * 4, latent)       # centre of the cloud
        self.to_logvar = nn.Linear(256 * 4 * 4, latent)   # log of its spread, squared
        # Decoder: the mirror image
        self.from_z = nn.Linear(latent, 256 * 4 * 4)
        self.dec = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1), nn.BatchNorm2d(128), nn.LeakyReLU(0.2),  # 8x8
            nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1), nn.BatchNorm2d(64), nn.LeakyReLU(0.2),    # 16x16
            nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1), nn.BatchNorm2d(32), nn.LeakyReLU(0.2),     # 32x32
            nn.ConvTranspose2d(32, 3, 4, stride=2, padding=1), nn.Sigmoid(),                               # 64x64
        )

    def encode(self, x):
        h = self.enc(x)
        return self.to_mu(h), self.to_logvar(h)

    def reparameterise(self, mu, logvar):
        std = torch.exp(0.5 * logvar)      # sigma = e^(logvar / 2)
        eps = torch.randn_like(std)        # the dice roll, made outside the network
        return mu + std * eps

    def decode(self, z):
        return self.dec(self.from_z(z).view(-1, 256, 4, 4))

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

reparameterise is the whole trick in three lines. The final Sigmoid keeps output pixels between 0 and 1, matching the input. The model has about 3 million parameters, small by modern standards.

Step 3: the loss. Both terms are summed over each image and then averaged over the batch, so they are on the same scale. This detail matters enormously, as Arjun is about to discover.

def vae_loss(recon, x, mu, logvar, beta=1.0):
    n = x.size(0)
    recon_loss = F.mse_loss(recon, x, reduction="sum") / n                  # per image
    kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / n        # per image
    return recon_loss + beta * kl, recon_loss, kl

The kl line is the formula you computed by hand, applied to all 128 dimensions of every image at once.

Step 4: train.

model = VAE().to(device)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(1, 11):
    for x, _ in loader:
        x = x.to(device)
        recon, mu, logvar = model(x)
        loss, rec, kl = vae_loss(recon, x, mu, logvar)
        opt.zero_grad(); loss.backward(); opt.step()
    print(f"epoch {epoch}: loss {loss.item():.1f}  recon {rec.item():.1f}  kl {kl.item():.1f}")
    torch.save(model.state_dict(), "vae_celeba.pt")   # save each epoch; Colab can disconnect

Each epoch passes over all 162,770 training faces, which takes several minutes on a T4. Expect reconstruction loss to fall quickly in the first epoch and then slowly, while the KL term rises from near zero as the encoder learns to put real information into the codes. Your exact numbers will depend on settings, but a healthy run settles with the KL term at a small fraction of the reconstruction term, in the low hundreds or below for this setup.

Step 5: generate, and walk the map.

from torchvision.utils import save_image

model.eval()
with torch.no_grad():
    z = torch.randn(64, LATENT, device=device)            # 64 random points near the centre
    save_image(model.decode(z), "new_faces.png", nrow=8)

    # Walk in a straight line between two real faces' codes
    x, _ = next(iter(loader))
    mu, _ = model.encode(x[:2].to(device))
    steps = torch.linspace(0, 1, 8, device=device).unsqueeze(1)
    path = mu[0] * (1 - steps) + mu[1] * steps            # 8 points from face A to face B
    save_image(model.decode(path), "morph.png", nrow=8)

new_faces.png should show 64 blurry but recognisable faces: varied hair, skin tones, angles, backgrounds, none of them real people. morph.png shows one face smoothly turning into another. That smooth walk is only possible because the VAE’s map has no empty land between them. Try the same walk with Arjun’s plain autoencoder and the middle frames turn into monsters.

Where VAEs break

The identical beige face. Arjun’s first VAE run produced the same beige, featureless face for every random sample. The bug: he had used F.mse_loss(recon, x) with its default reduction="mean", averaging the error over all 12,288 pixel values to a tiny number around 0.02, while summing the KL over 128 dimensions. Rent was charged in rupees and reconstruction scored in paise. The KL term dominated, every cloud collapsed onto the standard cloud, codes carried no information, and the decoder learned the average face. This failure has a name, posterior collapse, and mismatched loss scales are its most common cause. Sum both, or average both, consistently.

Blurriness. VAE images are soft. The squared-error loss punishes being far from the right pixel value but is very forgiving of being “roughly grey” everywhere. When the decoder is unsure whether a strand of hair goes here or there, the lowest-loss answer is a smeared average of both. It is the same reason a network trained to match random digits with squared error produces a grey blob. GANs, in the previous lesson, produce sharper images precisely because they are not judged by pixel distance.

The beta dial. Turning $\beta$ above 1 (a beta-VAE) makes the map smoother and the latent directions more interpretable, sometimes one direction for “smiling” and another for “hair colour”, at the cost of blurrier reconstructions. Below 1, reconstructions sharpen and random samples get worse. There is no free setting.

The data decides who exists. A generator can only produce faces like the ones it learned from. CelebA is a collection of celebrity photos and is known to skew towards young, conventionally attractive faces and to under-represent many groups. A poster generator trained on it would quietly reflect that. Combined with the non-commercial licence, this is why Priya called it a learning project, not a marketing tool.

Where you will meet VAEs again

VAEs are rarely the final image generator today, but their idea is everywhere. Stable Diffusion, which you will meet in Module 33, uses a VAE-style encoder to compress each 512 × 512 image into a much smaller latent grid, does all its expensive work on that grid, and uses the decoder to turn the result back into pixels. The map you built this weekend is literally the floor that modern image generators stand on.

Try it yourself

1. Compute the KL rent by hand for one dimension with $\mu =0$ and $\log {\sigma }^2=0$. Then for $\mu =2$, $\log {\sigma }^2=0$.
2. Train for one epoch with beta=0.1 and one with beta=4, saving new_faces.png from each. Which samples look sharper, and which look more varied and face-like?

Show the answer

1. With $\mu =0$ and $\log {\sigma }^2=0$ (so ${\sigma }^2=1$): $-\frac{1}{2}(1+0-0-1)=0$. This is exactly the standard cloud, so no rent. With $\mu =2$: $-\frac{1}{2}(1+0-4-1)=2$. Rent grows with the square of the distance from the centre.
2. With beta=0.1, reconstructions sharpen, but random samples look patchier and more monster-like, because the map has more empty land. With beta=4, samples are consistently face-like but blurrier and more alike. Same tug of war, different balance.

Common mistakes

• Mixing summed and averaged losses. The classic cause of identical, average-face samples. Put both terms on the same per-image scale.
• Sampling with torch.rand instead of torch.randn. rand gives numbers between 0 and 1, all positive, which is not the bell curve the prior uses. Samples come out skewed and odd.
• Forgetting model.eval() before generating. Batch normalisation behaves differently in training mode, which can wash out generated images.
• Expecting sharp photos. VAEs trade sharpness for a smooth, well-behaved map. If you need crisp images, use a GAN or a diffusion model.
• Ignoring the dataset licence and bias. Check what you are allowed to do with a dataset, and who is missing from it, before building anything public.

Recap

• A plain autoencoder makes a good filing system but a poor map: random codes land in empty regions and decode into garbage.
• A VAE encodes each image as a cloud (a centre $\mu$ and spread $\sigma$) and samples a point from it during training.
• The KL term charges rent for clouds that are off-centre or too narrow, so clouds overlap and fill the map; the reconstruction term makes codes carry information.
• The reparameterisation trick, $z=\mu +\sigma \cdot ϵ$, moves randomness outside the network so gradients can flow.
• VAEs give smooth latent spaces and stable training, at the cost of blurry images, and their encoders live on inside modern diffusion models.

On Monday, Arjun showed Meera the grid of 64 faces. “None of these people exist. The model learned a map of faces, and every point near the middle of the map is a face. Pick a random point, you get a new person. Walk from one point to another, and one face melts into the next.”

“They’re a bit soft,” she said.

“That’s the price of the map being smooth. And we can’t use these for posters: the photos it learned from are licensed for research only.” He hesitated. “Also, honestly, I ran this fourteen times with different settings, and I’ve lost track of which run produced which faces. My notes are a mess of printed numbers.”

Priya laughed. “Then before you train anything else, you need a dashboard. Next lesson: TensorBoard.”

Go deeper

• PyTorch tutorials, which include generative model examples.
• Hugging Face docs, see the Diffusers library for how VAEs are used inside Stable Diffusion.
• TensorFlow tutorials, which include a convolutional VAE walkthrough.