Implementing VAEs: A Step-by-Step Guide

A variational autoencoder fits on one screen of code, and that is part of the problem. A short script that trains without errors can still produce a model that ignores its latent code, weighs its loss terms wrongly or generates only grey smudges, and nothing in the output tells you which. The way to avoid that is to build the model in small steps and check each one before moving on.

This guide builds a convolutional VAE for 28 by 28 grayscale images in PyTorch, from loading data to saving a checkpoint. The mathematical foundations of VAEs article explains where the loss comes from; here the focus is on getting a working model and knowing that it works. If the encoder, decoder and latent space are new to you, start with the introduction to variational autoencoders.

What you will build

The model has the standard VAE shape: an encoder that maps an image to the mean and log-variance of a Gaussian over a 16-dimensional latent code, a sampling step, and a decoder that maps a code back to an image. Training minimizes the negative evidence lower bound (ELBO), a reconstruction term plus a KL term.

flowchart LR
    X["Image 1x28x28"] --> C1["Conv 32, stride 2"]
    C1 --> C2["Conv 64, stride 2"]
    C2 --> FC1["Linear to 2 x 16"]
    FC1 --> MU["mu"]
    FC1 --> LV["log-variance"]
    MU --> Z["z = mu + sigma * eps"]
    LV --> Z
    Z --> FC2["Linear to 64x7x7"]
    FC2 --> D1["ConvTranspose 32"]
    D1 --> D2["ConvTranspose 1"]
    D2 --> L["Logits 1x28x28"]

The steps, in order:

1. Set up the environment and load the data.
2. Write the encoder and decoder and check their output shapes.
3. Write the loss and check its value on an untrained model.
4. Train with a short KL warm-up.
5. Evaluate: test loss, KL per latent dimension and active units.
6. Generate samples, interpolate between images and save the model.

Step 1: environment and data

You need Python, PyTorch and torchvision. A GPU speeds things up, but this model is small enough to train on a laptop CPU in reasonable time. Follow the PyTorch install selector to pick the right build for your machine, then:

python -m venv .venv
source .venv/bin/activate
pip install torch torchvision

The examples use MNIST, 60,000 training and 10,000 test images of handwritten digits. Fashion-MNIST has the same shape and drops in by changing one class name, which is a useful second test because its textures are harder to reconstruct. If you need a refresher on project layout and tooling, see setting up a simple ML project.

import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"

transform = transforms.ToTensor()   # uint8 [0, 255] -> float [0, 1], shape (1, 28, 28)
train_set = datasets.MNIST("data", train=True, download=True, transform=transform)
test_set = datasets.MNIST("data", train=False, download=True, transform=transform)

train_loader = DataLoader(train_set, batch_size=128, shuffle=True, drop_last=True)
test_loader = DataLoader(test_set, batch_size=512)

ToTensor already scales pixels to [0, 1], which is what the Bernoulli decoder below expects. Do not add the usual Normalize step with a dataset mean and standard deviation. That shifts pixel values outside [0, 1], and binary cross-entropy against those targets no longer means anything.

Step 2: encoder and decoder

The encoder uses two strided convolutions to go from 28 by 28 to 7 by 7, then a single linear layer that outputs both the mean and the log-variance. Predicting log-variance rather than variance or standard deviation keeps the output unconstrained: any real number is a valid log-variance, so you need no activation function to keep it positive.

class Encoder(nn.Module):
    def __init__(self, z_dim=16):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=4, stride=2, padding=1),   # 28x28 -> 14x14
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),  # 14x14 -> 7x7
            nn.ReLU(),
            nn.Flatten(),                                           # 64 * 7 * 7 = 3136
        )
        self.fc = nn.Linear(64 * 7 * 7, 2 * z_dim)

    def forward(self, x):
        mu, logvar = self.fc(self.conv(x)).chunk(2, dim=1)
        return mu, logvar

The decoder mirrors it with transposed convolutions. With kernel size 4, stride 2 and padding 1, each transposed convolution exactly doubles the height and width, so 7 becomes 14 and then 28 without any cropping.

