Explain the roles of the encoder, latent space, reparameterization trick, and decoder in a VAE. Derive and explain the two major components of its loss function and discuss why both reconstruction loss and KL-divergence are required.
What's actually being asked
Explain each VAE component's role, derive the ELBO/loss function's two terms from first principles, and justify why both terms are necessary.
Answer
Roles of each component (see the diagram below for the flow)
Encoder: takes input x and outputs the parameters μ(x), σ(x) of an approximate posterior q(z|x) = N(μ, σ²) over the latent space — it doesn't output a single point, it outputs a distribution.
Latent space: the compressed probabilistic representation z; because the KL term during training pulls every input's q(z|x) toward the shared prior p(z) = N(0,I), the space ends up continuous and structured, so nearby points decode to similar, meaningful outputs.
Reparameterization trick: sampling z ~ q(z|x) is normally a non-differentiable operation (you can't backpropagate through a random draw). The trick rewrites the sample as z = μ + σ ⊙ ε, with ε ~ N(0, I) drawn independently — now the randomness is isolated in ε, and z is a differentiable function of μ and σ, so gradients can flow back through the encoder during training.
Decoder: takes the sampled z and outputs (the parameters of) p(x|z), reconstructing a full data sample x̂ from the compressed, noisy latent code.
Deriving the two loss components from the ELBO
Start from the (generally intractable) log-likelihood of the data:
log p(x) = log ∫ p(x,z) dz
Multiply and divide inside the integral by the encoder's distribution q(z|x) — this doesn't change the value, just rewrites it as an expectation under q:
log p(x) = log 𝔼_{q(z|x)}[ p(x,z) / q(z|x) ]
Apply Jensen's inequality (log is concave, so log 𝔼[·] ≥ 𝔼[log ·]):
log p(x) ≥ 𝔼_{q(z|x)}[ log( p(x,z) / q(z|x) ) ] — this right-hand side is the ELBO (evidence lower bound).
Expand p(x,z) = p(x|z)·p(z) and split the log of a ratio into a difference:
ELBO = 𝔼_{q(z|x)}[ log p(x|z) + log p(z) − log q(z|x) ] = 𝔼_{q(z|x)}[ log p(x|z) ] − 𝔼_{q(z|x)}[ log q(z|x) − log p(z) ]
The second term is, by definition, the KL-divergence between q(z|x) and the prior p(z):
𝔼_{q(z|x)}[ log q(z|x) − log p(z) ] = D_KL(q(z|x) ‖ p(z))
So:
log p(x) ≥ 𝔼_{q(z|x)}[log p(x|z)] − D_KL(q(z|x) ‖ p(z)) — the canonical ELBO.
Training maximizes this lower bound, which is equivalent to minimizing its negative — the VAE loss:
L(θ,φ;x) = −𝔼_{q_φ(z|x)}[log p_θ(x|z)] + D_KL(q_φ(z|x) ‖ p(z))
This has exactly two components:
Reconstruction loss−𝔼_{q(z|x)}[log p(x|z)] — how well the decoder reconstructs x from the sampled z (in practice, often MSE for continuous data or cross-entropy for binary/categorical data, both of which fall out of specific choices of p(x|z)).
KL-divergence lossD_KL(q(z|x) ‖ p(z)) — how far the encoder's posterior is from the standard normal prior.
Why both are required
Reconstruction loss alone (as in a plain autoencoder) gives no pressure on the shape of the latent space — the encoder could map each input to an arbitrary, far-apart point with near-zero variance, and the model would still reconstruct perfectly, but sampling a random z ~ N(0,I) afterward would land in a region nothing was ever trained to decode — generation would fail. KL-divergence alone, with no reconstruction term, would just push every q(z|x) to collapse onto the prior regardless of x ("posterior collapse") — the decoder would learn to ignore z and output some generic average sample, since there's no penalty for ignoring the input. Only together do they balance: reconstruction keeps the latent code informative about x, and KL keeps the latent space well-structured and sample-able — which is exactly what makes the VAE usable as a generator, not just a compressor.
Exam tip
If asked to "derive" the loss, examiners want the Jensen's-inequality step shown explicitly — stating the ELBO formula without deriving it from log p(x) is a common way to lose derivation marks even when the final formula quoted is correct.