72 - Variational AutoEncoder (2): Building It on MNIST and Generating New Digits

In the previous lesson you derived the ELBO and the KL formula; here you turn them into code. A complete VAE in PyTorch fits in about thirty lines: an encoder that produces μ and log σ², one line of reparametrisation, a decoder, and a loss that adds reconstruction and KL. The real work is knowing what to measure: how many nats of information the code carries, whether samples from N(0, I) really are digits, whether all ten digits show up, whether interpolating in latent space produces gradual transitions, and how many dimensions of the code are actually needed. Every number below comes from real runs, executed twice with identical outputs, on 8×8 handwritten digits and not on MNIST: the run did not download any dataset. The most instructive result is not the one I expected, and I tell it as it came.

The words of this lesson

This lesson uses the terms of the previous one (prior, posterior, ELBO, KL, reparametrisation, encoder, decoder) and adds a few. The logit of a pixel is the real number the decoder produces before turning it into a probability with the sigmoid. Binary cross-entropy (BCE) is the loss −[x·log p + (1 − x)·log(1 − p)] for a pixel with value x and predicted probability p: summed over the pixels it is the negative logarithm of the Bernoulli likelihood, the reconstruction term of the ELBO for black-and-white images. Nats per image is the unit in which we measure the loss (natural logarithm), summed over the pixels and averaged over the images. Sampling from the prior means drawing z ~ N(0, I) and passing it to the decoder, the operation that a plain autoencoder could not do in any meaningful way.

Latent interpolation decodes the points on the segment between the codes of two images. An active latent dimension is a component of z whose mean value really varies from one image to another. Finally, the judge is a digit classifier trained separately that we use to "look at" the generated images: it does not measure perceptual quality, but it lets us count diversity and consistency.

The analogy: the apprentice who became a pastry chef

In the previous lesson the apprentice (encoder) looked at a cake and proposed "which recipe card" had produced it, with a margin of uncertainty; the pastry chef (decoder) learned to bake the cake from the card. Now, with the work done, we dismiss the apprentice, draw a card from the N(0, I) jar and ask the pastry chef to bake a new cake. What we hope is that the jar and the recipes the pastry chef has really learned coincide: that is why the apprentice was forced, through the KL, never to propose cards outside the jar.

The model, line by line

The architecture is deliberately small: a single hidden layer of 128 units for the encoder, two linear "heads" that produce μ and log σ² with 10 dimensions, and a symmetric decoder that produces 64 logits. The data are the 1797 digits of scikit-learn's load_digits (values from 0 to 16, here brought into [0, 1]), split 70/30 with stratification: 1257 for training, 540 for testing.

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

torch.set_num_threads(2)
X, y = load_digits(return_X_y=True)                                # 1797 digits, 8x8 pixels, values 0..16
Xtr, Xte, ytr, yte = train_test_split(X / 16.0, y, test_size=0.3, random_state=0, stratify=y)
Xtr, Xte = torch.tensor(Xtr, dtype=torch.float32), torch.tensor(Xte, dtype=torch.float32)
ytr, yte = torch.tensor(ytr), torch.tensor(yte)
D, LATENT = 64, 10

class VAE(nn.Module):
    def __init__(self, latent=LATENT, hidden=128):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(D, hidden), nn.ReLU())
        self.mu = nn.Linear(hidden, latent)                      # mean of q(z|x)
        self.logvar = nn.Linear(hidden, latent)                  # log of the variance of q(z|x)
        self.dec = nn.Sequential(nn.Linear(latent, hidden), nn.ReLU(), nn.Linear(hidden, D))

    def encode(self, x):
        h = self.enc(x)
        return self.mu(h), self.logvar(h)

    def reparametrize(self, mu, logvar):
        eps = torch.randn_like(mu)                               # noise, independent of every parameter
        return mu + torch.exp(0.5 * logvar) * eps                # z = mu + sigma * eps

    def decode(self, z):
        return self.dec(z)                                       # LOGITS of the pixel probabilities

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