class Decoder(nn.Module):
    def __init__(self, z_dim=16):
        super().__init__()
        self.fc = nn.Linear(z_dim, 64 * 7 * 7)
        self.deconv = nn.Sequential(
            nn.ReLU(),
            nn.Unflatten(1, (64, 7, 7)),
            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),  # 7 -> 14
            nn.ReLU(),
            nn.ConvTranspose2d(32, 1, kernel_size=4, stride=2, padding=1),   # 14 -> 28
        )

    def forward(self, z):
        return self.deconv(self.fc(z))  # logits, shape (N, 1, 28, 28)

The decoder’s last layer has no sigmoid. It returns logits, one per pixel, and the loss applies the sigmoid internally in a numerically stable way. You apply torch.sigmoid yourself only when you want to look at an image.

The model ties the two together with the reparameterization trick, z = μ + σ ⊙ ε, so gradients can flow back through the sampling step into the encoder:

class VAE(nn.Module):
    def __init__(self, z_dim=16):
        super().__init__()
        self.z_dim = z_dim
        self.encoder = Encoder(z_dim)
        self.decoder = Decoder(z_dim)

    def forward(self, x):
        mu, logvar = self.encoder(x)
        std = torch.exp(0.5 * logvar)
        z = mu + std * torch.randn_like(std)
        return self.decoder(z), mu, logvar

Before writing any training code, push one batch through and check the shapes:

model = VAE(z_dim=16).to(device)
x, _ = next(iter(train_loader))
logits, mu, logvar = model(x.to(device))
print(logits.shape, mu.shape, logvar.shape)
# torch.Size([128, 1, 28, 28]) torch.Size([128, 16]) torch.Size([128, 16])

If the logits are not exactly the shape of the input, the loss will either fail or, worse, broadcast silently into something meaningless. Catch it here.

Choosing the latent size

Sixteen dimensions is a reasonable default for MNIST. Fewer dimensions give a smoother, more interpretable latent space but blurrier reconstructions; more give sharper reconstructions, but past a point the extra dimensions go unused, which step 5 shows you how to measure. A 2-dimensional latent space is worth one run on its own, because you can plot every test image’s μ on a scatter chart and see how the digits cluster.

Step 3: the loss

The negative ELBO has two parts. The reconstruction term is the negative log-likelihood of the input under the decoder’s Bernoulli distribution, which is binary cross-entropy. The KL term measures how far the encoder’s Gaussian is from the standard normal prior, and for two Gaussians it has a closed form.

def vae_loss(logits, x, mu, logvar, beta=1.0):
    recon = F.binary_cross_entropy_with_logits(logits, x, reduction="none")
    recon = recon.flatten(1).sum(dim=1)                       # per example
    kl = 0.5 * (mu.pow(2) + logvar.exp() - logvar - 1)        # per dimension
    kl = kl.sum(dim=1)                                        # per example
    loss = (recon + beta * kl).mean()
    return loss, recon.mean(), kl.mean()

The order of reductions is the most common source of bugs in VAE code. Sum over pixels and over latent dimensions first, so each example gets one reconstruction value and one KL value in the same units, then average over the batch. If you use the default reduction="mean" in the cross-entropy, you divide the reconstruction term by 784 while the KL term stays a sum. The KL then dominates, the encoder collapses onto the prior and the samples turn into the average digit.

Returning the two terms separately costs nothing and is the most useful diagnostic you have, so log both throughout training.

A sanity check before training

An untrained model has a predictable loss. The decoder’s logits start close to zero, which means a predicted probability of about 0.5 for every pixel. Binary cross-entropy at probability 0.5 is ln 2, about 0.693, whatever the target, so the reconstruction term should start near 784 × 0.693, roughly 543. The KL term should start small, since μ and log-variance are both near zero.

loss, recon, kl = vae_loss(logits, x.to(device), mu, logvar)
print(f"loss {loss:.1f}  recon {recon:.1f}  kl {kl:.2f}")

If the reconstruction value is closer to 0.7, you are averaging over pixels. If it is in the tens of thousands, you are summing over the batch as well. Both are the reduction bug from above.

Step 4: training with a KL warm-up

