VAE architecture, ELBO derivation, and training losses

What's actually being asked

Full VAE architecture and working, a from-scratch ELBO derivation, how reconstruction loss and KL loss each pull on training, and a diagram of the whole pipeline.

Answer

Architecture

• Encoder (recognition network, params φ): maps input x to the parameters of an approximate posterior, outputting μ(x), σ(x), defining q(z|x) = N(μ, σ²).
• Latent sampling: reparameterization z = μ + σ⊙ε, ε~N(0,I) (see diagram) — keeps the path differentiable while still sampling.
• Decoder (generative network, params θ): maps z back to a distribution over x, p(x|z), typically producing x̂ as the mean of a Gaussian/Bernoulli used for reconstruction.
• Prior: p(z) = N(0,I) assumed over latent space, so generation at inference just needs a fresh z ~ N(0,I) fed through the decoder, no encoder required.

Working principle

At training time: x → Encoder → (μ,σ) → reparameterize → z → Decoder → x̂ (exactly the diagram above). Loss combines how well x̂ reconstructs x and how close q(z|x) stays to the prior p(z); both networks are updated jointly in one backward pass. At generation time: discard the encoder entirely, sample z ~ N(0,I) directly, run only the decoder to get new samples.

Deriving the ELBO

We want to maximize the marginal likelihood p(x) = ∫ p(x,z) dz, but this integral is intractable. Introduce the approximate posterior q(z|x) by multiplying and dividing inside the integral:

log p(x) = log ∫ p(x,z) dz = log ∫ q(z|x)·[p(x,z)/q(z|x)] dz = log 𝔼_{q(z|x)}[p(x,z)/q(z|x)]

Apply Jensen's inequality (log is concave, so log of an expectation ≥ expectation of the log):

log 𝔼_{q(z|x)}[p(x,z)/q(z|x)] ≥ 𝔼_{q(z|x)}[log(p(x,z)/q(z|x))]

Factor the joint p(x,z) = p(x|z)p(z) and split the log:

𝔼_{q(z|x)}[log p(x,z) − log q(z|x)] = 𝔼_{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 p(z) − log q(z|x)] = 𝔼_{q(z|x)}[log p(x|z)] − D_KL(q(z|x) ‖ p(z))

This gives the Evidence Lower BOund:

log p(x) ≥ 𝔼_{q(z|x)}[log p(x|z)] − D_KL(q(z|x) ‖ p(z)) = ELBO

Training maximizes the ELBO (a tractable proxy for the intractable log p(x)), equivalently minimizing the negative ELBO:

L(θ,φ;x) = −𝔼_{q_φ(z|x)}[log p_θ(x|z)] + D_KL(q_φ(z|x) ‖ p(z))

The two loss terms

• Reconstruction loss −𝔼_{q(z|x)}[log p(x|z)]: how well the decoder reconstructs x from the sampled z. In practice approximated with one Monte-Carlo sample of z (thanks to reparameterization), implemented as MSE (Gaussian decoder) or binary cross-entropy (Bernoulli decoder) between x and x̂.
• KL loss D_KL(q(z|x) ‖ p(z)): regularizes the encoder's output distribution to stay close to the prior N(0,I), keeping the latent space smooth and continuous so it's usable for generation. With Gaussian q(z|x)=N(μ,σ²) against the standard normal prior, this has a closed form per latent dimension j, summed over all dims:

D_KL = −½ Σ_j (1 + log σ_j² − μ_j² − σ_j²)

Total loss = reconstruction loss + KL loss. The two pull in opposite directions — reconstruction wants z to carry as much x-specific information as possible (pushing σ→0), while KL wants q(z|x) to collapse toward the prior everywhere. That tension is exactly what keeps the latent space both informative and smooth/generation-friendly.

Worked example

Illustrative 1-D latent: μ=0.3, σ=0.8. D_KL = −½(1 + log(0.8²) − 0.3² − 0.8²) = −½(1 + log(0.64) − 0.09 − 0.64) = −½(1 − 0.446 − 0.09 − 0.64) = −½(−0.176) = 0.088. If reconstruction loss (BCE) for this sample is, say, 45.2, total loss ≈ 45.288 — reconstruction dominates per-sample, but KL matters summed across the whole batch and all latent dims.

Exam tip

Always show the Jensen's inequality step explicitly — many students jump straight to "ELBO = reconstruction − KL" without deriving it, and that's exactly where marks are lost. Also state clearly that KL is a regularizer, not a measure of reconstruction quality.