Variational Autoencoders (VAEs)

Start from something you know: a plain autoencoder. In one sentence: it squeezes an input through a narrow bottleneck and unpacks it again, training itself by "the unpacked thing must match the original".

It compresses beautifully. And it is useless as a generator. Why?

Because nothing forces the bottleneck coordinates to be organized. Train one on faces and the latent space ends up a lumpy scatter:

code
Copy
latent space of a plain autoencoder:

    ·        x        ·
  x    ·        ·   x        x = decode(random point) = garbage
        ·     x              · = a real training face
  ·        ·        ·
             x   ·

Pick a random latent vector at random and decode it — between the training points there is no law. You get static, or a half-melted face. You cannot even walk from a smile to a frown by moving a little; the terrain is broken.

So the question becomes

That is exactly what the Variational Autoencoder (Kingma & Welling, 2013) sets out to build: not just a compressor, but a generative latent space.

A VAE is an autoencoder with one philosophical swap:

Three pieces:

• The encoder becomes an amortized approximate posterior q_phi(z|x): given input x, it emits the parameters of a Gaussian — a mean mu and a (log-)variance sigma^2. For a 2-d latent, that is 4 numbers: 2 for where the cloud centers, 2 for how puffy it is.
• The decoder becomes a likelihood p_theta(x|z): given a latent sample, what image is plausible?
• A prior p(z) = N(0, I): the shape we demand of the latent space — a single smooth bell centered at origin. Sampling z ~ q(z|x) and decoding it yields a reconstruction. Because the prior is smooth and the KL term penalizes posteriors that drift from it, nearby latent codes decode to nearby, meaningful images — an interpolatable latent space, the property people actually want for generation.

Arrow chain of the whole model:

x → encoder → (mu, sigma) → sample z → decoder → x-hat, plus: pull (mu, sigma) toward N(0, I).

Why does forcing every input's cloud toward one shared bell make the map smooth? Because overlapping clouds means no empty wastelands between them — the dice-random point now falls inside some cloud's territory, and decodes to something real-ish.

Take one training image — a smiling face. The encoder sets the dials:

mu = [2.0, 0.0], sigma = [0.1, 1.0]

Translation: "smiling lives near coordinate 2 on axis one — and I am sure about that (small sigma); on axis two I have no opinion (sigma = 1, as wide as the prior)."

Step 1 — draw a latent sample with the reparameterization trick. Roll standard-normal dice: epsilon = [-1.28, 0.55]. Then

z = mu + sigma * epsilon = [2.0 + 0.1·(-1.28), 0.0 + 1.0·0.55] = [1.87, 0.55]

Sampling became arithmetic: a fixed function of mu, sigma, and an external dice-roll. Keep that in your pocket — Section 4 explains why it is the linchpin.

Step 2 — the KL term, with numbers. The penalty KL( q(z|x) || p(z) ) for a Gaussian vs the N(0, I) prior is computable in closed form. Axis one (mu=2, sigma=0.1):

KL = ln(1/0.1) + (0.1² + 2² - 1)/2 ≈ 2.30 + 1.51 = 3.81

— expensive! This cloud is tiny and parked far from the origin. Axis two (mu=0, sigma=1):

KL = ln(1/1) + (1 + 0 - 1)/2 = 0

— free; that axis is the prior. Training pushes axis one: slide mu back toward 0, or puff sigma up — pay reconstruction quality or pay KL. That tension is the dial.

Step 3 — the loss itself. ELBO = reconstruction − KL:

code
Copy
ELBO = E_{q(z|x)}[ log p(x|z) ]  −  KL( q(z|x) || p(z) )
         └── decode me back! ──┘    └── keep the clouds smooth ──┘

Maximize both? You cannot fully — squeezing the clouds to exactly N(0, I) would erase the information needed to reconstruct. The trained VAE sits at a chosen compromise between the two terms (beta-VAE, Section 6, makes that dial explicit).

Before (plain autoencoder) and after (VAE), each dot is one training image's latent location:

code
Copy
BEFORE: dots with dead space between     AFTER: overlapping clouds around origin
  ·          ·                              (   ·  )
      ·  ·        ·                           ( ··☁·· )     · = sampled z
 ·       ·   ·         ✗ random point         (  ·  )       ☁ = its cloud
   ·  ✗           ·     decodes to junk          ✗← every z lands in *some* cloud
                                              ────┼────
                                            N(0,I): one shared bell

The walk now works: pick two faces, average their z-vectors, decode — you land on a plausible face halfway between the expressions. That is "interpolatable latent space", and it is what the cloud-overlap discipline bought you.

