Skip to content

4. VAE

Deadline and Submission

📅 22.oct (thursday) 2026

🕐 Commits until 23:59

Individual

Click here to submit

Activity: Variational Autoencoders (VAEs)

This activity is designed to test your skills in Variational Autoencoders (VAEs).

The thread running through the activity is the tug of war inside the loss. A VAE pays for two things at once: reconstructing its input, and keeping its latent space close to a prior you can sample from. First a real VAE on MNIST, logging the two terms separately; then a 2-D latent space you can actually look at, to see what that prior buys you — and what it costs.

Technical rules (they apply to the whole activity)

  • Fix the seeds — torch.manual_seed(42) before building each model, and the split generator given in Exercise 1. Results that cannot be reproduced score no points;
  • Every plot must have a title and axis labels; every scatter, a class legend;
  • Allowed libraries: torch and torchvision (for loading MNIST), numpy, matplotlib/seaborn, and scipy/scikit-learn for utilities (quantiles, PCA). Autograd is allowed this time — the point of the activity is the VAE, not backpropagation. But the VAE-specific parts — the reparameterization, both terms of the loss, and the sampling — MUST BE WRITTEN BY YOU. No VAE class from a library (pythae, diffusers' AutoencoderKL, …) and no torch.distributions.kl_divergence: the closed-form KL is written out. Using one voids the implementation criterion;
  • Report both loss terms separately, in nats per image (summed over the 784 pixels, averaged over the images). A total that goes down says nothing about which of the two terms is winning;
  • The test set is used only for the final numbers of each model. Every decision during training looks at the validation set;
  • Whenever the statement asks for a number (a loss term, a KL, a test index), report the number in the text — not only in the code output;
  • Organize the report with one heading per exercise and one subheading per item (Exercise 1, A, B, …), in the same order as the statement, and number the figures as indicated. The last section of the report must be the Results summary described at the end of this page.

Exercise 1

A VAE on MNIST

A — Load and split the data

Use MNIST from torchvision, with pixels in \([0, 1]\) (transforms.ToTensor() already does this), flattened into 784-vectors. Split the 60 000 training images into 50 000 for training and 10 000 for validation, exactly like this:

from torch.utils.data import random_split

train_full = datasets.MNIST("data", train=True, download=True, transform=transforms.ToTensor())
test = datasets.MNIST("data", train=False, download=True, transform=transforms.ToTensor())
train, val = random_split(train_full, [50_000, 10_000],
                          generator=torch.Generator().manual_seed(42))

The 10 000 official test images are the test set.

B — Build the model

Use this architecture, so that the numbers in the report can be compared with each other and with the reference values below:

Part Layers
Encoder \(784 \to 512 \to 256\), ReLU after each
Heads two linear layers \(256 \to d\): one for \(\boldsymbol{\mu}\), one for \(\log \boldsymbol{\sigma}^2\)
Decoder \(d \to 256 \to 512 \to 784\), ReLU after the hidden layers, sigmoid at the output

with latent dimension \(d = 16\). The loss per image, summed over pixels and over latent dimensions:

\[ L = \underbrace{-\sum_{j=1}^{784} \Big[x_j \log \hat{x}_j + (1 - x_j) \log (1 - \hat{x}_j)\Big]}_{\text{reconstruction (BCE)}} \;+\; \beta \cdot \underbrace{\tfrac{1}{2}\sum_{i=1}^{d} \left(\mu_i^2 + \sigma_i^2 - 1 - \log \sigma_i^2\right)}_{\text{KL}} \]

with \(\beta = 1\), averaged over the images of the batch.

Why BCE here and \(\|\mathbf{x} - \hat{\mathbf{x}}\|^2\) in the lecture

Both are \(-\log p(\mathbf{x} \mid z)\), the reconstruction term of the ELBO, under two different assumptions about the pixels. A Gaussian likelihood gives the squared error; a Bernoulli likelihood — each pixel a probability of being "on", which is what a sigmoid output means — gives the binary cross-entropy. MNIST pixels are almost all 0 or 1, so the Bernoulli fits. Use F.binary_cross_entropy(..., reduction="sum") and divide by the batch size; reduction="mean" would also average over pixels and shrink the reconstruction term 784 times relative to the KL, which is the same as secretly setting \(\beta \approx 784\).

C — Train

Train with Adam, learning rate \(10^{-3}\), batch size 128, for 20 epochs. At the end of each epoch, record the reconstruction term and the KL term on the training set (the average over the epoch) and on the validation set. Then:

  1. Report the test reconstruction, the test KL and their sum (the negative ELBO), in nats per image. A correct implementation lands around 81 nats of reconstruction and 20 nats of KL. On a laptop CPU it takes about half a minute;
  2. Produce Figure 1: two panels, reconstruction \(\times\) epoch and KL \(\times\) epoch, each with the training and the validation curve.

D — Reconstructions and samples

  1. Produce Figure 2: the first test image of each digit 0–9 on the top row, and its reconstruction on the bottom row. Reconstruct by decoding \(\boldsymbol{\mu}\) — do not sample \(z\).
  2. Produce Figure 3: an \(8 \times 8\) grid of images decoded from \(z \sim \mathcal{N}(0, I)\), drawn after torch.manual_seed(0).

E — Analysis

  1. Figure 1: the KL rises during the first epochs while the total loss falls. Why would the optimizer choose to pay more KL?
  2. Figure 3 looks worse than Figure 2 — blurrier digits, some that are no digit at all. Give two separate reasons, one about where the prior samples land in the latent space, and one about what the decoder outputs for a given \(z\) (the lecture's section on blur).
  3. Why decode \(\boldsymbol{\mu}\) and not a sampled \(z\) in Figure 2?

Exercise 2

A Latent Space You Can See

With \(d = 16\) the latent space cannot be drawn. With \(d = 2\) it can — at a price that this exercise asks you to measure.

A — Train a 2-D VAE

Train the same model as in Exercise 1, changing only \(d = 2\). Same data, seed, \(\beta\) and training budget. Report the test reconstruction and KL (expect roughly 135 and 7 nats) and put them next to the \(d = 16\) numbers.

B — The encoded test set

Produce Figure 4: a scatter of \(\boldsymbol{\mu}\) for all 10 000 test images, colored by digit, with a legend.

C — The decoded plane

Produce Figure 5: decode a \(15 \times 15\) grid of latent points and tile the images in the same layout. Do not space the grid linearly: use \(z = \Phi^{-1}(q)\), where \(q\) runs over 15 evenly spaced values in \([0.05, 0.95]\) and \(\Phi^{-1}\) is the inverse CDF of \(\mathcal{N}(0, 1)\) (scipy.stats.norm.ppf). Say in one sentence why this spacing gives every cell the same prior mass.

D — Interpolation

Choose two test images of different digits and report their indices. Produce Figure 6 with two rows of 10 images each, for \(t\) from 0 to 1:

  • top row, in pixel space: \((1 - t)\,\mathbf{x}_a + t\,\mathbf{x}_b\);
  • bottom row, in latent space, with the \(d = 16\) model of Exercise 1: decode \((1 - t)\,\boldsymbol{\mu}_a + t\,\boldsymbol{\mu}_b\).

E — Analysis

  1. Which digits overlap in Figure 4? Find the same transition in Figure 5 and describe what the in-between images look like.
  2. \(d = 2\) reconstructs clearly worse than \(d = 16\). What does two numbers per image force the model to throw away? Use Figure 5 as evidence.
  3. Compare the two rows of Figure 6. Which one passes through images that look like digits, and why does the other one not?

Results summary

Close the report with this table, filled in:

# Quantity Value
1 Ex. 1 — test reconstruction, KL and negative ELBO (\(d = 16\))
2 Ex. 1 — validation KL at the first and at the last epoch
3 Ex. 2 — test reconstruction and KL (\(d = 2\))
4 Ex. 2 — the two test indices interpolated

Evaluation Criteria

The deliverable for this activity is a report that includes:

  1. The code for your VAE, its loss and its training loop, commented.
  2. Figures 1 to 6, numbered as requested.
  3. Your answers to the analysis questions in items 1E and 2E.
  4. The Results summary table.

Important Notes:

  • The deliverable is a GitHub Pages site backed by a public repository — see Submission Format for the required layout, front matter and checklist;

  • There is a strict no-plagiarism policy. Any form of plagiarism will result in a zero grade for the activity and may lead to further disciplinary action under the university's academic integrity policies;

  • The deadline for each activity is not extended — NO EXCEPTIONS will be made for late submissions.

  • AI collaboration is allowed, but each student MUST UNDERSTAND and be able to explain every part of the submitted code and analysis. Any use of AI tools must be properly cited. ORAL EXAMS may be conducted.

Grading Criteria:

Each row is worth the points indicated, awarded in full, partially (half), or not at all: in full when the item is complete and correct; partially when it is implemented but missing the requested analysis, or when the analysis lacks the numerical result that supports it; zero when absent or incorrect.

How good the samples look is not graded, and neither are the loss values themselves. Data, architecture and seeds are fixed, so the numbers are what they are; what is graded is the method, the figures, and reporting honestly what you got.

Exercise 1 — VAE on MNIST (5 points):

Points Criterion What earns full credit
0.5 Data (A) The specified split and generator; the test set not used before the final numbers.
2.0 Model and loss (B) The specified architecture; reparameterization, BCE summed over pixels and the closed-form KL all hand-written and correct.
1.0 Training (C) Both terms logged separately on train and validation, test numbers reported, Figure 1 correct.
0.5 Reconstructions and samples (D) Figures 2 and 3 as specified — reconstruction from \(\boldsymbol{\mu}\), samples from the prior with the given seed.
1.0 Analysis (E) The rising KL explained; two distinct reasons for Figure 3 looking worse than Figure 2; why \(\boldsymbol{\mu}\) in Figure 2.

Exercise 2 — 2-D latent space (5 points):

Points Criterion What earns full credit
1.0 2-D model (A) Same model and budget with \(d = 2\), numbers reported next to \(d = 16\).
1.0 Encoded test set (B) Figure 4 with all test points and a legend.
1.0 Decoded plane (C) Figure 5 on the \(\Phi^{-1}\) grid, with the sentence on why.
1.0 Interpolation (D) Figure 6 with both rows, indices reported, latent row from the \(d = 16\) model.
1.0 Analysis (E) Overlap located in Figures 4 and 5; the \(d = 2\) cost explained; the two interpolations contrasted.