12.2.1 Latent Space Exploration

Overview

Pick two images. Encode each one into a vector. Draw a straight line between the vectors. Decode every point along the way. If your model is a regular autoencoder, what you get is a sequence of garbage frames between two recognisable endpoints — the path passes through regions of latent space the decoder was never trained on. If your model is a variational autoencoder, you get the figure above: a smooth, plausible transformation where every step looks like something the decoder could have learned to produce.

The difference is a single extra term in the loss function. Kingma and Welling [1] introduced it in 2013, the community spent the next decade building generative models around it, and even now the same idea sits underneath every diffusion model and latent-space generator. This lesson reduces it to its simplest possible form: a small VAE trained on four 16×16 line patterns. Train it for 200 epochs, then walk its latent space.

Learning objectives

1. Distinguish a VAE from a plain autoencoder, and name the loss term that does the work.
2. Read the reparameterisation trick and explain why it is structurally necessary for gradient-based training.
3. Run a trained VAE and produce the four canonical visualisations: reconstructions, latent-space scatter, interpolation, and training progression.
4. Predict what happens to the latent space when you change the KL weight (β-VAE) or the latent dimension.

Quick start — encode, decode, look at the gap

Before training, look at what a fresh VAE does. With random weights, the encoder produces meaningless distribution parameters and the decoder produces meaningless reconstructions — but the shapes line up, and that is the contract you’ll be training against.

import torch
import torch.nn as nn

class VAE(nn.Module):
    def __init__(self, input_dim=256, hidden_dim=128, latent_dim=8):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
        )
        self.fc_mu     = nn.Linear(hidden_dim, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim), nn.ReLU(),
            nn.Linear(hidden_dim, input_dim), nn.Sigmoid(),
        )

    def forward(self, x):
        h = self.encoder(x)
        mu, logvar = self.fc_mu(h), self.fc_logvar(h)
        std = torch.exp(0.5 * logvar)
        z = mu + std * torch.randn_like(std)   # reparameterisation
        return self.decoder(z), mu, logvar

vae = VAE()
x = torch.randn(4, 256)
recon, mu, logvar = vae(x)
print(f"Input  {tuple(x.shape)} -> Latent mu {tuple(mu.shape)} -> Recon {tuple(recon.shape)}")

The output reports Input (4, 256) -> Latent mu (4, 8) -> Recon (4, 256). The full training script in Exercise 1 turns that random pipeline into the trained model behind the opening figure.

Core concepts

Concept 1 — A latent space, and what makes a VAE’s “structured”

A latent space is the low-dimensional output of an encoder — the place where an input gets compressed before reconstruction. Plain autoencoders have them too. The interesting question is not “does the model use a latent space” (every encoder-decoder pair does) but “is the latent space navigable — can you sample from it and get useful outputs?”

For a vanilla autoencoder, the answer is usually no. The encoder learns a mapping that minimises reconstruction error, and it is free to use any region of the latent space it wants. Empty pockets between encoded points are not constrained at all, so decoding a random latent vector typically produces nonsense. Earlier linear techniques in the factor-analysis lineage share this limitation — PCA, ICA, and their relatives all give you compression without a useful generative prior.

A VAE adds one constraint: the encoded distribution q(z|x) for every input should be close to a fixed prior p(z) = N(0, I). Goodfellow, Bengio and Courville’s textbook chapter [3] frames this as turning the decoder into a proper generative model rather than a glorified lookup table. The KL divergence term in the loss function measures and penalises the distance to that prior at every training step. Two things follow:

• Encoded points cluster around the origin instead of scattering arbitrarily.
• Nearby points in latent space decode to similar outputs, because the encoder has been pressured to map similar inputs into overlapping distributions.

That second property — smoothness — is exactly what lets you do the interpolation in Figure 1. The VAE’s latent space is not just compressed, it is organised.

Concept 2 — Encoder, decoder, and the reparameterisation trick

