VAE: learn to reconstruct and to sample

An interactive field guide

A VAE replaces a single encoder output with a distribution. This creates a trainable connection between reconstruction and sampling from a prior.

Before this article: Autoencoders · Probability and tensors

What was already possible

An ordinary AE learns where each input should go, but it does not give a simple rule for drawing new codes. A latent-variable model instead defines p(z) and pθ(x|z). Its image density integrates over all possible z. For a nonlinear decoder, this integral and the true posterior pθ(z|x) are usually difficult to compute.

What changes, and why

The encoder qφ(z|x) approximates the true posterior. It is not the prior and need not equal the true posterior. AEVB jointly learns this inference network and the decoder using a variational lower bound, or ELBO. Reparameterization gives a practical gradient estimator. These are the central contributions, rather than merely adding noise to an AE.

Read from left to right. Select a module to connect its role, tensor shape, and calculation. Snippets are teaching pseudocode. Sources provide the paper method or audited implementation.

1. Posterior networkx → μ, log σ²: [B,D]2. Reparameterizeε,z: [B,D]3. Likelihood and KLscalar objective

Posterior network

Two outputs parameterize qφ. They are not two separate latent samples.

mu, logvar = encoder(x).chunk(2, dim=-1)

Reparameterize

Noise is independent of encoder parameters. Gradients flow through the affine transform.

std = exp(0.5 * logvar)
z = mu + std * randn_like(std)

Likelihood and KL

Sum latent dimensions per example; define pixel reduction consistently.

x_params = decoder(z)
kl = 0.5 * (mu**2 + exp(logvar) - logvar - 1).sum(-1)
loss = reconstruction_nll(x, x_params) + kl

Inspect the pinned source

$$\mathcal{L}=\mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)]-D_{KL}(q_\phi(z|x)\Vert p(z))$$

Work through the mechanism

For a diagonal Gaussian, the network emits μ and log σ². Compute σ=exp(0.5 log σ²), draw ε from N(0,I), and form z=μ+σε. Randomness now enters through ε; the transformation from μ and σ remains differentiable. With μ=1, σ=0.5, and ε=0.5, z=1.25.

The demo shows a one-dimensional posterior and a standard-normal prior. KL is zero at μ=0 and σ=1. Narrowing σ does not always help: an extremely narrow posterior pays a large KL cost even if its mean is zero.

Change a parameter, inspect the computed values, then compare this small example with the model described above.

Analytical teaching model · controls change the example, not a trained research model.

Mean μ

0

Standard deviation σ

1

Reset

Green: posterior q. Grey: prior p.

Computed values for the current example

Quantity	Value
KL(q || p)	0.000
z = μ + σ × ε, ε = 0.5	0.500
log variance	0.000

Training is not sampling

Maximize expected log pθ(x|z) minus KL(qφ(z|x)||p(z)), or minimize its negative. Reconstruction evaluates samples from q; new generation draws z from p and does not use the encoder. Using μ alone for deterministic reconstruction is an evaluation choice, not the complete training procedure.

What this does not establish

KL regularization does not guarantee that every latent region produces a good image. A strong decoder may ignore z, called posterior collapse. Pixel likelihood, latent dimension, and KL weight all matter. A β other than one changes the standard ELBO objective.

Open the derivation

Insert q(z|x) into the latent integral. Jensen’s inequality moves log inside the expectation and produces a lower bound. This assumes q covers the relevant support.

$$\log p_\theta(x)=\log\mathbb{E}_{q_\phi}\!\left[\frac{p_\theta(x,z)}{q_\phi(z|x)}\right]\geq\mathbb{E}_{q_\phi}[\log p_\theta(x,z)-\log q_\phi(z|x)]$$

Split p(x,z)=p(x|z)p(z). The result is expected reconstruction log-likelihood minus KL to the prior. The gap to log p(x) is instead KL to the true posterior. These are two different KL terms.

$$\log p_\theta(x)-\mathcal{L}=D_{KL}(q_\phi(z|x)\Vert p_\theta(z|x))\geq0$$

For a diagonal Gaussian q and standard-normal p, integrate KL analytically. With one coordinate μ=1 and σ=0.5, KL≈0.818. The decoder expectation still usually needs Monte Carlo samples.

$$D_{KL}=\frac12\sum_j\left(\mu_j^2+\sigma_j^2-1-\log\sigma_j^2\right)$$

Check your understanding

Where does the encoder appear when generating a new image from the prior?

It does not appear. Draw z from p(z), then use pθ(x|z). The encoder is needed for inference from an observed x.

Sources and next steps

• Method / pinned implementation source
• VQ-VAE
• DiT
• DiTo

Return to the learning map

Diagrams and snippets are teaching reconstructions. Paper reports, implementation observations, and analytical examples are distinguished. No research model runs in this page.