The training loop is ordinary PyTorch, covered in general terms in deep learning basics. The one VAE-specific addition is a warm-up on β: it starts at 0 and rises linearly to 1, so the encoder learns to put information into z before the KL penalty applies in full. This is the KL annealing remedy for posterior collapse, introduced by Bowman et al. for text VAEs. A plain convolutional decoder on MNIST rarely collapses even without it, but the warm-up costs nothing and the habit pays off once your decoder gets stronger.

flowchart TD
    A["Next batch x"] --> B["Forward: logits, mu, logvar"]
    B --> C["beta = min(1, step / warmup)"]
    C --> D["loss = recon + beta * kl"]
    D --> E["Backward and Adam step"]
    E --> F{"End of epoch?"}
    F -->|"No"| A
    F -->|"Yes"| G["Log loss, recon, kl"]
    G --> H["Evaluate on test set"]
    H --> A

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
epochs = 20
warmup_steps = 2 * len(train_loader)   # reach beta = 1 after two epochs
step = 0

for epoch in range(1, epochs + 1):
    model.train()
    for x, _ in train_loader:
        x = x.to(device)
        beta = min(1.0, step / warmup_steps)
        logits, mu, logvar = model(x)
        loss, recon, kl = vae_loss(logits, x, mu, logvar, beta)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        step += 1
    print(f"epoch {epoch:2d}  beta {beta:.2f}  loss {loss.item():7.2f}  "
          f"recon {recon.item():7.2f}  kl {kl.item():6.2f}")

Two things to keep in mind when you read the log:

• During warm-up the total loss is not the ELBO. With β below 1, the KL term is discounted, so the loss looks better than the model really is. Compare runs only on losses computed with β = 1, which is what the evaluation in the next step does.
• The KL term usually rises before it settles. Early on, the encoder learns to encode more about each image, which moves q away from the prior. A KL that falls steadily toward zero while reconstruction stalls is the warning sign.

Step 5: evaluate the model

A falling training loss tells you the optimizer is working. It does not tell you whether the model generalizes or whether it uses its latent code. Evaluate on the test set with β fixed at 1, and track two more numbers.

@torch.no_grad()
def evaluate(model, loader):
    model.eval()
    total, recon_sum, kl_sum, n = 0.0, 0.0, 0.0, 0
    mus = []
    for x, _ in loader:
        x = x.to(device)
        logits, mu, logvar = model(x)
        loss, recon, kl = vae_loss(logits, x, mu, logvar)   # beta = 1
        b = x.size(0)
        total += loss.item() * b
        recon_sum += recon.item() * b
        kl_sum += kl.item() * b
        n += b
        mus.append(mu.cpu())
    mus = torch.cat(mus)
    active = (mus.var(dim=0) > 0.01).sum().item()
    return total / n, recon_sum / n, kl_sum / n, active

neg_elbo, recon, kl, active = evaluate(model, test_loader)
print(f"test -ELBO {neg_elbo:.2f}  recon {recon:.2f}  kl {kl:.2f}  "
      f"active units {active}/{model.z_dim}")

The loss values are in nats per image, since the cross-entropy uses natural logarithms. Lower is better, and because the negative ELBO is an upper bound on the true negative log-likelihood, it is a fair basis for comparing two VAEs on the same data.

Active units counts the latent dimensions whose mean actually changes with the input. If a dimension’s μ has almost no variance across the test set, the encoder outputs the same value for every image and that dimension carries no information. The threshold of 0.01 comes from Burda et al., who introduced the measure. Expect fewer active units than z_dim on MNIST; that is the KL term pruning dimensions the model does not need, not a bug.

The KL per dimension tells a similar story from another angle and is cheap to compute:

@torch.no_grad()
def kl_per_dim(model, loader):
    model.eval()
    total, n = torch.zeros(model.z_dim), 0
    for x, _ in loader:
        mu, logvar = model.encoder(x.to(device))
        kl = 0.5 * (mu.pow(2) + logvar.exp() - logvar - 1)
        total += kl.sum(dim=0).cpu()
        n += x.size(0)
    return total / n

Dimensions with KL near zero are collapsed onto the prior. If every dimension is near zero, the whole model has collapsed, and you should look at the warm-up, the decoder’s strength and the loss reductions, in that order.

