Q70

VAE image generation: sampling new MNIST digits

VAEGenerationLatent Space

What's actually being asked

Same VAE mechanics as Q69, but the emphasis here is the generation side: after training, sample fresh latent vectors from the prior and decode them into brand-new digit images — not reconstructions of real inputs. This answer uses a convolutional encoder/decoder (vs the fully-connected one in Q69) and keeps the training loop brief since it's identical in structure to Q69's.

Answer

The key conceptual point

During training, z comes from the encoder: z = μ(x) + σ(x) ⊙ ε. The KL term in the loss specifically pushes q(z|x) toward matching the prior p(z) = N(0, I). That's exactly what makes it valid, after training, to skip the encoder entirely and sample z ~ N(0, I) directly — decoding that gives a genuinely new image, not a reconstruction of anything that existed in the training set.

Model (convolutional encoder/decoder this time)

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
latent_dim = 20
 
class ConvVAE(nn.Module):
    def __init__(self, latent_dim=20):
        super().__init__()
        # Encoder: 1x28x28 -> 32x14x14 -> 64x7x7
        self.enc_conv1 = nn.Conv2d(1, 32, kernel_size=4, stride=2, padding=1)
        self.enc_conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1)
        self.fc_mu = nn.Linear(64 * 7 * 7, latent_dim)
        self.fc_logvar = nn.Linear(64 * 7 * 7, latent_dim)
 
        # Decoder: latent -> 64x7x7 -> 32x14x14 -> 1x28x28
        self.fc_dec = nn.Linear(latent_dim, 64 * 7 * 7)
        self.dec_conv1 = nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1)
        self.dec_conv2 = nn.ConvTranspose2d(32, 1, kernel_size=4, stride=2, padding=1)
 
    def encode(self, x):
        h = F.relu(self.enc_conv1(x))
        h = F.relu(self.enc_conv2(h))
        h = h.view(h.size(0), -1)
        return self.fc_mu(h), self.fc_logvar(h)   # estimated mean and log-variance
 
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + std * eps
 
    def decode(self, z):
        h = F.relu(self.fc_dec(z))
        h = h.view(-1, 64, 7, 7)
        h = F.relu(self.dec_conv1(h))
        return torch.sigmoid(self.dec_conv2(h))
 
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar
 
def vae_loss(x_hat, x, mu, logvar):
    recon = F.binary_cross_entropy(x_hat, x, reduction="sum")                 # -E[log p(x|z)]
    kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())              # D_KL(q(z|x) || N(0,I))
    return recon + kl

Training (same shape as Q69 — brief here on purpose)

transform = transforms.ToTensor()
train_data = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
 
model = ConvVAE(latent_dim).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
 
model.train()
for epoch in range(10):
    for images, _ in train_loader:
        images = images.to(device)                  # keep as [B, 1, 28, 28], no flattening needed
        x_hat, mu, logvar = model(images)
        loss = vae_loss(x_hat, images, mu, logvar)
 
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Generation: sampling brand-new digits from the prior

model.eval()
with torch.no_grad():
    z_sample = torch.randn(64, latent_dim).to(device)   # z ~ N(0, I) -- no encoder involved
    generated = model.decode(z_sample).cpu()
 
grid = make_grid(generated, nrow=8)
plt.figure(figsize=(6, 6))
plt.imshow(grid.permute(1, 2, 0).squeeze(), cmap="gray")
plt.axis("off")
plt.title("Newly generated digits, z ~ N(0, I)")
plt.show()

Bonus: latent-space interpolation (shows the space is smooth)

model.eval()
with torch.no_grad():
    img_a, img_b = train_data[0][0].unsqueeze(0).to(device), train_data[1][0].unsqueeze(0).to(device)
    mu_a, _ = model.encode(img_a)
    mu_b, _ = model.encode(img_b)
 
    steps = torch.linspace(0, 1, 8).to(device)
    z_path = torch.stack([mu_a.squeeze(0) * (1 - t) + mu_b.squeeze(0) * t for t in steps])
    interpolated = model.decode(z_path).cpu()
 
grid = make_grid(interpolated, nrow=8)
plt.imshow(grid.permute(1, 2, 0).squeeze(), cmap="gray")
plt.axis("off")
plt.title("Interpolation between two digits in latent space")
plt.show()

Why this works — tying it back to the maths

The loss trains q_φ(z|x) to approximate N(0,I) on average over the whole dataset via the KL term D_KL(q(z|x)‖p(z)). Once that's converged, the decoder p_θ(x|z) has effectively learned "what any plausible z from that standard normal region should look like as an image" — so feeding it fresh noise samples z ~ N(0,I) (never seen, never derived from a real x) still decodes to a plausible-looking digit, which is the entire point of calling a VAE a generative model rather than just an autoencoder.

What to expect visually

The grid of 64 freshly-sampled digits should show recognizable (if slightly blurry) digit shapes with natural variety across the grid — some clearer, some ambiguous between two digit classes (e.g., a shape between a 4 and a 9), which is expected: the latent space is continuous, so regions between well-formed clusters decode to in-between-looking digits.

Exam tip

If the question asks to distinguish "reconstruction" from "generation" explicitly: reconstruction encodes a real x to get z, then decodes it back (tests how well the model compresses/rebuilds known data); generation samples z directly from the prior N(0,I) with no real x involved at all (tests whether the model learned a usable generative latent space). Mixing these two up is the single most common mistake on this topic.