The variational autoencoder

A greyscale digit is 784 numbers. A variational autoencoder squeezes it into a short code, often two numbers, few enough to plot on a page, and then rebuilds all 784 from that code.

How do you train a model that samples? You train two networks at once: an encoder that turns each picture into a cloud of codes, and a decoder that turns a code back into a picture. You score the pair with one number, a lower bound on how probable the training pictures are, and you push that number up. !Codes for ten classes of picture after training. Each cloud is the encoder's answer for one class, labelled with the class; the dashed contour is the standard normal prior the distance term pulls toward. The arrows are the pull. The three small pictures are what the decoder returns at codes drawn from the prior. Status: synthetic.

The networks are ordinary; the difficulty is the sampling. An encoder is a network that maps a picture to a distribution over codes, and a decoder is a network that maps a code to a distribution over pictures [S:192:3]. The code is z, the picture is x, and the two networks carry parameters theta and phi. !The training graph. The encoder turns a picture into a mean and a spread, and the noise enters as its own input, so the code is a random draw the gradient can pass around. The decoder rebuilds the picture, and the loss adds the rebuild cost to the distance from the prior. Status: schematic.

A decoder that draws from a cloud is not a deterministic function, and you cannot differentiate through a random draw in the usual way [S:74:353].

The fix is a change of variables. Let the encoder output a mean and a spread sigma. Draw a number from a fixed standard normal, multiply it by the spread, add the mean, and you have a sample from the cloud. The randomness now enters through a separate input, so the gradient flows around it rather than through it [S:74:353].

That manoeuvre is the reparameterization trick, and it is what makes a variational autoencoder trainable.

The objective has two terms. One rewards reconstruction, and the other keeps the code cloud near a fixed prior, a simple distribution such as a standard normal.

$$\text{loss}=-\mathbb{E}\left[\log P(x\mid z)\right]+D_{\mathrm{KL}}\left(q(z\mid x)\Vert P(z)\right)$$

Read the first term as the cost of rebuilding the picture from a sampled code, and the second as the cost of placing that code far from the prior [S:74:350]. The sum is the negative evidence lower bound, so minimising it maximises a lower bound on the probability of the data [S:74:350]. The bound is not the probability itself, and the gap is the price of the approximation [S:74:341].

The second term is a KL divergence, a non-negative measure of how much one distribution differs from another. For a one-dimensional Gaussian code, against a standard normal prior, it equals half of the mean squared, plus the variance, minus one, minus the logarithm of the variance [S:74:351]. The expectation in the first term is estimated with one sampled code per picture [S:74:351].

The bound is tight when the encoder's cloud matches the true posterior over codes, and training pushes the cloud in that direction [S:74:348]. At sampling time the encoder is not used at all. You draw a code from the prior, pass it through the decoder, and add noise [S:74:355]. That is why the prior must stay simple: it is the distribution you sample from when no picture is available to encode.

Worked example. One pixel, one code dimension, and a standard normal prior. The true pixel value is 1. The decoder is Gaussian with variance 0.5, so a unit of squared error costs 1 in the loss. The decoder outputs a mean of 0.8, so the reconstruction cost is (1 - 0.8)^2 = 0.04.

For encoder A the distance term is half of (1 + 1 - 1 - 0), which is 0.5. For encoder B it is half of (0.25 + 1 - 1 - 0), which is 0.125. The table holds the reconstruction term at 0.04 for both rows by assumption, to isolate the distance term. That assumption describes a decoder that ignores the code and outputs 0.8 whatever code it receives, which is the collapsed decoder the Deepen warns about.

A decoder that read the code would draw different codes for the two encoders and charge different reconstruction costs, so the totals 0.54 and 0.165 compare two encoders through a decoder that cannot tell them apart. The regulariser argument survives; the claim that both encoders rebuild equally well does not.

Encoder B wins, and that is the regulariser at work. It pulls the code toward zero even when reconstruction does not improve. Without that pull the encoder could give every picture a private code far from the prior, and sampling from the prior would then produce nothing sensible.

Sampling is one line of arithmetic. With encoder A, a draw of -0.5 from the standard normal gives a code of 1 + 1 times (-0.5) = 0.5. The decoder receives 0.5, not 1. In the table the cost was not recomputed at that code, because the table's decoder ignores it; with a decoder that reads the code, the reconstruction term is charged at 0.5 and the totals move [S:74:353].

