The ELBO bounds the log-evidence, and the gap is exactly the divergence between the approximate and the true posterior. Its second term prices the code, and a code priced too high is ignored.
Learn a distribution over codes, not a code, so that every code can be decoded and drawn. The price is a bound in place of the likelihood, and a term that can make the code worthless.
How this chapter is built
M3Load-bearing
The content is mathematics. Understanding is demonstrated by computation, not recall.
basics11/11what the words mean
concept2/2what to picture
theory4/4why it works, and when it does not
mathematics15/15derive it, then compute it
practice9/9build it, break it, read the papers
Five strands, not one. Mathematics is the spine; the other four are the body. A chapter cannot pay its way out of teaching with problems, nor out of problems with teaching.
Chapter I.12 ended with codes that are points: nothing says which other codes are valid, so nothing can be drawn and decoded. A latent-variable model puts a distribution on the code and asks the decoder to explain the data from it. The price is an evidence that is an integral nobody can compute. We bound it, differentiate the bound through a random draw, and measure on the same four points as I.12 what the bound buys and when the code is thrown away.
Fig. 13 — A variational autoencoder in training. The encoder returns a mean and a width, the code is drawn as μ + σε, and the loss adds the decoder's reconstruction error to the KL divergence from the encoder's distribution to the prior that generation will draw from.
What this chapter covers
Latent-variable models: prior, decoder distribution, evidence and posterior
Probabilistic PCA, the linear model where every quantity has a closed form
Why the evidence cannot be estimated by sampling the prior
The evidence lower bound, derived twice, and the exact size of its gap
The Gaussian KL in closed form, and its derivatives
Amortised inference: one encoder in place of one optimisation per example
The reparameterisation trick and the score-function estimator
Variance of gradient estimators, and where each one fails
β, rate and distortion, solved exactly for a linear model
Posterior collapse, and the holes that a prior does not fill
Apparatus
The mathematics this chapter leans on, held in Book 0 so it can be assumed here without being taught here. Not a gate — follow a link when a step stops making sense.
λA regularisation coefficient, or an eigenvalue where the context is linear algebra
dModel width
ℒThe loss
Definitions
Definition 1Latent variable
An unobserved random variable z, given a prior p(z) chosen before training, here the standard normal N(0, I). The model explains each observation as generated from some value of z.
Definition 2Evidence
The marginal likelihood of an observation, p_theta(x) = integral of p(z) p_theta(x | z) over z. It is the quantity maximum likelihood would maximise, and for a nonlinear decoder it has no closed form.
Definition 3Posterior
The distribution of the latent variable given an observation, p_theta(z | x) = p(z) p_theta(x | z) / p_theta(x). Computing it requires the evidence, so it is no easier than the evidence.
Definition 4Variational distribution
A distribution q_phi(z | x) from a tractable family, used in place of the posterior. In a VAE it is a Gaussian with diagonal covariance whose mean and variances the encoder computes from x.
Definition 5Amortised inference
Computing every observation's variational parameters with one shared network, the encoder, in place of a separate optimisation for each observation. The bound lost to the shared network is called the amortisation gap.
Definition 6Evidence lower bound (ELBO)
The quantity E_q[log p_theta(x | z)] - KL(q_phi(z | x) || p(z)). It is at most log p_theta(x) for every q, and equal to it exactly when q is the posterior.
Definition 7Reparameterisation
Writing a draw from q as a differentiable function of its parameters and of noise that does not depend on them, z = mu + sigma elementwise epsilon with epsilon from N(0, I). Gradients then pass through the draw.
Definition 8Score-function estimator
The unbiased gradient estimate f(z) times the gradient of log q_phi(z), from one draw of z. It needs neither a differentiable f nor a continuous z, and its variance is often far larger. Also called REINFORCE.
Definition 9Rate and distortion
The rate is the average KL term, the cost in nats of describing the code relative to the prior. The distortion is the average reconstruction term. The ELBO is minus their sum, and a weight beta on the rate moves along their trade-off.
Definition 10Posterior collapse
The state in which the encoder's distribution equals the prior for every input. The KL term is then zero, the code carries no information about x, and the decoder generates without it.
Definition 11Aggregate posterior
The average of the encoder's distributions over the data, q(z) = average of q_phi(z | x). Generation draws codes from the prior; where q(z) and the prior disagree, the decoder meets codes it was not trained on.
Definition 12Probabilistic PCA
The latent-variable model z from N(0, I), x given z from N(W z, sigma^2 I). Its evidence is Gaussian with covariance W W^T + sigma^2 I, and its maximum-likelihood W spans the principal subspace of Chapter I.12.
Formal results
I.13.T1
The ELBO falls short of the log-evidence by exactly a KL divergence
For any distribution q(z | x) that is positive exactly where the posterior is, log p_theta(x) = ELBO(q) + KL(q(z | x) || p_theta(z | x)). Hence ELBO(q) <= log p_theta(x), with equality exactly when q is the posterior. Maximising the ELBO over q, with the model fixed, is the same as minimising the KL divergence from q to the posterior.
What this does not promise
The gap is not computable on its own, so a larger ELBO may mean a better model or only a tighter bound. Maximising the ELBO over theta also pulls the model towards posteriors that q can represent. A tight bound says nothing about whether the model is good.
I.13.T2
Two unbiased gradients of an expectation
For z = mu + sigma epsilon with epsilon from N(0, 1) and differentiable f, the derivative of E_q[f(z)] with respect to mu is E[f'(z)] and with respect to sigma is E[f'(z) epsilon]. For any integrable f, the score-function estimate f(z) times the derivative of log q(z) is also unbiased. For f(z) = z^2 at mu = 1, sigma = 0.5, the two one-draw estimates of the derivative with respect to mu have variances 1 and 21.75.
What this does not promise
Neither estimator always has the smaller variance: for f(z) = sin(10 z) under N(0, 1) the pathwise estimate has variance 50 and the score function 0.5. The pathwise form needs a differentiable f and a continuous, reparameterisable q; it does not apply to discrete codes.
I.13.T3
The linear β-VAE switches its code off when β sigma^2 reaches the top eigenvalue
Take one latent dimension, the prior N(0, 1), a decoder x given z from N(w z, sigma^2 I) with sigma^2 fixed, centred data whose covariance has top eigenvalue lambda_1 along u_1, and the objective E_q log p(x | z) - beta KL(q || p) with a Gaussian encoder. At the optimum, w lies along u_1 with ||w||^2 = max(lambda_1 - beta sigma^2, 0). While beta sigma^2 < lambda_1 the rate is (1/2) log(lambda_1 / (beta sigma^2)) and the mean squared distortion is (lambda_2 + ... + lambda_d) + beta sigma^2. Otherwise w = 0 and the encoder returns the prior for every x. At beta = 1 the solution is probabilistic PCA and the bound is tight.
What this does not promise
A nonlinear decoder can ignore its code for a different reason: if it can model x well without z, the KL term is pure cost. Learning sigma^2 moves the threshold. A beta above one is not a guarantee of disentangled or interpretable codes.
I.13.T4
The rate is the information in the code plus a mismatch with the prior
Average the KL term over N training examples, indexed by n. Then the average of KL(q(z | x_n) || p(z)) equals I_q(n; z) + KL(q(z) || p(z)), where I_q is the mutual information between the index and the code under q and q(z) is the aggregate posterior. In particular the average KL is zero only if every q(z | x_n) is the prior, and then the code carries no information.
What this does not promise
The identity does not say which term training will reduce. The information term is at most log N, so for a large rate most of it is spent on the mismatch. A small mismatch term does not mean the decoder produces good samples.
Assumptions
Each is paired with the problem that shows what its removal costs. An assumption nobody tests is a disclaimer.
The encoder's family can represent the posterior. A factorised Gaussian cannot represent a correlated one, and the gap that remains does not close with training.I.13.X03
The code is continuous, and the draw can be written as a differentiable function of noise that does not depend on the parameters.I.13.X08
The reconstruction and KL terms are reduced consistently, both summed over their coordinates, or the KL is reweighted without anyone choosing to.I.13.X07
beta sigma^2 stays below the variance the code is meant to capture, or the optimum discards the code.I.13.B03
The encoder's variance stays away from zero. A point code has unbounded rate.I.13.X04
A gradient estimator is chosen for the function at hand. Reparameterisation lowers the variance for smooth, slowly varying integrands, not for all.I.13.X10
A mean reduction where a sum was meant. Averaging squared error over the 784 pixels of an image while summing the KL over the latent dimensions weights the KL by 784 relative to the ELBO. With sigma^2 = 0.5 every direction whose variance is below 392 is then switched off.
Posterior collapse by design. On the chapter's four points with sigma^2 = 1, beta = 4 sets w = 0: the encoder returns N(0, 1) for every input, the rate is 0 and the distortion is 5, the total variance.
A factorised encoder for a correlated posterior. With correlation 0.9 the best diagonal Gaussian reports variance 0.19 for a coordinate whose posterior variance is 1, and a gap of 0.830366 nats remains however long training runs.
A point code. The KL from N(0.866025, s^2) to the prior is 2.182585, 4.480220 and 6.782756 nats at s = 0.1, 0.01 and 0.001: each tenfold narrowing adds about log 10. A deterministic encoder has unbounded rate.
The score-function estimator where a pathwise one exists. For f(z) = z^2 under N(1, s^2) its variance is 21.75 against 1 at s = 0.5, and 10014 against 0.0004 at s = 0.01.
The pathwise estimator on an oscillating integrand. For f(z) = sin(10 z) under N(0, 1) the true gradient is 10 e^(-50), the pathwise estimate has variance 50 and the score-function estimate 0.5.
Differentiating through a discrete draw. For a Bernoulli code with probability 0.5 and f(z) = (z - 0.3)^2, the true gradient is 0.4; writing z = 1[u < pi] gives a derivative of 0 almost everywhere.
Estimating the evidence by sampling the prior. For the point (1, 2), 10 per cent relative error needs 132 prior samples in one dimension, about 4.5e5 in ten independent dimensions, and about 4.7e15 when the noise variance is 0.01.
Predicting a standard deviation with an unconstrained linear output. An output of -0.3 is not a width. Predict log sigma^2 and exponentiate, or pass the output through softplus as in I.3.X02.
Expecting the prior to fill every hole. On the four points the aggregate posterior is a two-component mixture whose KL to the prior is 0.107208 nats: codes drawn from the prior near zero are rarer under q, and the decoder was trained less there.