Reading the symptoms

Blurriness deserves a note. Some of it is inherent: a decoder that predicts independent per-pixel probabilities averages over the plausible images for a given code, and an average of digits is soft at the edges. That is one reason GANs produce sharper images, and the next article in this series covers training them.

Step 6: sample, interpolate and save

Generating new images uses only the decoder. Draw codes from the standard normal prior and decode them:

model.eval()
with torch.no_grad():
    z = torch.randn(64, model.z_dim, device=device)
    samples = torch.sigmoid(model.decoder(z))         # (64, 1, 28, 28)

To view them as a grid, torchvision.utils.save_image(samples, "samples.png", nrow=8) writes a PNG.

Interpolation is a direct test of whether the latent space is smooth. Encode two test images, walk in a straight line between their means, and decode each step:

with torch.no_grad():
    x, _ = next(iter(test_loader))
    mu, _ = model.encoder(x[:2].to(device))
    t = torch.linspace(0, 1, steps=10, device=device).unsqueeze(1)
    z = (1 - t) * mu[0] + t * mu[1]                   # (10, z_dim)
    frames = torch.sigmoid(model.decoder(z))

In a well-trained VAE, each frame looks like a plausible digit and the shape morphs gradually. A plain autoencoder without the KL term usually produces nonsense halfway between two codes, which is the clearest way to see what the KL term buys you.

Finally, save the weights rather than the whole model object. A state_dict is just tensors, and recent PyTorch releases load files with weights_only=True by default, which refuses arbitrary pickled objects. Storing the configuration next to the weights lets you rebuild the model later:

torch.save({"model": model.state_dict(), "z_dim": model.z_dim}, "vae.pt")

ckpt = torch.load("vae.pt", map_location=device, weights_only=True)
model = VAE(z_dim=ckpt["z_dim"]).to(device)
model.load_state_dict(ckpt["model"])

Passing weights_only=True explicitly keeps the behaviour the same on older versions. The PyTorch serialization notes explain why loading untrusted full pickles is unsafe.

Where to go from here

Once the basic model works, a few changes are worth trying, one at a time:

• Swap the dataset. Fashion-MNIST works with the same code. For colour images, change the input channels to 3 and consider a Gaussian decoder with a squared-error reconstruction term, keeping in mind that its scale relative to the KL then depends on the variance you assume.
• Grayscale versus binary targets. MNIST pixels are not strictly 0 or 1. Binary cross-entropy still trains well on them, but the loss is no longer an exact Bernoulli likelihood. If you want comparable numbers with published results, binarize the images, for example by sampling each pixel with its intensity as the probability.
• Condition on labels. Feeding the digit label to both encoder and decoder gives a conditional VAE that generates a chosen class, the same idea the series applies to GANs shortly.
• Change β on purpose. Values above 1 trade reconstruction quality for more disentangled latent dimensions, as described in the foundations article.

For a gentler end-to-end project before or after this one, building your first generative AI model walks through the full workflow, and the overview of generative models places VAEs alongside the other model families.

Further reading

• Auto-Encoding Variational Bayes, Kingma and Welling: the original VAE paper.
• Importance Weighted Autoencoders, Burda, Grosse and Salakhutdinov: tighter bounds and the active units measure.
• PyTorch documentation, in particular torch.nn, torch.distributions and the serialization notes.
• torchvision datasets, for MNIST, Fashion-MNIST and others.

Key takeaways

• Build a VAE in checked steps: shapes after the model, the starting loss after the loss function, then training.
• Predict log-variance, return logits from the decoder and use binary_cross_entropy_with_logits for stability.
• Sum the reconstruction loss over pixels and the KL over latent dimensions, then average over the batch. Averaging over pixels is the most common VAE bug.
• An untrained Bernoulli decoder should score about 543 nats per 28 by 28 image, 784 × ln 2. A very different number points to a reduction bug.
• A KL warm-up from β = 0 to 1 guards against posterior collapse, but only losses computed with β = 1 are comparable.
• Evaluate with the test negative ELBO, the KL per dimension and the count of active units, not the training loss alone.
• Generate by decoding prior samples, check smoothness by interpolating, and save a state_dict with its configuration.