def vae_loss(logits, x, mu, logvar):
    """Negative ELBO per image, in nats: (Bernoulli reconstruction) + (KL to N(0, I)), both SUMMED, then averaged over the batch."""
    recon = F.binary_cross_entropy_with_logits(logits, x, reduction="sum") / x.size(0)
    kl = (0.5 * (mu.pow(2) + logvar.exp() - logvar - 1)).sum() / x.size(0)
    return recon + kl, recon, kl

print("parameters:", sum(p.numel() for p in VAE().parameters()))

parameters: 20564

VAE.__init__ builds the four parts: self.enc (64 → 128 with ReLU), the two heads self.mu and self.logvar (128 → 10 each) and self.dec (10 → 128 → 64). encode returns the pair (μ, log σ²). reparametrize is the key line: mu + exp(0.5·logvar)·eps, where exp(0.5·logvar) is σ (the square root of σ² = exp(logvar)) and eps = torch.randn_like(mu) is the noise, drawn on every call and independent of all the parameters, so z is differentiable with respect to mu and logvar. decode returns the logits, not the probabilities: the sigmoid is applied inside the loss (in a numerically stable way) or at use time. forward chains the three steps and also returns mu and logvar because the KL needs them. The program prints 20,564 parameters; the count by hand checks out: 8320 + 2·1290 + 1408 + 8256.

vae_loss is the negative of the ELBO. The reconstruction term F.binary_cross_entropy_with_logits(logits, x, reduction="sum") / x.size(0) sums the BCE over the 64 pixels and over the images in the batch, then divides by the number of images: it is nats per image. The KL is the formula you derived, summed over the 10 dimensions and averaged over the images. The crucial point: both terms are summed over components and averaged only over images, so their ratio is the one in the ELBO.

Training

The loop is the usual one: mini-batches of 64 images, vae(x), vae_loss, backward, step with Adam (learning rate 0.001), for 100 epochs. On the test set, with vae.eval() and torch.no_grad(), we measure the loss and its two parts; for brevity we print seven epochs.

torch.manual_seed(0)
vae = VAE()
opt = torch.optim.Adam(vae.parameters(), lr=1e-3)

@torch.no_grad()
def evaluate(model):
    model.eval()
    logits, mu, logvar = model(Xte)
    loss, recon, kl = vae_loss(logits, Xte, mu, logvar)
    model.train()
    return loss.item(), recon.item(), kl.item()

for epoch in range(1, 101):
    perm = torch.randperm(len(Xtr))
    for i in range(0, len(Xtr), 64):
        x = Xtr[perm[i:i + 64]]
        logits, mu, logvar = vae(x)
        loss, recon, kl = vae_loss(logits, x, mu, logvar)
        opt.zero_grad(); loss.backward(); opt.step()
    if epoch in (1, 2, 5, 10, 20, 50, 100):
        t_loss, t_recon, t_kl = evaluate(vae)
        print(f"epoch {epoch:3d}   test: -ELBO {t_loss:6.2f} = reconstruction {t_recon:6.2f} + KL {t_kl:5.2f}   (nats per image)")

epoch   1   test: -ELBO  35.61 = reconstruction  34.91 + KL  0.70   (nats per image)
epoch   2   test: -ELBO  31.21 = reconstruction  30.62 + KL  0.59   (nats per image)
epoch   5   test: -ELBO  28.20 = reconstruction  27.86 + KL  0.34   (nats per image)
epoch  10   test: -ELBO  26.64 = reconstruction  25.36 + KL  1.28   (nats per image)
epoch  20   test: -ELBO  25.78 = reconstruction  23.76 + KL  2.02   (nats per image)
epoch  50   test: -ELBO  24.98 = reconstruction  22.14 + KL  2.85   (nats per image)
epoch 100   test: -ELBO  24.56 = reconstruction  21.47 + KL  3.10   (nats per image)

The total loss falls from 35.6 to 24.6 nats per image. Reconstruction falls from 34.9 to 21.5; the KL starts at 0.70, drops to 0.34 at the fifth epoch and then slowly rises to 3.10. The pattern is instructive: at first the model minimises the loss by making the encoder say almost nothing (the KL is expensive and the decoder learns the mean image first), then, as the decoder improves, the encoder starts to "say more" about each image and pays for it in KL.