The cost is quality. Samples from a plain variational autoencoder are generally low-quality, partly because the decoder noise model is a simple Gaussian [S:74:356]. A sharper decoder needs a richer likelihood, and the bound then becomes harder to compute. Diffusion models, in the next section, can be read as a stack of such autoencoders with the encoder fixed [S:74:384].

The formal statement. Write $q_{\varphi }(z\mid x)$ for the encoder, $P_{\theta }(x\mid z)$ for the decoder, and $P(z)$ for the prior. The evidence lower bound is

$$\log P_{\theta }(x)\ge {\mathbb{E}}_q\left[\log P_{\theta }(x\mid z)\right]-D_{\mathrm{KL}}\left(q_{\varphi }(z\mid x)\Vert P(z)\right)$$

The line says that the log-probability of a picture is at least its average reconstruction score minus the distance from the code cloud to the prior. Five hypotheses hold. The encoder has a density wherever the posterior does. The expectation exists and is finite. The KL is finite when $q_{\varphi }(z\mid x)$ puts no mass where the prior puts none, that is, when the encoder is absolutely continuous with respect to $P(z)$. The converse is not required: regions where the prior has mass and the encoder has none contribute nothing to the divergence, which is why this KL is zero-forcing and lets the encoder under-estimate the prior's support [S:62:259]. The decoder likelihood comes from a chosen family, Gaussian in the core's example, with variance 0.5. The expectation is estimated by sampling, with one draw per picture in the usual implementation [S:74:351].

The derivation. Start from the identity $\log P(x)=\log {\mathbb{E}}_q[P(x,z)/q(z)]$. The logarithm is concave, so Jensen's inequality turns the log of an expectation into an expectation of logs. Expanding the right side gives the expected log-likelihood plus the expected log of the prior over the encoder, and the second term is minus the KL. The gap between the two sides equals the KL from the encoder to the true posterior, so the bound is tight exactly when the encoder reproduces that posterior [S:74:348]. This account omits the proof of Jensen's inequality, which follows in one line from the definition of concavity.

The edge case. The core's worked example has two encoders and concludes that the regulariser helps. Push the same logic. If the encoder ignores the picture and returns the prior for every input, the KL term is zero and cannot fall further. The decoder then receives noise and must model the data alone. Training can settle here, and the phenomenon has a name: posterior collapse, in which the encoder always predicts the prior [S:74:360]. The term that improves the core's example becomes a failure when it wins outright.

What the core hid. The core gave the KL for a one-dimensional code. In several dimensions the same expression sums over dimensions: half the trace of the covariance, plus the squared length of the mean, minus the dimension count, minus the log determinant [S:74:351]. The core also said the bound is tight when the cloud matches the posterior. The Gaussian family may not contain the posterior at all, and then the gap has a floor that training cannot remove. The bound measures the model and the approximation together.

Where the analogy breaks. The core says the encoder squeezes a picture into a short code. The encoder returns a mean and a spread, and the code is a draw from that pair. Two runs on one picture give two codes, and the decoder sees a different input each time. That is why the loss is charged at the sampled code rather than at the mean [S:74:353].

Original sources. Kingma and Welling introduce the recognition model as a probabilistic encoder and the likelihood as a probabilistic decoder [S:192:3]. The same paper builds the bound, the stochastic estimator and the reparameterisation that makes the estimator differentiable [S:192:1]. They note that the KL term integrates analytically and acts as a regulariser on the encoder [S:192:4].

An open question. The field has not settled how much of the gap between the bound and the true likelihood comes from the model and how much from the Gaussian approximation. The gap is a KL to the posterior [S:74:348], and that posterior is the object nobody can compute. A bound that is loose by an unknown amount cannot tell you whether a better decoder or a better encoder would help. Evidence that would settle it for a model family: a case with a tractable posterior, where the exact log-likelihood and the bound are both evaluated, and the gap is measured as the encoder family grows.

A frontier result. Diffusion models reuse the bound with the encoder fixed. Prince describes them as a special case of a hierarchical VAE whose encoder is fixed and whose latent space is the same size as the data [S:74:384]. The price is a latent space as large as the picture, with one noisy copy per step [S:74:363]. Hierarchical priors and careful architecture also rescue the plain VAE, which by itself produces low-quality samples [S:74:356]. The bound therefore survived the change of model and became a training recipe rather than an evaluation.

What to read next. Read Kingma and Welling sections 2.2 to 2.4 for the bound and the trick [S:192:3-4]. Read Prince 17.4 to 17.7 for tightness and the encoder [S:74:348-353]. Then read Prince's note on posterior collapse [S:74:360] and Murphy's account of why a high-capacity decoder makes the code disposable [S:62:840].