VAEs: Generating with Probability

From a Point to a Cloud: Making the Latent Space Probabilistic

In the last guide we built a plain autoencoder: an encoder squeezes an image down to a short list of numbers called a latent code, and a decoder rebuilds the image from that code. It learned to compress and reconstruct beautifully — but we hit a wall when we tried to generate. The trained codes sat like scattered islands in the latent space, with large empty stretches between them. Pick a random point in one of those gaps, hand it to the decoder, and you get mush. The space was full of holes, and a space full of holes is no place to sample from.

The variational autoencoder (VAE) fixes this with one deceptively simple change. Instead of encoding an image to a single point, it encodes the image to a small fuzzy cloud — a probability distribution. Think of pinning a photo onto a map. The old autoencoder stuck a sharp pin into one exact spot. The VAE instead paints a soft, blurry blob over a small region and says 'this photo lives roughly here.' That blurriness is not sloppiness; it is the whole trick.

Why does fuzziness close the holes? Because when every photo claims a small region rather than a single dot, neighbouring clouds overlap. Where two clouds overlap, the points in between belong a little to each image — so the decoder is forced to make those in-between points decode into something sensible (a blend of the two images). Do this for the whole dataset and the clouds knit together into a continuous fabric with no dead gaps. Every point you might land on now means something, which is exactly the property we need before we dare to sample.

So the encoder's job changes. It no longer spits out a single code z. Instead it predicts the recipe for a cloud: two vectors, a mean μ saying where the centre of the cloud sits, and a variance σ² saying how spread out (how fuzzy) the cloud is. We write this as the distribution below.

The encoder defines a Gaussian distribution over latent codes, given an input image.

Read this left to right. The symbol q(z | x) means 'the distribution over latent codes z, given the image x' — the vertical bar is read as 'given'. The 𝒩 stands for a Gaussian (the classic bell curve). μ (mu) is the cloud's centre: a vector the encoder predicts, one number per latent dimension. σ² (sigma-squared) is the variance: how wide the bell is in each dimension, also predicted by the encoder. The crucial shift is this — the encoder no longer outputs z itself. It outputs the recipe (μ, σ) for a whole cloud of possible z's. For example, if the latent space is just 2-dimensional and the encoder predicts μ = (1.0, −0.5) and σ = (0.2, 0.3), then this photo's cloud is centred at the point (1.0, −0.5) and is slightly fuzzier along the second axis than the first.

A Tidy Latent Space: The Gaussian Prior

Encoding each image to a cloud fills the local gaps, but it does not yet guarantee that the whole collection of clouds is organized. Imagine the encoder, left to its own devices, parking some clouds way out at coordinate (50, 70), others crammed near the origin, others scattered off in distant corners. Each individual cloud might be fine, but the overall pile would be a sprawling, lopsided mess — and we still wouldn't know where to throw a dart to get a valid image. We need to control the global shape of all the latent codes taken together.

Here is the goal, stated plainly: we want the entire pile of latent codes, across all images, to collectively look like one simple standard bell curve centred at the origin — a neat round ball of probability with no faraway outliers and no gaps. Picture taking all those fuzzy clouds and gently herding them so they pack together into a tidy sphere around the centre. This target shape is called the prior: the distribution we want the codes to follow, decided in advance.

Why insist on this specific simple shape? Because of what happens at generation time. Once training is done we will throw away the encoder entirely and just draw a random code and decode it. If the codes were supposed to live in some weird, lumpy, unknown distribution, we wouldn't know how to draw from it. But if they're supposed to look like a plain standard bell curve, drawing a fresh code is trivial — it's the most basic random number any library can give you. We are deliberately choosing an easy target now so that sampling later is effortless.

The prior: our chosen target distribution for latent codes — a standard multivariate Gaussian.

