VAEs from Scratch — Stir Probability into "Compress and Restore" and You Get a Generator

An autoencoder that only compresses and restores cannot invent anything new. This walks through why a single drop of probability turns it into a generative model — ELBO, the reparameterization trick, and walking the latent space — assuming no prior knowledge.

Start with a packing problem

Picture two people moving house: one packs the contents of a room into a suitcase, the other opens it at the new place and puts everything back. The smaller the suitcase, the harder the packer has to choose — every discarded item is a bet that the unpacker can infer it from context. Which means the reverse is also true: if a small case was enough to restore the room, its contents were a summary of what actually mattered.

That is an autoencoder. The packer is the encoder, the unpacker is the decoder, and the small case in between is the latent variable $z$. The only training signal is "did it come back correctly" — typically pixel-wise squared error. The label is the input itself, so no human annotation is involved.

The narrowness of the case is the whole point. If the case were as big as the room, the packer could dump everything in unchanged and learn nothing. Squeezing the passage is what forces a meaningful summary.

An autoencoder cannot generate

Take a trained autoencoder, throw away the encoder, feed the decoder some arbitrary $z$. Do you get a new image? You do not. You get noise, or a mangled echo of something in the training set.

The reason is that this training promises nothing about where $z$ lives. Encode ten thousand training images and you get ten thousand points scattered through the latent space, and the decoder was only ever drilled on those points. Nobody wrote anything into the terrain between them. It is a map that marks the towns but stays silent about whether the ground between two towns is a plain or a cliff.

Worse, you do not even know roughly where those ten thousand points sit. With no answer to where do I draw $z$ from, there is nothing to sample. To get a generative model you need both problems solved at once: no holes in the latent space, and a known place to draw from.

Stir in a drop of probability

The idea behind the VAE (variational autoencoder) is almost anticlimactic: make the encoder emit a blur instead of a point. Its output is no longer $z$ itself but a center $\mu$ and a spread $\sigma$, and the actual $z$ is drawn afresh each time from the Gaussian those two define. Feed the same image ten times and ten slightly different $z$ come out.

That solves the first problem. From the decoder's side, the same image now arrives at a slightly different address every time, and it is penalized unless it can restore that image from any of them. So the whole neighborhood around $\mu$ gets painted in as "the region for this image". Training points that used to be dots acquire area.

There is still an escape hatch, though: push $\sigma$ toward zero and the blur collapses back into a point. So a second pressure is applied — a penalty that says every input's blur must stay close to one shared, pre-agreed distribution. That distribution is normally the standard Gaussian $\mathcal{N}(0, I)$, and it is called the prior.

This closes the escape hatch and, in the same stroke, settles the second problem. If every blur has been pulled toward the standard Gaussian, then to generate you simply draw one $z$ from the standard Gaussian and run it through the decoder. The tap was decided in advance. That is what "stir probability into compress-and-restore and you get a generator" means. Kingma and Welling proposed this in 2013, and Rezende and colleagues arrived at an equivalent formulation independently at almost the same time.

The mechanism: one equation called the ELBO

What we actually want to maximize is $p(x)$, the probability that our data $x$ comes out of the model. But $p(x) = \int p(x \mid z)\, p(z)\, dz$ is an integral that sweeps over every possible $z$ — hopeless once $z$ has a few dozen dimensions.

So the VAE gives up on touching $p(x)$ directly and instead lifts a quantity guaranteed to sit below it. Raise the plank propping something up from underneath and the thing on top rises too. That lower bound is the ELBO (Evidence Lower BOund).

Symbol by symbol: $q(z \mid x)$ is the blur the encoder emits, i.e. the Gaussian with center $\mu$ and spread $\sigma$. $p(x \mid z)$ is the decoder — how plausibly $x$ falls out of a given $z$. $p(z)$ is the prior $\mathcal{N}(0, I)$. $\mathbb{E}$ is an average, and $D_{\mathrm{KL}}$ measures how far apart two distributions are.

Equation (1), in words: "how well a $z$ drawn from the blur restored the input" minus "how far the blur drifted from the tap" is a floor under the quantity we really want to raise. The left side is uncomputable, the right side is computable — that is the entire trick. Training flips the sign and descends $-\text{ELBO}$ as the loss.

KL divergence itself is covered in KL Divergence From Scratch; the one line you need here is that it is a fine that is zero for identical distributions and grows as they separate. And because both $q$ and $p(z)$ are Gaussian, the KL term has a closed form — no integration required.

$d$ is the number of latent dimensions and $j$ indexes them one by one. This is a formula which says nothing more than the further $\mu$ strays from 0 the bigger the fine, and the further $\sigma$ strays from 1 the bigger the fine — plug in $\mu_j = 0,\ \sigma_j = 1$ and the bracket becomes $1 + 0 - 0 - 1 = 0$. Since this term is exact without sampling, the only Monte Carlo noise in training comes from the reconstruction term.

The reparameterization trick — move the randomness off the path

Here the implementation hits a wall: "sample $z$ from $q(z \mid x)$" is not differentiable.

Neural nets learn by propagating output error back toward the input through the chain rule (see Backpropagation from Scratch). Put a random number generator partway along that path and the road ends there. Gradient reaches the decoder and cannot be handed to the encoder.

The fix is to push the randomness off the path entirely.

$\odot$ is elementwise multiplication. Draw $\varepsilon$ from a standard Gaussian, scale it by $\sigma$, shift it by $\mu$. The resulting distribution over $z$ is exactly the Gaussian with center $\mu$ and spread $\sigma$ — identical to before. What changed is the shape of the path: $\varepsilon$ is now an input that falls in from outside, unrelated to the network, and between $z$ and $\mu$, $\sigma$ there is nothing but a multiply and an add. Gradient flows straight through. Roll the die and then process it, or roll it midway through processing: same distribution, but only in the first case can you differentiate the processing.