The VAE has three moving parts: an encoder that maps x → (μ, log σ²), a sampler that draws z ~ N(μ, σ²), and a decoder that maps z → x̂.

The sampling step is where the trick happens. A naive implementation might write z = torch.randn(latent_dim) * std + mu and call it done — but this breaks gradient flow. The sampling operation is not differentiable with respect to μ or σ in the way PyTorch’s autograd needs.

The fix Kingma and Welling [1] (and independently Rezende, Mohamed and Wierstra [5]) proposed is the reparameterisation trick: factor the randomness out into a separate input.

def reparameterize(mu, log_var):
    std = torch.exp(0.5 * log_var)
    epsilon = torch.randn_like(std)   # noise sampled *outside* the gradient path
    return mu + std * epsilon         # deterministic w.r.t. (mu, std); ε is just an input

Now z is a deterministic function of μ, σ, and a noise tensor ε. The noise still randomises the sample, but the gradient through z to μ and σ is clean. PyTorch’s autograd [8] handles the rest.

Training minimises a two-part loss:

# Reconstruction term — does the decoder reproduce the input?
recon_loss = F.binary_cross_entropy(recon, x, reduction='sum')

# KL term — is q(z|x) close to N(0, I)?
kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())

total = recon_loss + beta * kl_loss     # β = 1 is the standard VAE

The reconstruction term wants the decoder to be faithful. The KL term wants the latent space to be tidy. The β coefficient [4] lets you tilt the trade-off; β > 1 produces a more disentangled but blurrier model (the β-VAE), β < 1 produces a sharper but less navigable model. The default β = 1 is the original Kingma-Welling formulation.

Concept 3 — Three ways to explore the trained latent space

Once trained, the VAE supports three operations the original autoencoder doesn’t.

Sampling. Draw z ~ N(0, I), run it through the decoder, get a brand new pattern. Doersch’s 2016 tutorial [2] is the canonical walkthrough of why this works: the KL penalty during training has shaped q(z|x) to overlap with N(0, I), so samples from the prior land in regions the decoder knows how to handle.

Interpolation. Encode two real patterns to z_a and z_b, then decode z = (1-α) z_a + α z_b for α swept from 0 to 1. This is what Figure 1 shows. The intermediate frames are not blends of the images — they are decodings of the intermediate latent vectors, which the decoder turns into images that smoothly transform structural features (a diagonal stroke softening, a horizontal stroke appearing, the cross forming).

Latent arithmetic. With more complex data, certain directions in latent space tend to correspond to specific attributes — a “smile” direction for faces, a “stroke-width” direction for handwriting. You can extract them by averaging encodings, then add and subtract them to combine features. Bank, Koenigstein and Giryes [9] survey the recurring patterns and what makes them stable across architectures.

Exercises

EXECUTE I.

Train the VAE and inspect every visualisation

Run the full training script. It generates 1,000 patterns (250 of each type: diagonal, horizontal, vertical, cross), trains the VAE for 200 epochs, and produces five PNG outputs.

latent_space_exploration.py — full training + visualisation pipeline vae_model.py — VAE architecture (imported by the training script)

bash Copy

python latent_space_exploration.py

The script prints loss values every 50 epochs and saves five PNGs to the working directory: vae_architecture.png, latent_space_visualization.png, reconstruction_comparison.png, interpolation_sequence.png, and training_progression.png. The architecture and the first three are already embedded above; the training progression looks like this:

Reflection

1. Why do the reconstructions at epoch 1 look like undifferentiated noise rather than at least the right colour balance?
2. In the latent-space scatter (Figure 3), the four clusters have soft boundaries. What would you expect them to look like in a plain autoencoder of the same capacity?
3. The “cross” pattern in the reconstructions (Figure 2) is the noticeably blurriest one. Why might that be, given how it’s constructed from horizontal + vertical?
4. If you increased LATENT_DIM from 8 to 32, what would you expect to happen — and what is the catch?

