Interactive Explainer

VAE Latent-Space Explorer

The β knob in a VAE trades reconstruction for latent structure. At β=0 you have a vanilla autoencoder — great reconstruction, chaotic latent space. At β=10 you have a strong VAE — clean unit-Gaussian latent, but samples get blurry. Drag the slider and see it.

A VAE's loss has two terms: reconstruction (how well can I decode this back?) and a KL divergence pulling the latent distribution q(z|x) toward the standard normal prior. The β-VAE (Higgins 2017) simply weights the KL term: L = recon + β · KL. Slide β to see what it does.

The playground

What's happening: at β=0 the latent points cluster in a few tight clumps — one per class. At β=1 the clumps become blobs of width 1, roughly tiling a unit circle. At β=10 everything collapses to the origin (posterior collapse) — the model stops using the latent.

The reparameterization trick

z = μ(x) + σ(x) ⊙ ε    ε ~ N(0, I)

The encoder outputs μ and σ (per input). We sample ε from a fixed noise distribution and compute z deterministically. The noise is in ε, which has no parameters, so gradients flow through μ and σ fine. This is the trick that makes VAE training work.

The ELBO

log p(x) ≥ E_{q(z|x)}[log p(x|z)]  −  KL(q(z|x) ∥ p(z))

The ELBO is a lower bound on the log-marginal likelihood. We can't compute log p(x) directly (intractable marginal), but we CAN compute and maximize its lower bound. That's what VAE training does.

ELBO, derived step by step

We want $\log p(x)$ but it's intractable. Introduce an approximate posterior $q_\phi(z\mid x)$ (the encoder) and use the fact that for any $q$:

log p(x) = E_{q_φ(z|x)} [ log p(x, z) − log q_φ(z|x) ] + KL( q_φ(z|x) ∥ p(z|x) )

The KL term on the right is $\ge 0$, so dropping it gives a lower bound — the ELBO. Rearranging:

ELBO(x) = E_{q_φ(z|x)} [ log p_θ(x|z) ] − KL( q_φ(z|x) ∥ p(z) )

Two terms with intuitive readings: the first wants $z$ to reconstruct $x$ (the decoder loss); the second wants $z$'s posterior to match the prior. Maximising ELBO trades off both.

The reparameterisation trick · why VAEs train

Sampling $z \sim \mathcal{N}(\mu_\phi(x), \sigma_\phi^2(x))$ doesn't have a gradient with respect to $\mu_\phi, \sigma_\phi$. Reparameterise:

z = μ_φ(x) + σ_φ(x) ⊙ ε,    ε ~ N(0, I)

Now the sampling step is just an elementwise multiply and add; $\mu_\phi, \sigma_\phi$ receive gradients. Kingma & Welling 2014 — the trick that made VAEs trainable end-to-end.

Posterior collapse · and three fixes

A common failure: the encoder outputs $q_\phi(z\mid x) \approx p(z)$ for every $x$, and the decoder learns to ignore $z$. The KL term is zero; reconstruction is mediocre but stable; the latent is useless. Symptoms: low KL, blurry samples, samples that don't depend on $z$.

• β-VAE. Weight the KL term by $\beta < 1$ (or $\beta > 1$ for disentanglement, Higgins et al. 2017). KL warm-up is the simpler practical fix.
• Free bits. Apply KL term only when above a threshold per latent dim (Kingma et al., 2016). Stops dims that could be informative from collapsing.
• Stronger decoder family. Use PixelCNN / autoregressive decoder so the latent has somewhere useful to live. Or stronger prior (VQ-VAE 2017, Vamp-prior).

VAE family in 2026

• VQ-VAE. Discrete latent via vector-quantised codebook. Standard image tokeniser; powers most discrete-token image / audio models.
• NVAE / VDM. Hierarchical VAEs that match diffusion quality on smaller scales.
• Latent-Diffusion VAE. The tiny VAE inside Stable Diffusion (KL-f8): pixels → 4-channel latent at $1/8$ resolution. Trained once, frozen.
• Conditional VAEs. $q_\phi(z\mid x, c)$ for class-conditional generation; or for tabular / time-series with structured context.

Companion: diffusion · Bayesian posterior · info theory (KL refresher).