Artificial Neural Networks and Deep Learning · Insper

Variational Autoencoders

A compressor with a shaped latent space

Two verdicts on the same model

Judged as an image generator, the VAE lost: blurry samples, beaten by GANs and then by diffusion.
Judged by deployment, every latent diffusion image you have seen was decoded by one.
How can both be true?

  • A VAE is not really a generator. It is a compressor with a shaped latent space.
  • Its blur is a choice of objective — and in production somebody else (a perceptual and an adversarial loss) corrects it downstream.

Start with an autoencoder

An autoencoder: the input image goes through an encoder into a small code, the bottleneck, and a decoder turns the code back into an image that the loss compares with the input. x input image E encoder BOTTLENECK code z a few numbers D decoder o = D(E(x)) reconstruction L = ‖x − o‖²compares the output with the input

An encoder squeezes into a small code, the bottleneck; a decoder brings it back, and the loss compares the reconstruction with . The bottleneck is the whole mechanism: the network can only reconstruct if the code keeps what matters.

Decode a random point

Train it, pick a random point of the latent space and decode it. What comes out?

  • Garbage, usually. Reconstruction only asks that each input map to some code the decoder can invert.
  • The codes spread wherever training pushed them, with holes in between that nobody trained the decoder on.
  • To sample you need a distribution to sample from. An autoencoder gives you points.

The VAE: encode a distribution, not a point

A VAE: the encoder outputs a mean and a log-variance, the code is z = μ + σ ⊙ ε with noise ε drawn from N(0, I), and the decoder reconstructs x̂. The loss has a reconstruction term between x and x̂ and a KL term that pulls each Gaussian toward the prior. x E qφ(z | x) μ(x) mean log σ²(x) log-variance KL( q(z | x) ‖ 𝒩(0, I) ) pulls each Gaussian toward the prior z z = μ + σ ⊙ ε ε ~ 𝒩(0, I) the noise enters as data D pθ(x | z) x̂ reconstruction: ‖x − x̂‖²

The encoder outputs and : a Gaussian per input, . A code is drawn from it and decoded; the KL term pulls every Gaussian toward a prior you can sample from, . : the encoder's weights · : the decoder's · : mean 0 and variance 1 in every latent dimension.

Points versus clouds

Two latent spaces. Left, an autoencoder: tight clusters of codes far apart, and samples from N(0, I) falling in the empty space between them. Right, a VAE: one Gaussian per input, overlapping inside the prior, and the same samples landing inside the clouds. AUTOENCODER · one point per input 𝒩(0, I) The codes go wherever training pushed them: the samples × fall in the holes. VAE · one Gaussian per input 𝒩(0, I) The clouds overlap inside the prior: every × lands where the decoder was trained. a sample from 𝒩(0, I): no code was ever here the same sample, inside a cloud

The Gaussians of nearby inputs overlap. The overlap fills the space between training examples, and that is what makes a sample from land somewhere the decoder knows. Why that loss and not another? The next slides derive it in five steps.

Step 1 · What we would like to maximize

A generative story: draw , then . Maximum likelihood asks for the evidence

How much the model believes in : the E of ELBO.

  • Estimate it by sampling from the prior? Almost every such decodes to some other image: the average is decided by a few lucky draws.
  • The codes that matter for are the ones the posterior points to…
  • …and computing the posterior needs . The problem is circular.

Step 2 · Guess the posterior, then measure the guess

The encoder guesses , one mean and one variance per dimension. Then, with equalities only:

  1. A KL is never negative, so the ELBO is a lower bound on the evidence: the Evidence Lower BOund.
  2. The gap is exactly how wrong the encoder's guess is.
  3. Raising the ELBO shrinks the gap (encoder) and raises itself (decoder): one objective trains both networks.

How it comes out: (no in it) (Bayes); multiply and divide by inside the log. (appendix).

Step 3 · The two terms you will code

Put inside the ELBO of step 2 and it splits into two terms:

term if it were the only one
reconstruction , codes anywhere: an autoencoder
KL to the prior for every , the code says nothing: posterior collapse

Two different KLs. The loss has the KL to the prior (known, closed form). The KL to the posterior is the gap — nobody computes it, and the ELBO never needs it.

A map of the terms

A map of the terms of the ELBO. On the left, the path of an image: x, the encoder q(z | x), the code z, the decoder p(x | z) and the reconstruction x̂. On the right, the KL to the prior, the reconstruction term and the ELBO, one minus the other. Below, the evidence and the posterior, which belong to the model but cannot be computed. q · the approximation we train p · the generative model ALSO PART OF THE MODEL, BUT INTRACTABLE used when generating gap = KL(q ‖ p(z | x)) ≥ 0 x real image Encoder qφ(z | x) guess of the posterior: μ, σ z · latent code z = μ + σ · ε Decoder pθ(x | z) likelihood x̂ reconstructed image KL( qφ(z | x) ‖ p(z) ) term 2 · stay near the prior Prior p(z) 𝒩(0, I) 𝔼q[ log pθ(x | z) ] term 1 · reconstruct the input ELBO = term 1 − term 2 step 5 puts a weight β on term 2 Posterior p(z | x) what q tries to imitate Evidence log pθ(x) the ELBO stays below it

Step 4 · An ELBO you can compute by hand

, decoder , one observation . Here : mean , variance .

  • Evidence: is Gaussian with variance , so .
  • Posterior (Bayes with Gaussians): .
  • For a guess both terms are closed forms:

The 4v comes from 𝔼q[(x − 2z)²] = (x − 2m)² + 4v. It is the tug of war in one symbol: reconstruction wants v → 0, the KL wants v → 1.

Exercise · one row in full

Compute the ELBO of the guess , that is, and . How far below is it?

term formula, with value
reconstruction ?
KL to the prior ?
ELBO reconstruction − KL ?
gap ?
KL to the posterior ?

Help: · · between Gaussians, (appendix)

One row in full · answer

term computation value
reconstruction
KL to the prior
ELBO
gap
KL to the posterior

The last two rows are the same number: the gap is the distance from the guess to the posterior, as step 2 said.

Four guesses side by side

Each row with the formulas of step 4; gap .

guess reconstruction KL to prior ELBO gap
= prior −7.4189 0.0000 −7.4189 4.7953
= point −0.9389 2.9326 −3.8715 1.2479
= guess −1.9189 0.8181 −2.7371 0.1134
= posterior −1.4989 1.1247 −2.6237 0.0000
  • The prior pays no KL and explains nothing: collapse in miniature.
  • The point, at (it decodes exactly to ), reconstructs best — it is the autoencoder — and pays 2.93 nats of KL for it.
  • The posterior wins neither column and has the best ELBO, equal to : the ELBO rewards the best trade, not the best reconstruction.

nat: the unit of a natural log (base e) · 1 nat ≈ 1.44 bits · every log in this deck is natural, so KL, ELBO and log p come out in nats

Step 5 · From the ELBO to the loss you code

  1. The expectation becomes one sample: per image.
  2. becomes a familiar loss: Gaussian decoder → squared error; Bernoulli decoder → binary cross-entropy.
  3. The KL is written in closed form, summed over the latent dimensions.
  4. The KL gets a weight, . The ELBO is ; leaving free is the β-VAE (Higgins et al., 2017).

With this is the negative ELBO up to a constant (a Gaussian decoder of variance , whose is ). Below 1 the reconstruction rules; above it, the prior. That is where the behaviour lives.

x̂: the decoder's output · i runs over the latent dimensions · σ²i = exp(logvari), the output of the second head

In code: shapes and one classic bug

A batch of images flattened to pixels, a -dimensional latent.

tensor shape
x, x_hat input and reconstruction
mu, logvar, eps, z one Gaussian per image, a fresh per image and dimension
reconstruction, KL summed over pixels and over latent dimensions
loss scalar averaged over the batch