At the end of training the code carries about 3.1 nats (≈ 4.5 bits) of information about each digit, and the total loss of 24.56 nats corresponds to 35.4 bits per image, that is 0.55 bits per pixel. An honest caveat: the pixels are grey tones in [0, 1] used as BCE targets, not binarised; formally this is not a valid Bernoulli likelihood, but it is the most common convention in tutorials, and these numbers are not comparable with those published on binarised datasets.

How well it reconstructs and how much it looks like a digit

We measure reconstruction as MSE (the metric we already had), decoding the mean of the code with no noise. Then we draw 2000 codes from the prior, decode them, and measure the mean distance from the nearest training image (in pixel space; smaller = closer to a real digit), with three references: the real test digits, uniform noise, and the mean image.

# Reconstruction quality (MSE on the mean decoded image)
with torch.no_grad():
    logits, mu, logvar = vae(Xte)
    recon_from_mean = torch.sigmoid(vae.decode(mu))                   # decode the MEAN code, no noise
    print(f"VAE reconstruction MSE (latent {LATENT}):", round(F.mse_loss(recon_from_mean, Xte).item(), 4))

# Generation: decode z ~ N(0, I). Is it digit-like? Compare with real digits via distance to the nearest TRAINING image.
def nn_distance(images):
    return torch.cdist(images, Xtr).min(dim=1).values.mean().item()

with torch.no_grad():
    samples = torch.sigmoid(vae.decode(torch.randn(2000, LATENT)))
    real = Xte
    noise = torch.rand(2000, D)
    mean_img = Xtr.mean(0).expand(2000, D)
    print("mean distance to the nearest training image (pixel space):")
    print(f"  real test digits        : {nn_distance(real):.2f}")
    print(f"  VAE samples, z ~ N(0,I) : {nn_distance(samples):.2f}")
    print(f"  uniform noise images    : {nn_distance(noise):.2f}")
    print(f"  the mean image          : {nn_distance(mean_img):.2f}")

VAE reconstruction MSE (latent 10): 0.0282
mean distance to the nearest training image (pixel space):
  real test digits        : 1.08
  VAE samples, z ~ N(0,I) : 1.21
  uniform noise images    : 3.56
  the mean image          : 1.51

The VAE's MSE is 0.0282. A plain autoencoder with the same 10-number code (tried further below) gets 0.0112: the VAE reconstructs much worse, and that is the price of the KL regularisation, which prevents you from using the latent space as you please. As for the distances: a real test digit lies on average 1.08 from the nearest training digit, a uniform-noise image 3.56, the mean image 1.51. The VAE samples sit at 1.21: much closer to real digits than noise is, a little farther than a real digit, and closer than the mean image. This says something honest: the samples are plausible but blurry. The distance to the nearest neighbour does not measure perceptual quality, and a low value does not mean "pretty": it is a rough check.

Generating: z ~ N(0, I), with a judge

To find out whether all the digits appear, we train a small classifier (an MLP of 64 → 64 → 10, 60 epochs) and use it as a judge: which digit it sees in each sample, and how often. First we check that the judge is reliable on real digits: its accuracy on the test set is 97.8%. Then we apply it to the 2000 VAE samples, and to an interpolation between a 3 and a 7 from the test set (the means μ_a and μ_b of the two codes, nine points on the segment (1 − t)·μ_a + t·μ_b).

# A classifier is used as a "judge": which digit does it see, and how sure is it?
torch.manual_seed(0)
clf = nn.Sequential(nn.Linear(D, 64), nn.ReLU(), nn.Linear(64, 10))
opt_c = torch.optim.Adam(clf.parameters(), lr=1e-3)
for _ in range(60):
    perm = torch.randperm(len(Xtr))
    for i in range(0, len(Xtr), 64):
        idx = perm[i:i + 64]
        loss = F.cross_entropy(clf(Xtr[idx]), ytr[idx]); opt_c.zero_grad(); loss.backward(); opt_c.step()
with torch.no_grad():
    print("judge accuracy on real test digits:", round((clf(Xte).argmax(-1) == yte).float().mean().item(), 4))
    probs = F.softmax(clf(samples), dim=-1)
    hist = torch.bincount(probs.argmax(-1), minlength=10).float() / len(samples)
