Module 2 — Autoencoders and latent space

Before we build a generator we build something simpler: a network that takes an image, squeezes it through a narrow layer, and tries to reconstruct it on the way out. That narrow layer is a latent space, and the whole of this course orbits around it. This module trains an autoencoder on the MNIST digits set up in module 1, inspects the latent space it carves out, and shows why sampling from that latent space fails — a failure that motivates the variational autoencoder of module 3.

An encoder, a decoder, and a bottleneck between them

An autoencoder is two networks glued together. The encoder $f$ maps an image $x$ to a low-dimensional code $z = f(x)$. The decoder $g$ maps the code back to an image $\hat{x} = g(z)$. Training minimises the reconstruction error:

$\mathcal{L}(x) = \| x - g(f(x)) \|^2$

If the code $z$ had the same dimension as $x$, the identity function would score zero and learn nothing useful. The bottleneck — a code $z$ much smaller than $x$ — forces the encoder to throw information away and the decoder to fill it back in. What survives is the structure the network judged most reconstructable, and for handwritten digits that means digit identity, stroke thickness and slant.

import torch
import torch.nn as nn

class Autoencoder(nn.Module):
    def __init__(self, latent_dim=16):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(28 * 28, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, latent_dim),
        )
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 64), nn.ReLU(),
            nn.Linear(64, 256), nn.ReLU(),
            nn.Linear(256, 28 * 28), nn.Tanh(),
        )

    def forward(self, x):
        z = self.encoder(x)
        x_hat = self.decoder(z).view(-1, 1, 28, 28)
        return x_hat, z

The final Tanh matches the $[-1, 1]$ normalisation established in module 1. The latent dimension of 16 is a common starting point on MNIST: enough capacity to reconstruct all ten digits, small enough to force compression.

Training on MNIST in a few epochs

The training loop looks like any regression problem, because reconstruction is one:

from torch.utils.data import DataLoader
from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,)),
])
loader = DataLoader(datasets.MNIST("./data", train=True, download=True, transform=transform),
                    batch_size=128, shuffle=True)

model = Autoencoder(latent_dim=16).to("cuda" if torch.cuda.is_available() else "cpu")
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(10):
    for x, _ in loader:
        x = x.to(next(model.parameters()).device)
        x_hat, _ = model(x)
        loss = ((x - x_hat) ** 2).mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f"epoch {epoch}: loss = {loss.item():.4f}")

Labels are ignored — that is what makes an autoencoder self-supervised. After ten epochs on a laptop CPU the reconstructions are recognisable, though a shade blurry, and the training loss reaches roughly $0.02$ per pixel.

Squared error blurs, absolute error stays crisper

On images, the squared error loss pulls each pixel towards the average of plausible values and produces a smoothed reconstruction. The absolute error l1_loss pulls each pixel towards the median and preserves sharp edges. Neither is right in absolute terms; both are wrong when the true target is "a plausible sharp image", which is why generative losses do not stop at pixel error.

The latent space, seen with a two-dimensional projection

Set latent_dim=2 and every image becomes a point in the plane. Plotting the whole test set, coloured by digit class, reveals whether the autoencoder has discovered the class structure without being told about it.

import matplotlib.pyplot as plt

model = Autoencoder(latent_dim=2).to("cuda" if torch.cuda.is_available() else "cpu")
# ... train as above ...

with torch.no_grad():
    codes, labels = [], []
    for x, y in loader:
        _, z = model(x.to(next(model.parameters()).device))
        codes.append(z.cpu())
        labels.append(y)
codes = torch.cat(codes).numpy()
labels = torch.cat(labels).numpy()

plt.figure(figsize=(6, 6))
plt.scatter(codes[:, 0], codes[:, 1], c=labels, cmap="tab10", s=3)
plt.colorbar(label="digit")
plt.title("MNIST in a 2D latent space")

The clusters are visible but unevenly sized, and above all they leave large empty regions between them. That geometric fact is the whole point of this module: an autoencoder was trained to compress, not to fill the space.

Interpolation reveals what the network has learnt

Picking two images $x_A$ and $x_B$, encoding both, and decoding intermediate points $z(t) = (1-t) z_A + t z_B$ shows what the network considers a smooth path between them:

import torch.nn.functional as F

def interpolate(model, x_a, x_b, n=8):
    with torch.no_grad():
        _, z_a = model(x_a.unsqueeze(0))
        _, z_b = model(x_b.unsqueeze(0))
        ts = torch.linspace(0.0, 1.0, n).unsqueeze(1)
        zs = (1 - ts) * z_a + ts * z_b
        return model.decoder(zs).view(-1, 1, 28, 28)

For two images of the same digit the interpolation is clean: the stroke slowly morphs. For two different digits, say a "3" and an "8", the middle images are often not digits at all — smudges that look like neither of the two. That is a direct symptom of the empty regions in the latent space: the decoder has never had to reconstruct anything in that neighbourhood.

Why an autoencoder generates poorly

There is one obvious way to turn an autoencoder into a generator: draw $z \sim \mathcal{N}(0, I)$ and decode it. Try it, and the outputs are noise-like blobs, occasionally reminiscent of a digit but never a clean one.

Three reasons combine:

1. The training never told the encoder where to place codes. They sit wherever reconstruction is easiest, which is neither centred at zero nor of unit variance.
2. The latent distribution is unknown. After training we could try to fit a Gaussian mixture to the observed codes, but nothing forced them to be Gaussian in the first place.
3. The decoder was never asked to be robust off-distribution. A code that lies in an empty region of the training latent space is out of distribution for the decoder; its output is undefined.

A high reconstruction score does not mean a good generator

An autoencoder can achieve near-zero reconstruction error and still produce noise from random codes. The two tasks are decoupled — reconstruction rates the round trip from real data, generation rates a sample from a distribution the model never saw. The whole story of the next module is: how do we make training and sampling agree on the same latent distribution?

The fix in one sentence

Force the encoder to place codes in a known, well-behaved distribution — say $\mathcal{N}(0, I)$ — while keeping the reconstruction objective. That single sentence contains the variational autoencoder of module 3, its reparameterisation trick, and its KL divergence term.

In summary

• An autoencoder compresses through a bottleneck and reconstructs; the loss is a pixel error like squared or absolute, and no labels are needed.
• The latent space clusters similar inputs but leaves empty regions between them; interpolating between two different digits often produces something that is neither.
• Sampling a random code from $\mathcal{N}(0, I)$ and decoding it does not produce plausible images: the training never enforced a known latent distribution, and the decoder is undefined off-distribution.
• Reconstruction quality and generation quality are two different metrics; low reconstruction error is a necessary but not sufficient condition for the model to generate.

Next module: the variational autoencoder, which turns this failure into a working generator by constraining the latent distribution.