reduction="mean" on the reconstruction also averages over the 784 pixels — and silently multiplies by 784.

Simulator · the ELBO in code, sample by sample

guess q = 𝒩(m, v)
m  v
eps · draws of 𝒩(0, 1)

The reparameterization trick

is drawn at random. What is ?

  • As written, nothing: sampling has no derivative, and the encoder would get no gradient from the reconstruction.
  • Move the randomness out of the path: , (: elementwise product).
  • Now and, with (what the head outputs) and , . The noise enters as data.

Kingma & Welling, Auto-Encoding Variational Bayes (arXiv 2013, ICLR 2014); independently, Rezende, Mohamed & Wierstra (ICML 2014). Encoder, decoder and variational bound already existed in the Helmholtz machine (1995): what this line added is a gradient with little noise.

Is the trick necessary? Measured

The score function (REINFORCE) estimates the same gradient without differentiating through : .

Gradient of , , at the guess (, , ): one sample per estimate, 100 000 estimates, std = standard deviation of one estimate (lab, section 2).

estimator mean of std mean of std
exact — —
reparameterized 2.001 2.009 −0.505 0.878
score function 2.032 7.084 −0.527 3.631
  • Both are unbiased. The spread of an average falls like : and — the score function needs 12 to 17× more samples for the same noise, in one dimension.
  • It only sees how good was. The trick also sees which way to move it, through .

Simulator · one VAE pass, forward and back

The tug of war, on a trained network

One draw of can push either way. What does the average do — and what should it do once training has converged?

  • The KL part is the derivative of in . Lab, section 4: a VAE trained on Fashion-MNIST, 256 test images × 16 dimensions = 4,096 entries.
  • Reconstruction part: positive in 71% of the entries, negative in the rest; mean +0.311 — noise costs reconstruction, shrink .
  • KL part: mean −0.335 — widen toward 1. The two nearly cancel: a trained encoder is where the tug of war is a draw.

β is the whole story

Turn from 0 up to a very large value. What happens to the reconstruction, to the KL — and to samples from ?

  • β → 0: an autoencoder. Near-perfect reconstruction, codes anywhere, prior samples land in holes.
  • β in a good band: the latent matches the prior closely enough to sample from, and still carries enough to reconstruct.
  • β too large: posterior collapse. The encoder returns for every , the KL goes to zero, the decoder outputs the dataset average. The loss still looks like it is going down.

Recognising collapse: KL near zero and reconstruction near the "average image" value. Always log the two terms separately.

Simulator · β, trained in your browser

Measured in the lab: β on Fashion-MNIST

16-dimensional latent, 5 epochs each. A classifier judges 2 000 decoded prior samples (real test images: confidence 0.86, 6.72 effective classes).

β rec KL active dims confidence eff. classes
0 219.7 254.5 16 0.49 1.88
0.25 222.4 31.7 16 0.62 3.27
1 230.6 15.8 14 0.67 3.86
4 248.0 6.8 7 0.70 4.34
16 279.4 2.7 3 0.66 4.02
64 365.0 0.3 1 0.43 1.63
256 385.4 0.0 0 0.33 1.00

rec and KL in nats per image (rec: cross-entropy summed over the 784 pixels; the average image for everything gives 385.0) · active dims: with KL > 0.1 nat · confidence: the classifier's top probability · eff. classes: from 1 (every sample alike) to 10 · −ELBO = rec + KL is smallest at β = 1: 230.6 + 15.8 = 246.4.

Collapse with β = 1: a strong decoder

Give the decoder a hint: the image itself, average-pooled, next to . β stays 1.