print("classes the judge sees in 2000 VAE samples:", [round(v, 2) for v in hist.tolist()])
print("entropy of that distribution (max ln 10 = 2.30):", round(-(hist[hist > 0] * hist[hist > 0].log()).sum().item(), 2))

# Interpolation in latent space between a "3" and a "7": does the path go through plausible digits?
i3 = (yte == 3).nonzero()[0].item(); i7 = (yte == 7).nonzero()[0].item()
with torch.no_grad():
    mu_a, _ = vae.encode(Xte[i3:i3 + 1]); mu_b, _ = vae.encode(Xte[i7:i7 + 1])
    for t in torch.linspace(0, 1, 9):
        img = torch.sigmoid(vae.decode((1 - t) * mu_a + t * mu_b))
        p = F.softmax(clf(img), dim=-1)
        print(f"t={t:.2f}  judge sees a {p.argmax().item()} with probability {p.max().item():.2f}")

judge accuracy on real test digits: 0.9778
classes the judge sees in 2000 VAE samples: [0.09, 0.07, 0.08, 0.16, 0.12, 0.1, 0.12, 0.09, 0.09, 0.08]
entropy of that distribution (max ln 10 = 2.30): 2.27
t=0.00  judge sees a 3 with probability 0.83
t=0.12  judge sees a 3 with probability 0.84
t=0.25  judge sees a 3 with probability 0.86
t=0.38  judge sees a 3 with probability 0.85
t=0.50  judge sees a 3 with probability 0.74
t=0.62  judge sees a 7 with probability 0.41
t=0.75  judge sees a 7 with probability 0.90
t=0.88  judge sees a 7 with probability 0.98
t=1.00  judge sees a 7 with probability 1.00

In 2000 samples the judge sees all ten digits, each with a frequency between 7% and 16%; the entropy of this distribution is 2.27 nats out of a maximum of 2.30 (ln 10): almost uniform. The VAE has no labels for the digits, but the latent space organised by the KL lets you draw from every "zone". The judge does tend to assign a digit even to blurry images, so I read the result as "no class is missing", not as "all of them are sharp". The interpolation is gradual and without holes in the sense that the judge sees a 3 up to t = 0.50 and a 7 from t = 0.62, without passing through other digits; the greatest uncertainty sits between the two, where the probability of the 7 at t = 0.62 is only 0.41. Note that even the endpoints are not certain: the 3 reconstructed from the mean of the code has probability 0.83, not 1, because the VAE reconstructs in a blurry way. The interpolation uses μ (without noise), that is, the centres of the distributions.

A plain autoencoder for comparison, with the same test

To make the comparison fair we train a plain autoencoder with the same 10-number code, the same architecture and the same 100 epochs, and apply exactly the same test: draw z ~ N(0, I) and decode. The judge is the same as before.

# The same generation test with a PLAIN autoencoder of the same size (code of 10 numbers, same layers, BCE loss)
class AE(nn.Module):
    def __init__(self, k=LATENT, hidden=128):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(D, hidden), nn.ReLU(), nn.Linear(hidden, k))
        self.dec = nn.Sequential(nn.Linear(k, hidden), nn.ReLU(), nn.Linear(hidden, D))

torch.manual_seed(0)
ae = AE()
opt_a = torch.optim.Adam(ae.parameters(), lr=1e-3)
for epoch in range(100):
    perm = torch.randperm(len(Xtr))
    for i in range(0, len(Xtr), 64):
        x = Xtr[perm[i:i + 64]]
        loss = F.binary_cross_entropy_with_logits(ae.dec(ae.enc(x)), x)
        opt_a.zero_grad(); loss.backward(); opt_a.step()

with torch.no_grad():
    rec = torch.sigmoid(ae.dec(ae.enc(Xte)))
    print(f"AE reconstruction MSE (code {LATENT}):", round(F.mse_loss(rec, Xte).item(), 4))
    codes = ae.enc(Xtr)
    print("AE code statistics: mean of means", round(codes.mean().item(), 2), " mean std", round(codes.std(0).mean().item(), 2))
    ae_samples = torch.sigmoid(ae.dec(torch.randn(2000, LATENT)))                 # the SAME recipe used for the VAE
    probs = F.softmax(clf(ae_samples), dim=-1)
    ae_hist = torch.bincount(probs.argmax(-1), minlength=10).float() / 2000
    print(f"AE samples,  z ~ N(0,I): nearest-training-image distance {nn_distance(ae_samples):.2f}")
    print("classes the judge sees:", [round(v, 2) for v in ae_hist.tolist()])
    print("entropy of that distribution:", round(-(ae_hist[ae_hist > 0] * ae_hist[ae_hist > 0].log()).sum().item(), 2))

