Variational Autoencoders¶
Summary — A VAE learns an encoder that maps data to a distribution over latent codes and a decoder that maps codes back to data, trained by maximizing the ELBO. The reparameterization trick makes the sampling step differentiable. VAEs give you a structured latent space and fast one-shot sampling at the cost of blurry outputs — which is why they survive today not as standalone generators but as the compressor inside latent diffusion and as the tokenizer (VQ-VAE) in front of autoregressive image, audio and video models.
Prerequisites: → Probability & information theory · Next: → GAN
1. The setup¶
Assume data is generated by a two-step process:
\(z\) is a low-dimensional latent code (the "concept"), \(x\) the observed data (the "rendering"). To train by maximum likelihood we need
which is intractable — it integrates a neural network over a \(d_z\)-dimensional space.
Why not Monte Carlo?
\(p_\theta(x) \approx \frac{1}{K}\sum_k p_\theta(x \mid z_k)\) with \(z_k \sim p(z)\). In high dimensions almost every random \(z\) decodes to something nothing like \(x\), so almost every term is ≈ 0. You'd need astronomically many samples. The fix is importance sampling from a learned proposal \(q_\phi(z\mid x)\) — which is exactly the encoder.
2. The ELBO, and what each half wants¶
That learned proposal \(q_\phi(z\mid x)\) is the encoder, and plugging it into the intractable likelihood gives you something you actually can optimize: the ELBO.
(full derivation: → Probability & information theory §7)
ENCODER DECODER
q_φ(z|x) = N(μ,σ²) p_θ(x|z)
┌──────────────────────┐ ┌────────────────┐
x ──────────────►│ conv/MLP ──► μ(x) │ │ │
(784-dim) │ ──► σ(x) │ │ deconv/MLP │──► x̂
└──────────┬───────────┘ └───────▲────────┘
│ │
│ z = μ + σ⊙ε │
└──────────►(+)────────────┘
▲
ε ~ N(0,I)
(reparameterization)
LOSS = ‖x − x̂‖² + β · KL(N(μ,σ²) ‖ N(0,I))
"be accurate" "be close to the prior"
← pulls σ→0, μ spread → pushes σ→1, μ→0
The tension, stated plainly
The reconstruction term wants each input to get its own private, tightly-concentrated region of latent space (\(\sigma \to 0\), codes far apart): that makes decoding unambiguous. The KL term wants every input to map to the same standard Gaussian (\(\mu \to 0, \sigma \to 1\)): that makes the latent space smooth and samplable. The equilibrium is a latent space where codes overlap just enough that sampling \(z \sim \mathcal{N}(0,I)\) lands somewhere the decoder understands.
Without the KL term you have a plain autoencoder: great reconstruction, but the latent space is full of holes. Sample a random \(z\) and you decode garbage. The KL term is what makes it generative.
The rate–distortion reading¶
The KL term is literally the rate: the expected number of nats needed to transmit \(z\) under a code optimized for the prior. The reconstruction term is the distortion. A VAE is a learned lossy compressor, and \(\beta\) (below) slides you along the rate–distortion curve.
3. The reparameterization trick¶
The ELBO is a well-defined objective, but training it means backpropagating through a sampling step — \(z \sim q_\phi(z\mid x)\) — and gradients don't flow through randomness by default. The reparameterization trick is what makes that step differentiable at all.
We need \(\nabla_\phi \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x\mid z)]\), but the distribution being sampled from depends on \(\phi\). Sampling is not differentiable.
Solution: move the randomness to a parameter-free source.
Now \(z\) is a deterministic, differentiable function of \(\phi\) and a fixed random variable.
✗ BLOCKED ✓ FLOWS
φ ──► (μ,σ) ──► [SAMPLE] ──► z φ ──► (μ,σ) ──┐
▲ ├──► z = μ+σ⊙ε
stochastic node; ε ─────┘
no gradient path (constant w.r.t. φ)
Why it's correct — it is a change of variables inside the expectation:
The right-hand side's distribution no longer depends on \(\phi\), so \(\nabla_\phi\) passes straight inside.
Variance: the reparameterized estimator typically has variance orders of magnitude below REINFORCE, which is the difference between "trains in an hour" and "does not train."
The standard implementation detail: the encoder outputs \(\log\sigma^2\), not \(\sigma\), so that the output is unconstrained (any real number) and \(\sigma = \exp(\frac12\log\sigma^2) > 0\) automatically.
mu, logvar = self.encoder(x).chunk(2, dim=-1)
std = torch.exp(0.5 * logvar)
z = mu + std * torch.randn_like(std) # reparameterize
kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)
4. The loss, concretely¶
With a differentiable sampling step in hand, the ELBO turns into two ordinary terms you can compute and add up on real data — which is worth doing by hand once, numbers included.
The choice of reconstruction term is a choice of decoder likelihood:
| Decoder likelihood | Reconstruction loss | Use for |
|---|---|---|
| \(\mathcal{N}(\hat x, \sigma^2 I)\), fixed \(\sigma\) | MSE | continuous data (default) |
| Bernoulli | binary cross-entropy | binarized MNIST |
| Categorical (256-way per channel) | cross-entropy | discretized pixels |
| Laplace | L1 | robustness to outliers |
| + LPIPS + adversarial | perceptual + GAN | latent diffusion VAEs |
This is the blurriness, explained
With a Gaussian likelihood, the optimal decoder output is \(\mathbb{E}[x \mid z]\) — the average of all images consistent with \(z\). If \(z\) doesn't pin down the exact location of an edge, the average of "edge at pixel 10" and "edge at pixel 11" is a smeared edge at 10.5. MSE literally rewards hedging.
Two plausible reconstructions MSE-optimal output = their mean
██▓▓░░ ░░▓▓██ ▓▓▒▒▒▒▓▓
sharp, at left sharp, at right blurry, in between
← lower MSE than either!
The fixes, in increasing order of effectiveness
- Use a richer likelihood (discretized logistic mixture) — helps a little.
- Use a perceptual loss (LPIPS: distance in a pretrained VGG's feature space) — helps a lot, because feature-space averaging is not pixel-space averaging.
- Add an adversarial loss on the decoder — a discriminator punishes blur directly. This is what the Stable Diffusion VAE does, and it is why its reconstructions are sharp.
- Use discrete latents (VQ-VAE) + an expressive autoregressive prior — moves the "which of the plausible images" decision to a model that can represent multimodality.
5. The two classic failure modes¶
Blur is one specific symptom of a VAE's Gaussian likelihood. It isn't the only place where the theory and the practice diverge — the encoder and the prior can also drift apart in a way that has nothing to do with reconstruction quality.
Posterior collapse¶
The KL term is minimized at \(q_\phi(z\mid x) = p(z)\) — i.e. the encoder ignores \(x\) entirely. If the decoder is powerful enough to model \(p(x)\) on its own (e.g. an autoregressive decoder), this is a perfectly good local optimum: KL = 0, reconstruction handled by the decoder alone.
Result: \(I(x; z) = 0\). The latent is dead. You have trained an expensive autoregressive model with a useless noise input.
| Fix | Mechanism |
|---|---|
| KL annealing | ramp \(\beta\) from 0 → 1 over training; let reconstruction establish a useful code first |
| Free bits | \(\max(\lambda, \text{KL}_j)\) per dimension — no gradient pressure below \(\lambda\) nats |
| Weaken the decoder | limit capacity / receptive field so it needs \(z\) |
| Skip connections from \(z\) | inject \(z\) at every decoder layer |
| Discrete latents (VQ) | no KL term to collapse at all |
Free bits, the most robust fix:
kl_per_dim = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp()) # (B, d_z)
kl = torch.clamp(kl_per_dim, min=free_bits).sum(dim=1).mean() # λ ≈ 0.03–0.5 nats/dim
Holes in the prior (the aggregate posterior mismatch)¶
The KL term only constrains \(q_\phi(z\mid x)\) per example. What matters for sampling is the aggregate posterior \(q_\phi(z) = \mathbb{E}_{x}[q_\phi(z\mid x)]\). These can differ a lot.
Prior p(z) = N(0,I) Aggregate posterior q(z)
╭───────────╮ ╭───────────╮
│ ░░░░░░░ │ │ ██ ░ ██│ ← clusters + gaps
│ ░░░░░░░░░ │ vs │ ██ ██ │
│ ░░░░░░░ │ │ ██ ░░ █ │
╰───────────╯ ╰───────────╯
sample from here but the decoder only knows here
⇒ samples from the gaps decode to garbage
Fixes: a learned prior (VampPrior, a flow, or — most effectively — a second generative model trained on the latents). That last option is precisely latent diffusion: train a VAE, then train a diffusion model on \(q_\phi(z)\) instead of assuming it is Gaussian. → Latent diffusion
6. Beta-VAE and disentanglement¶
Posterior collapse and the prior-hole problem are both failure modes to fix. A more ambitious question is whether you can shape the latent space on purpose — for instance, so that each dimension corresponds to one independent factor of variation.
| \(\beta\) | Behaviour |
|---|---|
| \(\beta = 0\) | plain autoencoder: sharp reconstructions, unsamplable latent |
| \(\beta = 1\) | the true ELBO |
| \(\beta > 1\) | more disentangled, more compressed, blurrier |
| \(\beta \to \infty\) | posterior collapse |
Intuition
Why large \(\beta\) encourages disentanglement — a large KL penalty forces the model through a narrow information channel. To reconstruct well under a tight bit budget, the encoder must use the most efficient code, and for data generated by independent factors (pose, lighting, identity) the most efficient code is one that allocates a separate dimension to each factor. The isotropic Gaussian prior additionally prefers axis-aligned, independent dimensions.
The honest caveat
Locatello et al. (2019), Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations, proved that unsupervised disentanglement is impossible without inductive biases, and showed empirically that reported gains were largely attributable to random seed variation and unprincipled model selection. Treat disentanglement claims sceptically; if you need disentangled factors, use some supervision.
7. VQ-VAE: discrete latents¶
Disentanglement asks the latent space to be structured in a particular continuous way. A more radical change is to make it discrete altogether — trading the Gaussian prior for a learned codebook, and the VAE's usual role for a completely different one.
Replace the continuous Gaussian latent with a lookup into a learned codebook \(\mathcal{E} = \{e_1,\dots,e_K\} \subset \mathbb{R}^{d}\).
x ──►[encoder]──► z_e ──► ┌─────────────┐ ──► z_q ──►[decoder]──► x̂
(H,W,3) (h,w,d) │ CODEBOOK │ (h,w,d)
│ e₁ e₂ … e_K │
└─────────────┘
nearest-neighbour lookup
output is a grid of INTEGERS (h,w) ∈ {1..K}
The loss (three terms):
where \(\mathrm{sg}[\cdot]\) is stop-gradient. Term 2 moves codebook vectors toward encoder outputs; term 3 stops the encoder from drifting away from the codebook (\(\beta \approx 0.25\)).
Straight-through estimator — \(\arg\min\) has zero gradient everywhere, so copy the gradient across the quantization boundary:
The compression numbers — a \(256\times256\times3\) image → a \(32\times32\) grid of indices into a 512-entry codebook:
A 170× compression into 1024 discrete tokens. Now run a Transformer over those 1024 tokens — an entirely tractable sequence length. That two-stage recipe is DALL·E 1, VQGAN, Parti, MusicLM, and most video generation.
| Model | What VQ gave it |
|---|---|
| VQ-VAE-2 | hierarchical codes → high-res images with a PixelCNN prior |
| VQGAN | + adversarial and perceptual loss → sharp, low token count |
| DALL·E 1 | text tokens + image tokens in one autoregressive stream |
| EnCodec / SoundStream | audio → tokens → language model over audio |
| Modern video models | spatio-temporal VQ + Transformer |
Codebook collapse
Most entries go unused; effective vocabulary shrinks to a handful. Fixes: EMA codebook updates, restarting dead codes to random encoder outputs, low-dimensional codes with L2 normalization, or FSQ (finite scalar quantization), which replaces the learned codebook with fixed per-dimension rounding and sidesteps the problem entirely.
8. Implementation¶
All of that — ELBO, reparameterization, the failure modes, VQ — comes together in surprisingly little code. Training one end to end on MNIST is worth doing once, both because it works and because the numbers it produces tell you what a healthy VAE should look like.
A complete VAE on MNIST, end to end:
import torch, torch.nn as nn, torch.nn.functional as F
class VAE(nn.Module):
def __init__(self, x_dim=784, h=512, z_dim=20):
super().__init__()
self.enc = nn.Sequential(nn.Linear(x_dim, h), nn.ReLU(),
nn.Linear(h, h), nn.ReLU(),
nn.Linear(h, 2 * z_dim)) # -> mu, logvar
self.dec = nn.Sequential(nn.Linear(z_dim, h), nn.ReLU(),
nn.Linear(h, h), nn.ReLU(),
nn.Linear(h, x_dim)) # -> logits
def forward(self, x):
mu, logvar = self.enc(x).chunk(2, dim=-1)
logvar = logvar.clamp(-8, 8) # numerical safety
z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)
return self.dec(z), mu, logvar
def loss_fn(x_logits, x, mu, logvar, beta=1.0):
# Bernoulli likelihood: sum over pixels, mean over batch
recon = F.binary_cross_entropy_with_logits(x_logits, x, reduction='none').sum(1).mean()
kl = (-0.5 * (1 + logvar - mu.pow(2) - logvar.exp())).sum(1).mean()
return recon + beta * kl, recon, kl
# training loop
model, opt = VAE(), None
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(20):
for x, _ in loader:
x = x.view(x.size(0), -1)
x_logits, mu, logvar = model(x)
# KL annealing over the first 10 epochs guards against collapse
beta = min(1.0, epoch / 10)
loss, recon, kl = loss_fn(x_logits, x, mu, logvar, beta)
opt.zero_grad(); loss.backward(); opt.step()
# sampling: no encoder involved
with torch.no_grad():
z = torch.randn(64, 20)
samples = torch.sigmoid(model.dec(z)).view(-1, 1, 28, 28)
Expected numbers on binarized MNIST, \(d_z = 20\), after ~20 epochs: reconstruction ≈ 70 nats, KL ≈ 25 nats, total ≈ 95 nats/image. Reported state-of-the-art ELBOs are around 80 nats; PixelCNN-decoder VAEs reach ~78.
Sanity check on the KL value: 25 nats = 36 bits of information in the latent code. For MNIST that is plausible — digit identity (3.3 bits) plus ~30 bits of style. If your KL is near 0, you have posterior collapse. If it is near the reconstruction loss, \(\beta\) is too low.
9. Where VAEs actually live today¶
Those healthy-looking numbers come from a toy 20-dimensional MNIST model. At the scale modern systems actually operate, the VAE covered in this page is rarely the end product anymore — it's a component inside something bigger.
| Role | System | Why a VAE |
|---|---|---|
| Latent compressor | Stable Diffusion, SDXL, Flux | 8× spatial downsample → 64× fewer pixels to diffuse. This is the single most important VAE deployment. |
| Tokenizer | DALL·E 1, VQGAN, Parti, video models | discrete codes → a Transformer can model images/video |
| Audio codec | EnCodec, SoundStream, DAC | residual VQ → audio as a token stream |
| Molecule design | JT-VAE, chemical VAEs | continuous latent enables gradient-based optimization of properties |
| Anomaly detection | industrial, medical | high reconstruction error = out of distribution |
| World models | Dreamer, model-based RL | compact latent state for planning |
The lesson worth generalizing
VAEs lost as generators but won as representation learners. Splitting a hard generative problem into "compress to a good space" + "generate in that space" is one of the most reusable architectural patterns in the field.
10. Exercises¶
Problem 1 — Gaussian KL, new numbers. Compute \(D_{KL}(\mathcal{N}(-1.5, 0.25)\|\mathcal{N}(0,1))\) in nats using the boxed formula from §2 (note: variance \(0.25\), so \(\sigma=0.5\)). Is the resulting latent code "expensive" (far from the prior) or "cheap"? Compare with the §8 sanity checks (\(\mu{=}0,\sigma{=}1 \to 0\) nats; \(\mu{=}2,\sigma{=}1\to2\) nats, from the info-theory page).
Solution
Comparing directly in nats against the \(\mu{=}2,\sigma{=}1\) reference point from → Probability & information theory §8 (which costs exactly 2 nats): this code, at \(1.443\) nats, is actually cheaper — even though \(|\mu|=1.5\) is a large shift, because \(\sigma=0.5\) is tighter than the prior's \(\sigma=1\), and the \(-\log\sigma^2\) term rewards tight (confident) codes even as the \(\mu^2\) term penalizes the shift. This is exactly the tension described in §2: the decoder wants tight, informative codes (small \(\sigma\)) while the KL term pulls toward \(\sigma{=}1,\mu{=}0\).
Problem 2 — posterior collapse, diagnose it. After training, you observe: reconstruction loss \(=210\) nats, KL \(=0.003\) nats, and samples from the prior look reasonable but bear no relationship to any specific input when you check reconstructions (every input reconstructs to roughly the same generic output). Using §5, what's happening, and name two fixes with their mechanism.
Solution
\(KL\approx0\) is the signature of posterior collapse (§5): the encoder has learned to ignore \(x\) and output \(q_\phi(z|x)\approx\mathcal{N}(0,I)\) regardless of input, satisfying the KL term at zero cost — while the decoder, if expressive enough, learns to model \(p(x)\) unconditionally (hence the generic reconstructions and the fact that prior samples still "look reasonable" — the decoder is a passable unconditional generator, just not a conditional one).
Fix 1 — KL annealing: ramp \(\beta\) from 0→1 over training, so early on the model is rewarded purely for reconstruction and is forced to use \(z\) to succeed, before the KL penalty has a chance to zero it out.
Fix 2 — free bits: clamp the per-dimension KL to a floor \(\lambda\)
(torch.clamp(kl_per_dim, min=free_bits)), which removes the gradient pressure to shrink KL
below \(\lambda\) nats/dim — so the optimizer can't reach the zero-information solution even if it
would otherwise prefer to.
Problem 3 — VQ-VAE compression, a different codebook. Redo §7's compression-ratio calculation for a \(128\times128\times3\) image compressed to a \(16\times16\) grid of codes from a 1024-entry codebook (note: different image size, grid size, and codebook size from the worked example). What's the compression ratio, and how does the token count compare to the original \(256\times256\)/\(32\times32\)/512-entry example?
Solution
Raw bits: \(128\times128\times3\times8 = 393{,}216\) bits. Encoded bits: \(16\times16\times\log_2(1024) = 256\times10 = 2{,}560\) bits.
Slightly less compression than the worked example's \(170\times\) — because this codebook is larger (1024 entries vs 512), costing 10 bits/code instead of 9, which pushes toward less compression even though the smaller image (4× fewer raw pixels) and smaller grid (\(16^2{=}256\) vs \(32^2{=}1024\) codes, i.e. 4× fewer codes) scale proportionally and roughly cancel that out. Net: \(153.6\times\) vs \(170\times\) — comparable order of magnitude, and in both cases the resulting token count (256 vs 1024) is small enough for a Transformer to model directly, which is the operative point from §7, not the exact ratio.
11. Key takeaways¶
| # | Takeaway |
|---|---|
| 1 | ELBO = reconstruction − KL. The two terms fight; that fight is the entire model. |
| 2 | Reparameterization (\(z = \mu + \sigma\epsilon\)) is what makes the gradient flow. |
| 3 | Blurriness comes from the Gaussian likelihood: MSE's optimum is the mean of plausible outputs. |
| 4 | Fix blur with perceptual + adversarial decoder losses (what latent-diffusion VAEs do). |
| 5 | Posterior collapse = KL → 0 = dead latent. Fix with annealing or free bits. |
| 6 | The prior/aggregate-posterior gap causes bad samples; fix it by learning the latent distribution. |
| 7 | VQ-VAE gives discrete tokens → unlocks Transformers over images, audio and video. ~170× compression. |
| 8 | Disentanglement from \(\beta\) alone is not reliably achievable without inductive bias or supervision. |
Further reading¶
- Kingma & Welling, Auto-Encoding Variational Bayes (2013) — the original.
- Kingma & Welling, An Introduction to Variational Autoencoders (2019) — the better tutorial.
- van den Oord et al., Neural Discrete Representation Learning (VQ-VAE, 2017).
- Esser et al., Taming Transformers for High-Resolution Image Synthesis (VQGAN, 2020).
- Higgins et al., β-VAE (2017); Locatello et al., Challenging Common Assumptions… (2019) — read both.
- Alemi et al., Fixing a Broken ELBO (2018) — the rate–distortion view.