hint rec KL
none 230.6 15.83
4×4 228.5 7.41
7×7 223.0 1.82
14×14 213.2 0.10
28×28 204.8 0.02
  • The more the decoder can do alone, the less the code is worth its KL.
  • With the whole image the prior is the optimum — nothing failed.
  • Free bits (λ = 0.5 nat: a dimension's KL below λ is not charged) on the 14×14 hint: KL 0.10 → 6.94 and rec 213.2 → 211.9, but −ELBO = rec + KL goes from 213.3 to 218.9. It keeps the code alive by leaving the optimum.

The gap, on a real VAE

Average the weights before the log, over samples: still a lower bound, and tighter (Burda et al., 2016).

nats per image (500 test images)
ELBO −243.72
, samples from −241.06
, samples from the prior −267.40
ELBO with optimized per image −241.93
  • From the prior, the estimate is useless: −267.40 is 23.7 nats below the ELBO itself — the circular problem of step 1, measured.

Where the gap comes from

−244−243−242−241nats per imagesamples from the prior: −267.40−243.72ELBO−241.93ELBO, q per image−241.06L1000log p(x): here or to the rightamortization ≥ 1.79approximation ≈ 0.87gap ≥ 2.66 nats

  • The gap nats.
  • Amortization : one encoder for all images; optimizing image by image recovers it.
  • Approximation : no diagonal Gaussian, however well chosen, is this decoder's true posterior. An estimate, not a bound: both of its ends are lower bounds.

Why the samples are blurry

The decoder is trained with . Several sharp images are consistent with the same code. What is the loss-optimal output?

  • Their average — the optimum of squared error is the conditional mean. Averages of images are blurry.
  • The ELBO is maximum likelihood, i.e. forward KL, which is mode-covering: when the model cannot represent the data exactly, it puts mass between the modes.
  • No bug: the objective got exactly what it asked for. So change the objective.

The fix that production uses

  • Perceptual: compares deep features ( here is a pretrained network, not the encoder), so a texture shifted by one pixel is not punished.
  • Adversarial: a GAN discriminator on the decoder output punishes exactly the smoothness rewards.
  • This is the VQGAN / latent-diffusion recipe: every latent autoencoder under a diffusion model is trained this way.

A convolutional VAE: the latent is a small image

A convolutional autoencoder: a 28×28×1 image goes through two stride-2 convolutions to 14×14×32 and 7×7×64, then to a latent of 7×7×4 that is still a small image with channels; transposed convolutions bring it back to 28×28×1. x · 28×28×1 14×14×32 7×7×64 z · 7×7×4 7×7×64 14×14×32 x̂ · 28×28×1 ENCODER each layer halves the resolution and adds channels DECODER transposed convolutions bring it back LATENT a small image with channels Stable Diffusion's VAE has the same shape: 512×512×3 → 64×64×4, 48× fewer numbers.
  • Every production latent autoencoder is convolutional: each stride-2 layer halves the resolution and adds channels; the decoder mirrors it with transposed convolutions.
  • The latent keeps the layout of the image, so the generative model can run on it as if it were a picture (a U-Net or a DiT).

What a VAE is actually for

Latent diffusion runs in the VAE's compressed space, never on pixels.

at 512 × 512 pixels SD 1.5 / SDXL latent SD 3 / FLUX latent
shape 512 × 512 × 3 64 × 64 × 4 64 × 64 × 16
numbers 786,432 16,384 65,536
compression 1× 48× = 786,432 / 16,384 12× = 786,432 / 65,536
  1. It caps the quality: what the encoder drops, diffusion cannot generate (small text, fine texture — hence 16 channels).
  2. Its β is tiny (~): a compressor, never sampled from its prior.
  3. It is frozen: one latent space shared by a whole ecosystem of fine-tunes.

The other branch: discrete latents

VQ-VAE snaps each latent vector (the encoder's output) to the nearest entry of a learned codebook: .

KL-regularized (continuous) VQ (discrete)
latent real-valued vectors codebook indices
used by latent diffusion: SD, SDXL, SD3, FLUX autoregressive and masked: VQGAN, MaskGIT, Chameleon
failure mode blur; limited capacity per channel codebook collapse
why diffusion needs somewhere to add noise Transformers need tokens to predict

Once an image is a grid of integers it is a sequence of tokens — the subject of the autoregressive generation chapter.

Appendix · the KL divergence

  1. Never negative, zero only when .
  2. Not symmetric: — not a distance.
  3. The integrand can be negative where — so a one-sample estimate can be too (lab: 8.6% of them for posterior vs prior). That is why the loss uses the closed form.

With and : , the VAE term.

Simulator · two Gaussians to drag

Lab · the ELBO on a real VAE

Colab notebook, ~2 h, CPU, Fashion-MNIST. Each section ends in a checkpoint with a check value.

  1. the evidence by sampling, and Jensen
  2. the trick against autograd and against the score function
  3. the KL three ways
  4. a VAE on Fashion-MNIST, and the gradient at its heads
  1. β, measured: reconstruction, KL and samples
  2. collapse with a strong decoder, and free bits
  3. the gap on a real VAE: importance weights and amortization
  4. your turn: does the best likelihood give the best samples?

open in Colab →

Takeaways

  1. An autoencoder compresses; it does not generate: points, not a distribution.
  2. : a lower bound whose gap is the encoder's error. One objective trains both networks.
  3. ELBO = reconstruction − KL to the prior. It rewards the best trade.
  4. The reparameterization trick makes it differentiable — and far less noisy than the score function.
  5. β slides between an autoencoder and posterior collapse; a strong decoder collapses even at β = 1. Log both terms.
  6. Blur is the objective: fixed with perceptual and adversarial losses.
  7. In 2026 a VAE is the frozen compressor under latent diffusion; VQ-VAE is the discrete branch.

Outline: (1) why an autoencoder cannot generate; (2) the ELBO in five steps, with the worked example; (3) the trick and one pass by hand; (4) beta, collapse and the gap, measured; (5) blur, and what a VAE is for in 2026. Closes with the Colab lab and the KL appendix.

In the lab (section 7): 1000 samples from the prior give a bound 23.7 nats worse than the ELBO itself.

Derivation: E_q[log p(x)] = E_q[log p(x,z)/p(z|x)], multiply and divide by q. Every line is an equality.

Step through the five stages: the generative story, what cannot be computed, the encoder's guess, the two terms, and the gap. Hover a box to see its arrows. Purple is q, the one we train; green is the model p.

In the in-page lab: one sample is unbiased; the spread falls like 1/sqrt(K): 1.29, 0.38, 0.12, 0.04 for K = 1, 10, 100, 1000. A batch of 128 images already averages 128 draws.

Default case: q_C, the row of the exercise. Step line by line: the exact column on the right is the table of the previous slides (ELBO −2.7371, KL 0.8181, gap 0.1134); the four ε give elbo −2.7455, L_K −2.6401, gap 0.1219. Then: "the posterior" (every log w = −2.6237, gap 0), "the prior" (log w from −1.90 to −16.60, elbo −7.40 but L_K −2.73: averaging w first rescues a bad q), "the point" (narrow q: the reconstruction barely moves, −0.92 to −0.95, but the gap is 1.25). K = 1 makes L_K equal to elbo.

Use the Next button (the arrow keys step the diagram while it has focus; click outside it to move to the next slide). Steps 2, 3, 5, 6 are an ordinary MLP. Step 4 is the only new operation. At step 10, drag epsilon: the reconstruction part of dL/ds changes sign, its average over epsilon is +0.100.

In the worked example: beta -> 0 leads to the point q_B, beta -> infinity to the prior, beta = 1 to the posterior.

Eight clusters on a ring, a 2-D latent. Headless runs of this panel: beta = 0 -> 90% of prior samples off the data, KL 10.9 nats; beta = 1 -> 21%, KL 1.85; beta = 7.5 or more -> collapse, every sample off the data and a reconstruction error of 0.52, the variance of the data (0.53).

Presets: the worked example's posterior against the prior gives 1.1247, the KL column of its table. Narrow q inside p: 0.318 one way, 0.807 the other. Far apart with equal widths: 8 both ways. Click outside the panel before using the arrow keys to change slide.