AE reconstruction MSE (code 10): 0.0112
AE code statistics: mean of means 0.26  mean std 2.49
AE samples,  z ~ N(0,I): nearest-training-image distance 1.88
classes the judge sees: [0.03, 0.12, 0.15, 0.04, 0.19, 0.17, 0.02, 0.26, 0.01, 0.03]
entropy of that distribution: 1.94

First, the autoencoder reconstructs much better (0.0112 against 0.0282), as expected. Second, its codes do not have the scale of the prior: the mean standard deviation of each component is 2.49, not 1. Third, the samples from N(0, I) fall farther from training digits (distance 1.88, against 1.21 for the VAE and 1.08 for real digits). Fourth, and decisive: the classes are unbalanced: the judge sees 26% of 7s, 19% of 4s and 17% of 5s, but only 1% of 8s and 2% of 6s, with entropy 1.94 against the VAE's 2.27. Sampling "at random" from an autoencoder produces only some of the digits, because N(0, I) falls in a part of the space of its codes that differs from the one where the decoder learned. The VAE, with the KL, made the jar coincide with the recipes. One seed per network: the order of the entropies is clear, the individual frequencies are not.

How many latent dimensions does it really use?

We granted 10 dimensions. Does it use them all? For each dimension we compute the variance of μ across the test images (if it is close to zero, that dimension has the same value for every image and carries no information) and the mean KL.

# How many of the latent dimensions does the VAE actually use?
with torch.no_grad():
    mu, logvar = vae.encode(Xte)
    var_of_means = mu.var(dim=0)                       # how much the MEAN code varies across images, per dimension
    kl_per_dim = (0.5 * (mu.pow(2) + logvar.exp() - logvar - 1)).mean(dim=0)
print("variance of mu per dimension :", [round(v, 2) for v in var_of_means.tolist()])
print("KL per dimension (nats)      :", [round(v, 2) for v in kl_per_dim.tolist()])
print(f"active dimensions (var of mu > 0.01):", int((var_of_means > 0.01).sum()), f"of {LATENT}")

variance of mu per dimension : [0.08, 0.79, 0.01, 0.85, 0.03, 0.67, 0.06, 0.14, 0.85, 0.05]
KL per dimension (nats)      : [0.05, 0.72, 0.0, 0.86, 0.02, 0.52, 0.03, 0.08, 0.78, 0.03]
active dimensions (var of mu > 0.01): 9 of 10

Here the result is not what I expected. With a threshold of 0.01 on the variance of μ the "active" dimensions are 9 out of 10, but the KL per dimension says something else: only four dimensions (the second, fourth, sixth and ninth, counting from one) carry more than 0.5 nats, with 0.72, 0.86, 0.52 and 0.78, while the other six stay below 0.1. The sum is about 3.1, which is the whole training KL.

The VAE has therefore used 4 dimensions out of 10 in a substantial way: this is a partial posterior collapse, and the threshold on the variance of μ is too permissive to notice it, because a dimension can vary a little without carrying useful information. The next lesson shows this phenomenon and how to control it with the weight β. One seed and a single latent size: I cannot say how much of the collapse depends on the small dataset and how much on the recipe.

Trade-off

Reconstruction has a price: the VAE reconstructs worse than the autoencoder (0.0282 against 0.0112) but generates in a balanced way and allows interpolations without holes, which is the compromise of the KL term. The images are soft, because the loss averaged over the pixels and the diagonal Gaussian favour means; for sharper images you need more powerful decoders or other losses. The latent dimension carries a double risk: too small gives poor reconstructions, too large leaves dimensions unused. On the plus side, a VAE trains stably with few hyperparameters.

Real pitfalls

