Q69

Variational Autoencoder on MNIST, step by step

VAEMNISTPyTorch

What's actually being asked

Build and train a working VAE on MNIST in PyTorch, with every piece the question names present and correct: encoder, reparameterized latent sampling, decoder, both loss terms, a training loop, and a plot comparing original vs reconstructed digits.

Answer

Architecture, at a glance

image x (784-d) → encoder → μ(x), logσ²(x) → z = μ + σ⊙ε → decoder → x̂ with ε ~ N(0,I).

Full code

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
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)
 
class VAE(nn.Module):
    def __init__(self, latent_dim=20):
        super().__init__()
        self.fc1 = nn.Linear(784, 400)
        self.fc_mu = nn.Linear(400, latent_dim)
        self.fc_logvar = nn.Linear(400, latent_dim)
 
        self.fc2 = nn.Linear(latent_dim, 400)
        self.fc_out = nn.Linear(400, 784)
 
    def encode(self, x):
        h = F.relu(self.fc1(x))
        return self.fc_mu(h), self.fc_logvar(h)
 
    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.fc2(z))
        return torch.sigmoid(self.fc_out(h))
 
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        x_hat = self.decode(z)
        return x_hat, mu, logvar
 
model = VAE(latent_dim=20).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
 
def loss_function(x_hat, x, mu, logvar):
    recon_loss = F.binary_cross_entropy(x_hat, x, reduction="sum")
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + kl_loss
 
epochs = 10
model.train()
for epoch in range(epochs):
    total_loss = 0.0
    for images, _ in train_loader:
        images = images.view(images.size(0), -1).to(device)
        x_hat, mu, logvar = model(images)
        loss = loss_function(x_hat, images, mu, logvar)
 
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
 
    print(f"Epoch {epoch+1}/{epochs}  Loss: {total_loss/len(train_data):.4f}")
 
# --- Display original vs reconstructed ---
model.eval()
with torch.no_grad():
    sample_images, _ = next(iter(train_loader))
    sample_images = sample_images[:8].view(8, -1).to(device)
    reconstructed, _, _ = model(sample_images)
 
fig, axes = plt.subplots(2, 8, figsize=(12, 3))
for i in range(8):
    axes[0, i].imshow(sample_images[i].cpu().view(28, 28), cmap="gray")
    axes[0, i].axis("off")
    axes[1, i].imshow(reconstructed[i].cpu().view(28, 28), cmap="gray")
    axes[1, i].axis("off")
fig.suptitle("Top: original  |  Bottom: reconstructed")
plt.show()

Mapping code to theory

Encoder (encode): fc1 + ReLU gives a hidden representation; fc_mu/fc_logvar are two separate linear heads producing μ(x) and log σ²(x) — predicting log-variance rather than variance directly keeps it unconstrained in sign while exp() guarantees a positive variance.
Reparameterization trick (reparameterize): std = exp(0.5*logvar) recovers σ from logσ², then z = μ + σ ⊙ ε, ε ~ N(0,I) — sampling is pushed onto ε, which carries no gradient requirement, so gradients can backpropagate through μ and σ even though z itself involves a random draw.
Decoder (decode): mirrors the encoder in reverse, ending in a sigmoid since MNIST pixels are normalized to [0,1] and are modeled as Bernoulli probabilities.
Reconstruction loss: F.binary_cross_entropy(x_hat, x, reduction="sum") is −𝔼[log p(x|z)] under a Bernoulli decoder assumption — exactly the first term of the VAE loss.
KL loss: -0.5 * sum(1 + logvar - mu.pow(2) - logvar.exp()) is the closed-form D_KL = −½ Σ_j (1 + log σ_j² − μ_j² − σ_j²) against a standard normal prior, summed over latent dimensions — logvar.exp() is exactly σ_j².
Total loss: recon_loss + kl_loss is the negative ELBO, L(θ,φ;x) = −𝔼_{q_φ(z|x)}[log p_θ(x|z)] + D_KL(q_φ(z|x)‖p(z)); minimizing it is equivalent to maximizing log p(x) ≥ 𝔼[log p(x|z)] − D_KL(...).
Training loop: standard — forward pass, compute loss, zero_grad → backward → step, repeated over epochs.

What the displayed images should show

Early epochs: reconstructions look blurry/averaged. By epoch ~10 with this architecture, reconstructed digits should be clearly recognizable but slightly smoothed compared to the sharp originals — a known VAE characteristic (blurrier than GAN outputs), because the pixel-wise BCE loss and the KL regularization both push toward an "average-looking" reconstruction rather than a sharp one.

Exam tip

If asked why z = μ + σ⊙ε instead of just sampling z ~ N(μ,σ²) directly, the answer is the reparameterization trick: directly sampling isn't differentiable w.r.t. μ/σ, but rewriting the randomness as a fixed, parameter-free ε makes the whole thing differentiable, so backprop can train the encoder.