A VAE in code

The heart of one training step is about ten lines.

mu, logvar = encoder(x)                  # emit log σ², not σ
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)              # randomness arriving from outside the net
z = mu + std * eps                       # reparameterization

x_hat = decoder(z)
recon = F.mse_loss(x_hat, x, reduction="sum") / x.size(0)
kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp()) / x.size(0)
loss = recon + beta * kl                 # beta=1 is the plain VAE

Emitting $\log \sigma^2$ rather than $\sigma$ is the practical trick. $\sigma$ must be positive, but a network's output happily goes negative. With $\log \sigma^2$ the whole real line is legal and exp guarantees positivity. It is the kind of stabilization that never appears in the equations and is always required in the code.

Walking the latent space

Encode two images to get $z_1$ and $z_2$, step along the line between them, decode each step: with a VAE you get a smooth morph through the middle, where a plain autoencoder gave you garbage. Subtract the mean $z$ of "faces without glasses" from the mean $z$ of "faces with glasses" and you get a vector pointing along an attribute direction — add it to any face's $z$, decode, and glasses appear. Same structure as arithmetic on word vectors.

There is a high-dimensional trap waiting here. Vectors drawn from a $d$-dimensional standard Gaussian have lengths tightly concentrated near $\sqrt{d}$; the region around the origin is essentially empty. So the midpoint of a straight line between $z_1$ and $z_2$ is noticeably shorter than either end, and the path cuts through origin-adjacent territory the model never visited during training. When only the middle of an interpolation goes mushy, this is usually why — which is why spherical linear interpolation (slerp), which rotates the angle while preserving length, is the standard choice.

Why VAE outputs look blurry

VAEs have long had a reputation for producing blurry images. That is not sloppy implementation; it follows logically from the choice of loss.

Using squared error for the reconstruction term is equivalent to assuming $p(x \mid z)$ is a Gaussian with fixed variance. And when several outputs are plausible for a given $z$, the maximum-likelihood answer under a Gaussian is their mean. If an edge could be here or there, the model draws the average of both — a soft edge in between. Picking one of the crisp candidates would cost more expected squared error than splitting the difference.

The answer to this averaging problem is essentially the history of generative modeling since. GANs handed the crispness judgment to an adversarial discriminator; VQ-VAE replaced continuous latents with a discrete codebook; diffusion models avoid the one-shot average by stacking hundreds of small denoising steps (see Diffusion Models from the Ground Up). And yet the VAE did not disappear. In latent diffusion systems like Stable Diffusion it stepped out of the generator role and stayed on as the compressor that folds images into a small latent tensor.

How this shows up on the job

Who touches it, and when. First, anyone running an image generation pipeline: work with Stable Diffusion-family models long enough and you will be swapping AutoencoderKL checkpoints, handling latent tensors, and budgeting VRAM for the VAE decode. Second, anomaly detection engineers — training a VAE on normal data only and scoring anomalies by reconstruction error or ELBO is still in service for visual inspection lines. Third, representation learning, where you compress unlabeled logs or purchase histories into low-dimensional features.

Parameter names you will actually set.

• beta — the weight on the KL term. 1 is the plain VAE; raise it and the latent space becomes tidier and interpolation smoother at the cost of reconstruction fidelity. Tuning this knob explicitly is what the β-VAE line of work is about
• latent width (latent_dim / latent_channels) — the size of the suitcase. Wider reconstructs better, but the latent carries less structure
• free_bits — a floor on each latent dimension's KL, so dimensions do not collapse entirely
• KL warm-up / annealing — a schedule that starts beta at 0 and ramps it up over early training
• scaling_factor — the constant that standardizes VAE latents in latent diffusion. It lives in the model config and differs between models

Failure modes that will bite you.

1. Posterior collapse. The KL loss pins to zero and stops moving: the encoder starts returning the prior regardless of input, so $z$ carries no information. It happens most readily when the decoder is powerful (autoregressive decoders especially), and the VAE degenerates into an unconditional generator that ignores its input. Log the KL term separately — watching only the total loss will hide this. Remedies are KL warm-up, free_bits, and weakening the decoder
2. Mismatched loss scales. Taking the reconstruction error with sum versus mean changes its ratio to the KL term by the number of dimensions — thousands of times over, for images. When you cannot reproduce a paper's numbers, check this first
3. Exploding logvar. If training destabilizes, $\log \sigma^2$ runs away and exp returns inf. Clamping the range is the standard guard
4. Over-trusting anomaly scores. A VAE will often reconstruct inputs it never trained on reasonably well, so a threshold on reconstruction error alone lets anomalies through
5. Forgetting scaling_factor in latent diffusion. The latent variance no longer matches what the noise schedule assumes and outputs fall apart — easy to hit when swapping in your own VAE

Questions you will be asked to answer. "Why is reparameterization necessary?" — to take the gradient of an expectation without pushing it through a sampling operation. "What happens if the KL term hits zero?" — $q(z \mid x)$ has become identical to the prior regardless of input, meaning the latent carries no information at all.

Wrap-up

• A plain autoencoder cannot generate because the space between latent points, and the place to draw from, are both undefined
• Making the encoder emit a distribution ($\mu$ and $\sigma$) instead of a point, plus a penalty pulling it toward the prior, solves both at once
• The ELBO is "reconstruction quality minus drift from the prior" — a plank propping up the uncomputable $\log p(x)$
• The reparameterization $z = \mu + \sigma \odot \varepsilon$ rewrites sampling so that randomness sits off the gradient path
• Blurriness comes from squared error implying a Gaussian that returns the mean, and the effort to escape it runs through GANs, VQ-VAE, and diffusion