Discussion

1. At epoch 1 the encoder has barely been trained, so μ and log σ² are still close to their random init — meaning every input maps to nearly the same point in latent space. The decoder, also random, produces one near-constant output regardless of input. There is no “right colour balance” yet because there is no signal flowing.

2. Sharper but more chaotic. A plain autoencoder has no pressure to keep the encoded distribution tidy, so the four pattern types would map to four well-separated regions with empty space between them. Reconstruction error on the training set would be lower, but the empty regions would decode to garbage — try this by setting beta = 0 in the loss.

3. The cross is a superposition of the horizontal and vertical patterns. In latent space it sits closer to the boundary between those two clusters, and the VAE’s averaging behaviour makes superpositions inherently soft — there are more “plausible” cross-like images for the decoder to average over than there are for a single thick stroke.

4. Sharper reconstructions and slightly more capacity for variation. The catch: many of the extra dimensions are likely to become “dead” — the network learns to ignore them by collapsing μ to 0 and σ to 1 for those coordinates. This is called posterior collapse in its mild form, and you can detect it by measuring the per-dimension variance of μ across the training set after training.

LATENT_DIM = 2     # tiny: visualisable directly, limited capacity
LATENT_DIM = 8     # default
LATENT_DIM = 32    # large: expect dead dimensions
LATENT_DIM = 64    # excessive for 4 pattern types

def vae_loss(recon_x, x, mu, log_var, beta=4.0):  # try 0.1, 1.0, 4.0, 10.0
    ...
    return recon_loss + beta * kl_loss, recon_loss, kl_loss

# In SimpleVAE:
def encode(self, x):       # TODO: x -> ReLU(fc1) -> ReLU(fc2) -> mu
    pass
def decode(self, z):       # TODO: z -> ReLU(fc1) -> ReLU(fc2) -> Sigmoid(out)
    pass

# Module-level:
def interpolate(vae, pattern_a, pattern_b, num_steps=8):
    # TODO: encode both, walk z = (1-α)z_a + α z_b, decode each
    pass

def encode(self, x):
    h = torch.relu(self.encoder_fc1(x))
    h = torch.relu(self.encoder_fc2(h))
    return self.fc_mu(h)

def decode(self, z):
    h = torch.relu(self.decoder_fc1(z))
    h = torch.relu(self.decoder_fc2(h))
    return torch.sigmoid(self.decoder_out(h))

def encode(self, x):
    h = torch.relu(self.encoder_fc1(x))
    h = torch.relu(self.encoder_fc2(h))
    return self.fc_mu(h)

def decode(self, z):
    h = torch.relu(self.decoder_fc1(z))
    h = torch.relu(self.decoder_fc2(h))
    return torch.sigmoid(self.decoder_out(h))

def interpolate(vae, pattern_a, pattern_b, num_steps=8):
    vae.eval()
    frames = []
    with torch.no_grad():
        z_a = vae.encode(pattern_a)
        z_b = vae.encode(pattern_b)
        for alpha in np.linspace(0, 1, num_steps):
            z = (1 - alpha) * z_a + alpha * z_b
            frames.append(vae.decode(z).numpy().reshape(IMAGE_SIZE, IMAGE_SIZE))
    return frames

Summary

Common pitfalls

• Forgetting torch.no_grad() during evaluation — memory and time leak fast when you’re running many decode-only operations.
• β tuned to 0 chases sharper reconstructions but turns the model into a vanilla autoencoder that scatters its latent codes, breaking every operation in Concept 3.
• Dead latent dimensions are not a bug. They are the network telling you LATENT_DIM is bigger than the intrinsic dimensionality of your data — useful diagnostic, not a failure mode.
• VAE samples will not look as sharp as GAN samples, by design. Pick the model whose trade-off matches your goal.
• Reconstruction-loss numbers aren’t comparable across runs with different β values, because you’ve changed what the model is being asked to optimise in the first place.

References