Decoding the symbols: p(z) is the prior — the target distribution we picked for the latent codes. 𝒩(0, I) is a standard multivariate Gaussian. The 0 means the cloud is centred at the origin (every dimension's mean is zero). The I is the identity covariance matrix, which carries two messages at once: every dimension has a spread (variance) of exactly 1, and the dimensions are independent of one another — no axis leans on any other. In plain words the whole equation says: 'we want the codes to look like ordinary random noise drawn from a standard bell curve.' Concretely, in a 2-D latent space, p(z) = 𝒩(0, I) is a round blob centred at (0, 0) where about 99% of the probability sits within a radius of roughly 3 — that round, predictable shape is exactly what makes sampling later so easy.

KL Divergence: A Tax for Straying from the Prior

To push the encoder's clouds toward the prior, we first need a way to measure how far off they are. That measuring stick is the KL divergence (Kullback–Leibler divergence). Build the intuition first: KL is a kind of distance between two probability distributions. It is zero exactly when the two distributions match perfectly, and it grows larger the more they disagree. It's a little unusual in that it is asymmetric — the 'distance' from A to B is not necessarily the same as from B to A — so it's a directed mismatch score rather than a true symmetric ruler. But for our purposes you can safely think of it as 'how different is the encoder's cloud from the target prior?'

Here's the analogy that makes it click: think of KL as a tax. Every training step, the encoder is allowed to place each image's cloud wherever it likes and make it any size it likes — but it must pay a penalty proportional to how far that cloud strays from the standard normal. Park a cloud far out at (50, 70)? Big tax. Make a cloud razor-thin or balloon-wide instead of spread-1? Tax. The only way to pay zero tax is to make the cloud exactly the standard bell curve at the origin. This levy is precisely the force that herds all the clouds into the tidy, gap-free sphere we wanted in the last section.

The KL divergence between two Gaussians has a clean closed form — no integral needed.

This looks dense, but every piece earns its place. The sum Σ runs over j = 1 to d, the latent dimensions, so we compute a tax per dimension and add them up. Inside, μⱼ² penalizes a cloud whose centre has drifted away from the origin — the farther out, the bigger the square, the heavier the tax. σⱼ² penalizes a cloud that has grown too wide (large variance costs more). The −ln σⱼ² term pulls the other way: as σⱼ² shrinks toward zero, −ln σⱼ² shoots up toward infinity, so it punishes a cloud that collapses too narrow (a cloud with no fuzziness is just a sharp pin again — the very thing we were escaping). The −1 and the ½ are constants tuned so the whole expression bottoms out at exactly zero when μ = 0 and σ = 1.

Let's plug in a number to feel it. Take a single dimension that is already perfect: μ = 0, σ = 1, so σ² = 1 and ln σ² = ln 1 = 0. The bracket becomes 0 + 1 − 0 − 1 = 0. Zero tax — the cloud obeys, nothing to pay. Now nudge it off-centre to μ = 2 with σ still 1: the bracket is 4 + 1 − 0 − 1 = 4, times ½ gives a tax of 2. The encoder feels that 2 as a gradient pulling μ back toward 0. The penalty vanishes only when the cloud is exactly the standard normal — obey the prior and the tax disappears.

The ELBO: One Objective, Two Forces

We now have two demands — reconstruct faithfully, and stay close to the prior — and we need to fold them into a single number to maximize during training. That number is the evidence lower bound, almost always abbreviated ELBO. Let's unpack the intimidating name gently. 'Evidence' is just a fancy word for how likely our model thinks the real training images are; we'd love to make that as high as possible, but it's mathematically awkward to compute directly. 'Lower bound' means the ELBO is a quantity that always sits below that true likelihood. So if we push the ELBO upward, we drag the true likelihood up along with it — we're pushing on a floor that lifts the ceiling. Maximize the ELBO and you maximize the model's ability to have plausibly produced the real images.

The VAE's single training objective: a reconstruction reward minus a tidiness tax.

Let's decode every symbol. The first term, 𝔼_{q(z|x)}[ log p(x | z) ], reads 'the average, over codes z drawn from the encoder's cloud q(z|x), of log p(x | z)'. The 𝔼[...] is an expectation — just a weighted average. Inside, p(x | z) is the decoder's reconstruction quality: the probability the decoder assigns to rebuilding the original image x when fed the code z. It is high when the rebuilt image closely matches x, so this whole term is a reward for faithful reconstruction. The second term is exactly the KL divergence tax from the previous section, measuring how far the encoder's cloud q(z|x) strays from the prior p(z). The minus sign in front means we subtract that tax — straying is penalized. Put plainly: maximize ELBO = reconstruct the image well AND keep the cloud close to the prior. That is the entire VAE objective in one line.

Connect the first term back to guide 1. There we trained the plain autoencoder with a squared-error reconstruction loss — sum the squared pixel differences between the original and the rebuild, and push it toward zero. It turns out that 'log p(x | z) for a Gaussian decoder' is, up to a constant, just the negative of that very squared error. So maximizing 𝔼[log p(x | z)] is the same familiar move as minimizing pixel-wise squared error, only now averaged over codes sampled from the cloud. Nothing exotic — the same reconstruction goal you already know, dressed in probability.

Picture the two terms as a tug-of-war. The reconstruction term wants to pull every image's cloud apart and out to its own distinctive spot, so it can be told from the others and rebuilt precisely. The KL term wants to pull all the clouds together toward the origin into one neat ball. Pull too hard on reconstruction and you get the old holey islands; pull too hard on KL and every cloud collapses into mush. Good generation lives in the balance between them. In practice we often add a weighting knob, β, in front of the KL term to control the trade-off: turn β up and the latent space gets cleaner and more disentangled but reconstructions blur; turn β down and images get sharper but the space gets messier. Tuning that single knob is one of the central crafts of training a VAE.

The Reparameterization Trick: Making Randomness Trainable

Everything so far is elegant on paper, but there's a serious problem hiding in the middle of the pipeline, and it's worth slowing right down for. Neural networks are trained by backpropagation: we compute how the loss changes as we wiggle each parameter (the gradient) and nudge the parameters to reduce the loss. For gradients to flow, every step from input to loss must be a smooth, differentiable function. But look at the VAE's middle: the encoder produces a cloud (μ, σ), and then we have to sample a code z from that cloud — we draw a random point. And you simply cannot differentiate through a random draw. A coin flip has no derivative; there's no smooth knob saying 'if I nudge μ a little, the random outcome shifts this much.'

Here's an analogy. Imagine a slot machine with a lever (your learnable μ and σ) and, sealed inside, a set of dice that get rolled to decide the payout (the random sampling of z). You want to learn how to adjust the lever to improve your expected winnings. But the randomness happens inside the sealed box, downstream of your lever, so you can't trace a clean cause-and-effect line from the lever, through the dice, to the result. The gradient gets stuck at the dice. As long as the random step sits between your parameters and the loss, training is blocked.

The fix is the reparameterization trick, and it is beautifully simple once you see it: move the randomness out of the way. Instead of drawing z directly from the cloud 𝒩(μ, σ²), we draw plain, generic noise ε from a fixed standard normal that has no learnable parameters, and then build z deterministically out of μ, σ, and that noise. The randomness still happens — but now it enters from the side, as an external input, not on the path between the parameters and the loss.

Build the latent code deterministically from the encoder's outputs plus fixed external noise.

Symbol by symbol: μ and σ come from the encoder and carry the network's gradients — they are the learnable knobs. ε (epsilon) is fresh standard-normal noise, drawn anew on every step from 𝒩(0, I); it carries no gradient because it's just a random input from outside, with nothing learnable in it. The ⊙ symbol means elementwise multiplication: multiply each component of σ by the matching component of ε. So we take the noise, stretch it by σ, shift it by μ, and that's our code. The payoff is twofold. First, statistically this is identical to drawing z straight from 𝒩(μ, σ²) — scaling standard noise by σ and shifting by μ is exactly how you turn a standard bell curve into one with mean μ and spread σ. Second, and crucially, z is now a plain differentiable function of μ and σ: the derivatives ∂z/∂μ = 1 and ∂z/∂σ = ε both exist. Gradients can flow straight back through z into the encoder. The slot machine's dice now live outside the box, so we can finally learn the lever.

import torch

def reparameterize(mu, logvar):
    # The encoder predicts log-variance for numerical stability,
    # so recover the standard deviation sigma = exp(0.5 * logvar).
    sigma = torch.exp(0.5 * logvar)

    # Draw fresh standard-normal noise epsilon ~ N(0, I).
    # randn_like has NO learnable parameters -> it carries no gradient.
    epsilon = torch.randn_like(sigma)

    # Build z deterministically: z = mu + sigma (elementwise *) epsilon.
    # Gradients now flow through mu and sigma, but not through epsilon.
    z = mu + sigma * epsilon
    return z

# Why predict logvar instead of sigma directly?
# logvar can be any real number, while sigma must stay positive;
# exp(0.5 * logvar) is always > 0, so the network never has to
# fight a positivity constraint during training.

Sampling New Images and the Blur Problem

Now the payoff we've been building toward. Once the VAE is trained, generating a brand-new image is almost embarrassingly easy. We don't need the encoder at all anymore — its only job was to organize the latent space during training. To create a fresh image we just draw a latent code z straight from the standard normal prior 𝒩(0, I), feed it to the decoder, and out comes a new picture the model has never seen. And because the KL tax forced the latent space to be tidy and gap-free, wherever that random z lands it falls in meaningful territory — the decoder knows what to do with it. This is the precise thing the plain autoencoder could not do.

import torch

# Generation: the encoder is gone. We sample straight from the prior.
num_samples = 16
latent_dim = 128

# Draw codes from N(0, I) -- the exact distribution we trained toward.
z = torch.randn(num_samples, latent_dim)

# Decode each code into a brand-new image.
new_images = decoder(z)   # shape: (16, channels, height, width)

# --- Latent interpolation: walk smoothly between two images ---
z_a = torch.randn(1, latent_dim)   # code for an imagined face A
z_b = torch.randn(1, latent_dim)   # code for an imagined face B
for t in torch.linspace(0, 1, steps=10):
    z_mix = (1 - t) * z_a + t * z_b   # straight line through latent space
    frame = decoder(z_mix)            # face A morphs smoothly into face B

Here's the moment that makes people fall in love with VAEs: latent interpolation. Pick two codes, z_a and z_b, and walk in a straight line from one to the other, decoding each step along the way. Because the latent space is now continuous and meaningful everywhere, the decoded images morph smoothly — one face slowly grows a smile, rounds its cheeks, and turns into another face, with every in-between frame a plausible face of its own. There are no jarring jumps and no garbage frames in the gaps, because there are no gaps. That smoothness is the visible proof that the probabilistic latent space worked.

Now the honest limitation, because mastery means knowing the weaknesses too. VAE outputs tend to look soft and blurry — pleasant and coherent, but lacking crisp edges and fine texture. The reason traces straight back to the reconstruction loss. Remember it's a pixel-wise squared error: it rewards getting each pixel's value close to the target. But for any given code there are many equally plausible sharp images (a strand of hair could fall slightly left or slightly right, an edge could sit one pixel over). Since squared error punishes being wrong about any single pixel, the decoder plays it safe by predicting the average of all those plausible options — and the average of many sharp possibilities is a smooth, smeared one. Averaging blurs edges. The blur isn't a bug in the code; it's baked into 'match every pixel' as a goal.