But one plumbing detail got in the way. The pipeline literally includes sampling a dot from a cloud. And gradient descent cannot pass through a dice roll:

code
Copy
naive:  x → encoder → cloud → [ ROLL THE DICE ] → z → decoder → loss
                             ↑
                  gradient hits the randomness and STOPS.
                  the encoder never learns.

fixed:  x → encoder → (mu, sigma) ──┐
      dice: epsilon ~ N(0,I) ───────┼→ z = mu + sigma·epsilon → decoder → loss
                                    ↑
                  z is now plain arithmetic; the gradient flows
                  back to mu and sigma. epsilon just carries noise.

That single algebraic rewrite — pulling the dice outside the differentiable path — is the reparameterization trick, and the entire trainability of VAEs rests on it.

Carry this through: the VAE is a portrait studio with two staff and one rule.

• The encoder is the agent who looks at each sitting client and writes down dial settings (mu) plus how uncertain they are (sigma) — "nose-height: 2.0 ± 0.1".
• The decoder is the painter who can only work from dial settings: give it any settings and it produces a portrait.
• The KL prior is the studio's house style: every client's dials must look like they came from the same control panel's normal range — no private dial positions. Consequence: on any day, a technician can turn the dials at random and the painter still produces a plausible face. That is generation for free — the lumpy-scatter studio could only repaint booked clients.
• The reparameterization trick = the painter no longer waits for the dice inside the agent's office. The agent writes numbers; the dice live at the painter's desk; the settings combine with the rolls by simple addition. Because the combining is arithmetic, feedback ("too much blush, lower dial 3") flows all the way back to the agent's note-taking. Without the trick, feedback dies at the dice.
• Blurry output = a painter trained to minimize average mistake paints the average of everything plausible — a soft, generic face. Great for compression, not for sharpness.
• Posterior collapse = a painter so good they can restore any portrait from the client's name alone; the dials get ignored, and the KL house style wipes them to pure random. The agent stops contributing. (Fixes in Section 6.)

Two failure modes to know cold:

• Blurry samples: decoders typically model each pixel as an independent Gaussian, so the reconstruction is a conditional mean that averages plausible outcomes — sharper than nothing, softer than GANs/diffusion.
• Posterior collapse: with a powerful autoregressive decoder, the model can ignore z entirely; the KL term drives q(z|x) onto the prior and the latent contributes nothing. Fixes include KL annealing and free-bits (a floor on the per-dimension KL).

Rich extension family — each one turns a dial from the studio story:

• beta-VAE reweights KL by beta, trading reconstruction for disentangled, interpretable factors (stricter house style → cleaner dial meanings).
• Conditional VAE conditions encoder/decoder on labels y for targeted generation ("paint me a frown").
• Normalizing-flow VAEs / NVAE replace the Gaussian posterior with flexible ones to beat the blurry-posterior limitation.
• VQ-VAE swaps the continuous Gaussian latent for a discrete codebook learned via vector quantization + a straight-through estimator, yielding sharp, tokenizable latents. Plain words: instead of dial positions, the studio owns a fixed palette of a few thousand codes; every patch is snapped to the nearest palette entry — integers a transformer can predict.
• VQGAN adds adversarial and perceptual losses to the quantizer for photorealistic reconstruction.

def reparameterize(mu, log_var):
    std = torch.exp(0.5 * log_var)          # sigma
    eps = torch.randn_like(std)             # epsilon ~ N(0, I)
    return mu + eps * std                   # differentiable z

def vae_loss(recon, x, mu, log_var):
    recon_loss = F.mse_loss(recon, x, reduction="sum")
    kl = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
    return recon_loss + kl

The VAE's biggest modern impact is as the compression stage of latent diffusion (next topic's star player). Running a diffusion process directly on a 512x512x3 image is prohibitively expensive; Stable Diffusion instead trains a VQ-style autoencoder that compresses pixels into an 8x-smaller continuous latent (64x64x4), runs all denoising there, and decodes once at the end. The autoencoder's reconstruction quality and latent structure directly bound the fidelity of every downstream image.

The same encoder-decoder idea underpins tokenizers for autoregressive image/video models: a learned codebook turns pixels into integer tokens the transformer can predict.

So while "pure" VAEs have faded as headline image generators, the variational, reparameterized, codebook-based autoencoder is quietly the front door of nearly every 2024-2026 diffusion and unified multimodal stack.

Compare with the GAN you just met, in one table:

code
Copy
              GAN                     VAE
training      adversarial game        single smooth loss (ELBO)
samples       sharp, risky diversity  soft, guaranteed coverage
likelihood    none written down       tractable lower bound
modern role   heads, vocoders         latent compression backbone