1. Averaging the KL over components and summing the reconstruction over pixels (or the other way round). It changes the relative weight of the two terms by a factor equal to the number of pixels or components, and the network ends up ignoring one of them. Sum both over components, average over images.
2. Forgetting eval() and no_grad() in evaluation. Here the VAE has no dropout or BatchNorm, so eval() is harmless, but the sampling of eps makes the evaluation loss noisy; for a stable number use μ or average several samples.
3. Sampling from the posterior of an image instead of the prior in order to "generate". If you decode μ + σ·ε of an image you are reconstructing, not generating.
4. Judging generation from a few chosen images. Look at a large sample and use measures (distance, diversity), not just the prettiest images.
5. Thinking that a lower ELBO means prettier images. The ELBO measures likelihood under a model of independent pixels; it has no simple relationship with perceptual quality.
6. exp(logvar) blowing up. If logvar becomes large, exp overflows and the loss becomes nan; with high learning rates or unnormalised data you can limit logvar with a clamp. I did not provoke it here.

The six pitfalls share a cause: what the program prints (a loss, an image, a generation) does not coincide with what you want to measure. The case of the active dimensions in this lesson is an example: the threshold on the variance says "9 out of 10", the KL per dimension says "4".

A worked example by hand: one latent dimension and one pixel

The KL of one component. Suppose that the encoder, for a certain image, produces for component 3 of the code μ = 0.8 and log σ² = −0.5. Then σ² = e^(−0.5) = 0.607 and the KL of that component is ½(μ² + σ² − log σ² − 1) = ½(0.64 + 0.607 + 0.5 − 1) = 0.373 nats. A component with μ = 0 and log σ² = 0 has zero KL: it is an "inactive" dimension for that image. The reparametrisation. With noise ε = 0.3 we get z = μ + σ·ε = 0.8 + 0.779·0.3 = 1.034, with σ = √0.607 = 0.779; with ε = −1.2, z = 0.8 − 0.935 = −0.135: same image, different codes, and it is this randomness that makes the decoder learn a zone and not a point.

The loss of one pixel. Let x = 1 (pixel on) and the decoder produce the logit 2.0: the probability is σ(2) = 0.881 and the BCE is −ln 0.881 = 0.127 nats. If instead the pixel were grey, x = 0.3, with the same logit the BCE is −[0.3·ln 0.881 + 0.7·ln 0.119] = 1.527, much higher. And the minimum possible for x = 0.3 is not zero but 0.611 (the entropy of a Bernoulli with parameter 0.3), reached by predicting p = 0.3. This explains why the reconstruction cannot fall to zero with grey pixels, and why comparison with numbers published on binarised datasets is not valid. The calculations in this section are done by hand, not with the program.

Frequently asked questions

Why does the decoder return logits instead of probabilities? Because binary_cross_entropy_with_logits combines the sigmoid and the logarithm in a numerically stable way, avoiding log(0).

Why two linear heads instead of one network that produces 20 numbers? It is equivalent, but the two heads make the structure explicit. In real code a single Linear(128, 2·latent) is often used and the output is split with chunk(2, dim=-1).

How many epochs are needed? Here 100 lead to 24.56 nats and the curve is still slowly descending (24.98 at 50 epochs). With more epochs the loss would keep decreasing; I did not measure it.

Can I condition generation on a digit? Yes with a conditional VAE (you give the label to the encoder and the decoder), but it is not what we built: our VAE has never seen the labels. The judge serves only for evaluation.

Why not sample from the posterior to generate? The posterior depends on an image x that is not there during generation. The prior is the distribution that training has made compatible with the decoder.

What you can do now

You can write a VAE in PyTorch: an encoder that produces μ and log σ², one line of reparametrisation, a logit decoder and a loss of BCE summed over pixels plus KL summed over components. You measured that, after 100 epochs on 8×8 digits, the −ELBO is 24.56 nats (21.47 reconstruction, 3.10 KL), that samples from N(0, I) cover all ten digits (7-16% each) and that the 3→7 interpolation does not pass through other digits. You saw the same test on a plain autoencoder give an unbalanced distribution of digits, and you found out that the VAE uses only 4 dimensions out of 10 in a substantial way: a partial collapse that the next lesson tackles with the β weight and the other training instabilities.