71 - Variational AutoEncoder (1): From Likelihood to the ELBO, Fully Derived

At the end of the lesson on autoencoders we had a precise problem: the decoder knows how to turn codes into images, but we do not know which distribution to draw the codes from. The Variational AutoEncoder (VAE), proposed by Kingma and Welling in 2013, solves the problem by changing perspective: it stops being a machine that compresses and becomes a probabilistic model of the images, with a latent variable z whose prior distribution is fixed (a standard normal) and of which one learns how each image depends on it. This lesson is the theory, and I do it in full: no formula is declared without being derived, with the single exception of Jensen's inequality (for which I give a definition and a numerical example). Every numerical step is verified with a program, which I ran twice with identical outputs. In the next lesson we will build the VAE in PyTorch; if you understand this page, the code will be very short.

The words of this lesson

This is the lesson with the most vocabulary of the whole module, so the definitions should be read one at a time. A generative model describes how the data are "produced": here first a code z is drawn from a simple distribution, then the image x is generated from z. The latent variable z is not in the data but we postulate that it exists and "explains" them, like the ten words of the face description. The prior p(z) is the distribution that we choose for z before seeing any data: the standard normal N(0, I), mean zero, variance one, independent components.

The likelihood p(x | z) is the probability of an image given the code, and it is what the decoder computes. The evidence p(x) = ∫ p(x | z) p(z) dz is the probability of the image in absolute terms, summed over all possible codes, and it is the quantity we would like to maximise on the real data. The posterior p(z | x) is the distribution of the codes that could have produced x; by Bayes' rule it equals p(x | z) p(z) / p(x), and it is the "inverse" of the decoder.

The KL divergence KL(q ‖ p) measures how much a distribution q differs from p: it is E_q[log q(z) − log p(z)], the mean according to q of the difference of the logarithms; it is always ≥ 0, equals 0 only if q = p and is not symmetric. Jensen's inequality says that for a concave function such as the logarithm log(E[Y]) ≥ E[log Y]: if Y is 1 or 9 with probability ½, log(5) = 1.609 while ½·log(1) + ½·log(9) = 1.099. Variational inference approximates an intractable exact distribution (here the posterior) with a simpler q, chosen from a family and optimised. The ELBO (Evidence Lower BOund) is a computable lower bound of log p(x): maximising it pushes the likelihood up. Amortised inference, finally, uses a network (the encoder) that given x immediately produces the parameters of q, instead of optimising a separate q for every image.

The analogy: a pastry chef and the apprentice

A pastry chef (the model) prepares cakes by choosing a "recipe slip" z (not very sweet, very tall, with fruit…), drawn from a jar that contains the recipes they usually use: the prior. Given the slip, the cake x that comes out is never identical: there is a bit of variability (the likelihood p(x|z)). Now you are shown a cake and asked: "which slip produced it?". It is the inverse problem, the posterior, and for a pastry chef with a thousand ingredients it cannot be computed exactly: you would have to try all the recipes. So an apprentice (the encoder) is hired who, looking at the cake, answers: "in my opinion the slip is around this value, with this uncertainty": a central value μ and a width σ. Then there is a further rule: the apprentice is asked never to propose slips that the pastry chef would never use (far from the jar). That rule is the KL term. The result: the pastry chef can now be used on their own, drawing a slip from the jar, and the cakes will be sensible.

The problem: maximising p(x) is intractable

We want the model to assign high probability to real images, that is, to maximise log p(x). By definition:

p(x) = ∫ p(x | z) p(z) dz

The decoder p(x | z) is a neural network. The integral, over all possible z in a space of dozens of dimensions, of a function defined by a network has no closed form, and a direct numerical computation (sampling many z from the prior and averaging) has an enormous variance: almost all random z produce nothing resembling x, and an astronomical number of samples would be needed. So log p(x) can neither be computed nor differentiated directly.

The idea: a family q and a lower bound

We introduce a distribution q(z | x), which we call the variational approximation of the posterior, and which is exactly what the encoder produces. In the VAE a Gaussian with a diagonal covariance matrix is chosen: q(z | x) = N(μ(x), diag(σ²(x))), where μ(x) and σ²(x) are two vectors computed by the encoder. Now we derive a lower bound of log p(x), in two ways.

Derivation 1: with Jensen's inequality

We start from log p(x) and multiply and divide by q inside the integral (nothing changes):

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

The last step recognises the integral as an expected value with respect to q. Now we apply Jensen (log E[Y] ≥ E[log Y]) with Y = p(x, z) / q(z|x):

log p(x) ≥ E_{z~q} [ log p(x, z) − log q(z|x) ]  =: ELBO

Since p(x, z) = p(x | z) p(z), it can be split into two pieces:

ELBO = E_q[ log p(x|z) ]  +  E_q[ log p(z) − log q(z|x) ]
     = E_q[ log p(x|z) ]  −  KL( q(z|x) ‖ p(z) )

The last equality is the definition of KL (with the sign changed, because E_q[log p − log q] = −KL(q‖p)).

Derivation 2: with the exact decomposition

The second derivation is more instructive because it says how much the gap between the ELBO and log p(x) is. We start from log p(x), which does not depend on z, so it is equal to its mean with respect to q. Then we use Bayes' rule p(x) = p(x, z) / p(z | x):

log p(x) = E_q[ log p(x) ]
         = E_q[ log ( p(x, z) / p(z|x) ) ]
         = E_q[ log ( p(x, z) / q(z|x) ) ]  +  E_q[ log ( q(z|x) / p(z|x) ) ]
         =           ELBO                  +  KL( q(z|x) ‖ p(z|x) )

In the third step we multiplied and divided by q(z|x) inside the logarithm, splitting the logarithm of a ratio into two. The second piece is exactly the definition of KL between q and the true posterior. Since every KL is ≥ 0:

log p(x) = ELBO + KL( q(z|x) ‖ p(z|x) )   ≥   ELBO

Three important consequences. (1) The ELBO really is a lower bound, and by two different roads we arrive at the same result. (2) The gap is exactly the KL between q and the true posterior: the better q approximates p(z|x), the closer the ELBO is to log p(x); if q = p(z|x) the gap is zero. (3) Maximising the ELBO with respect to the encoder parameters reduces that gap, while maximising it with respect to the decoder raises log p(x). With a single objective both things are done together.

The two terms of the ELBO and their meaning

ELBO(x) =   E_{z~q(z|x)} [ log p(x | z) ]     −   KL( q(z|x) ‖ p(z) )
              reconstruction term                  regularisation term

The reconstruction term measures, for each z drawn by the encoder, how probable the decoder finds the real image. If the decoder is a Gaussian with fixed variance s², log p(x | z) = −‖x − x̂‖² / (2s²) + constant: it is the usual squared error; if the pixels are considered binary (Bernoulli) it is the negative of the binary cross-entropy. The regularisation term, the KL to the prior, measures how far the distribution of the codes of x departs from the standard normal: the more "specific" the encoder is about x (mean far from 0, small standard deviation), the higher the KL, and it is the cost, in nats of information, of the code.

This balance is the soul of the VAE. An ordinary autoencoder has only the first term, and therefore uses the latent space as it pleases. Here the second term forces the codes to stay in a zone close to the standard normal, and this makes the latent space samplable.

The KL between a Gaussian and the standard normal, derived by hand

In the VAE we need KL( N(μ, σ²) ‖ N(0, 1) ) component by component. It is derived in a few lines, and it is worth doing: it is the formula you will use in the code. For one component, with q = N(μ, σ²) and p = N(0, 1):

log q(z) = −½ log(2π) − ½ log σ² − (z − μ)² / (2σ²)
log p(z) = −½ log(2π) − z² / 2

KL = E_q[ log q(z) − log p(z) ]
   = E_q[ −½ log σ² − (z − μ)² / (2σ²) + z² / 2 ]

We need three means according to q. By the definition of variance, E_q[(z − μ)²] = σ², so E_q[(z − μ)²/(2σ²)] = ½. By another basic property, E_q[z²] = μ² + σ² (mean squared plus variance). Substituting:

KL = −½ log σ² − ½ + ½ (μ² + σ²)
   = ½ ( μ² + σ² − log σ² − 1 )

For a latent vector with k independent components, the KLs add up. We verify the formula by comparing it with a sample estimate (Monte Carlo): we draw a million z from q, compute log q(z) − log p(z) and take the mean.

import numpy as np

rng = np.random.default_rng(0)

# 1) KL( N(mu, sigma^2) || N(0, 1) ): closed form versus a Monte Carlo estimate
def kl_closed(mu, log_var):
    return 0.5 * (mu ** 2 + np.exp(log_var) - log_var - 1.0)

def kl_monte_carlo(mu, log_var, n=1_000_000):
    sigma = np.exp(0.5 * log_var)
    z = mu + sigma * rng.standard_normal(n)                              # samples from q
    log_q = -0.5 * np.log(2 * np.pi) - 0.5 * log_var - (z - mu) ** 2 / (2 * sigma ** 2)
    log_p = -0.5 * np.log(2 * np.pi) - 0.5 * z ** 2
    return np.mean(log_q - log_p)                                        # E_q[ log q - log p ]

for mu, lv in [(0.0, 0.0), (1.0, 0.0), (0.5, -1.0), (2.0, -3.0)]:
    print(f"mu={mu:4.1f} log_var={lv:5.1f}   closed {kl_closed(mu, lv):.4f}   Monte Carlo {kl_monte_carlo(mu, lv):.4f}")

mu= 0.0 log_var=  0.0   closed 0.0000   Monte Carlo 0.0000
mu= 1.0 log_var=  0.0   closed 0.5000   Monte Carlo 0.5008
mu= 0.5 log_var= -1.0   closed 0.3089   Monte Carlo 0.3087
mu= 2.0 log_var= -3.0   closed 3.0249   Monte Carlo 3.0243

Explanation. kl_closed is the derived formula (with log σ² as the parameter, exp(log_var) is σ²). kl_monte_carlo draws n samples z = μ + σ·ε, computes the two log-densities written out explicitly and averages their difference. The two methods coincide to the third or fourth digit (the small difference is sampling noise). To read the cases: μ = 0, σ = 1 is the prior itself, KL = 0; moving the mean to 1 costs 0.5 nat (½·1²); reducing the variance to e⁻³ ≈ 0.05 with mean 2 costs 3.02 nat.

What the cost means. The KL measures how much information the code carries about x. An encoder that wants to be very precise (small σ) and use a code far from the origin pays dearly: with μ = 3 and σ = 0.01 the KL is 8.61 nat; with μ = 3 and σ = 1 only 4.5; with μ = 0 and σ = 0.1 it is 1.81. The encoder has to convince the decoder that that precision (those nats) is worth the improvement in reconstruction.

The reparametrisation trick

To train with gradient descent we must differentiate the ELBO with respect to the encoder parameters (μ, σ). But the reconstruction term is an expected value with respect to z ~ q(z|x), that is, with respect to a random sampling whose distribution depends on μ and σ. How do you differentiate a random draw? There are two roads.

Score-function estimator (REINFORCE). One differentiates the density: ∇_μ E_q[f(z)] = E_q[ f(z) · ∇_μ log q(z) ]. It works with any f, even a non-differentiable one, but the estimate has a large variance.

Reparametrisation. One writes z = μ + σ ⊙ ε with ε ~ N(0, I). Now the randomness is all in ε, which does not depend on any parameter, while z is a deterministic and differentiable function of μ, σ and ε. One can therefore differentiate inside the expected value: ∇_μ E[f(μ + σ ε)] = E[f'(z)]. We compare the two estimators on a case where we know the exact answer: f(z) = z² with z ~ N(1.5, 0.7²). Here E[z²] = μ² + σ², so the derivative with respect to μ is exactly 2μ = 3.

# 2) Why the reparametrisation trick: two ways to estimate  d/d(mu) E_{z~N(mu,s^2)} [ f(z) ]   with f(z) = z^2
mu, sigma, n = 1.5, 0.7, 200_000
true_grad = 2 * mu                                      # E[z^2] = mu^2 + sigma^2  ->  derivative 2*mu

eps = rng.standard_normal(n)
z = mu + sigma * eps

# (a) pathwise / reparametrised: z = mu + sigma*eps, so d f(z)/d mu = f'(z) * 1 = 2 z
g_reparam = 2 * z

# (b) score-function (REINFORCE): d/dmu E[f] = E[ f(z) * d log q(z)/d mu ] = E[ f(z) * (z - mu) / sigma^2 ]
g_score = (z ** 2) * (z - mu) / sigma ** 2

print(f"true gradient            : {true_grad:.3f}")
print(f"reparametrised estimator : mean {g_reparam.mean():.3f}   std of one sample {g_reparam.std():.3f}")
print(f"score-function estimator : mean {g_score.mean():.3f}   std of one sample {g_score.std():.3f}")
print(f"variance ratio score/reparam: {g_score.var() / g_reparam.var():.1f}x")

true gradient            : 3.000
reparametrised estimator : mean 2.997   std of one sample 1.398
score-function estimator : mean 2.982   std of one sample 6.979
variance ratio score/reparam: 24.9x

Explanation. 200,000 samples of ε are drawn. The reparametrised estimator is simply f'(z) = 2z computed on each sample; the score-function one is f(z) · (z − μ)/σ². Both are correct on average (2.997 and 2.982 against 3.000: the difference is noise). But the standard deviation of a single sample is 1.40 for the reparametrised one and 6.98 for the other: the variance is 24.9 times larger. In practice, with a single sample per image (as is done in VAEs), the score-function estimator would give gradients so noisy as to make training unusable; the reparametrised one works. It is the reason why the trick is "the" ingredient of the VAE. Condition: the trick requires that f be differentiable with respect to z (here it is, the decoder is a network), and that z can be written as a differentiable function of the parameters plus an independent noise (true for Gaussians).

An ELBO on a toy case, computed in three ways

To see every piece work together, a case with one observed variable and one latent variable, in which everything can be computed exactly. Imagine the "image" is a single pixel x = 0.8. The prior is z ~ N(0, 1). The decoder is p(x | z) = N(x; z, 0.3²) (the pixel is z plus a bit of noise). The encoder says, for this x, q(z | x) = N(0.5, e⁻¹), that is μ = 0.5 and log σ² = −1.

# 3) The ELBO on a toy problem with ONE number x and a 1-D latent variable, computed three ways
x = 0.8                      # the observed "image" (a single pixel)
mu, log_var = 0.5, -1.0      # what the ENCODER says about z given x:  q(z|x) = N(mu, exp(log_var))
s = 0.3                      # the DECODER's fixed noise:  p(x|z) = N(x; z, s^2)
sigma = np.exp(0.5 * log_var)

# (i) closed form: reconstruction term  E_q[log N(x; z, s^2)]  and KL to the prior
recon_closed = -0.5 * np.log(2 * np.pi * s ** 2) - ((x - mu) ** 2 + sigma ** 2) / (2 * s ** 2)
kl = kl_closed(mu, log_var)
print(f"closed form : reconstruction {recon_closed:.4f}   KL {kl:.4f}   ELBO {recon_closed - kl:.4f}")

# (ii) Monte Carlo with the reparametrisation trick
z = mu + sigma * rng.standard_normal(1_000_000)
recon_mc = np.mean(-0.5 * np.log(2 * np.pi * s ** 2) - (x - z) ** 2 / (2 * s ** 2))
print(f"Monte Carlo : reconstruction {recon_mc:.4f}   KL {kl_monte_carlo(mu, log_var):.4f}   ELBO {recon_mc - kl_monte_carlo(mu, log_var):.4f}")

# (iii) the quantity the ELBO bounds: log p(x) = log integral p(x|z) p(z) dz  (here computable: x ~ N(0, 1 + s^2))
log_px = -0.5 * np.log(2 * np.pi * (1 + s ** 2)) - x ** 2 / (2 * (1 + s ** 2))
print(f"exact log p(x) = {log_px:.4f}   (ELBO must be <= this; the gap is KL(q(z|x) || p(z|x)))")

# the best possible q for this x gives a gap of zero: the exact posterior p(z|x) = N( x/(1+s^2), s^2/(1+s^2) )
mu_star, var_star = x / (1 + s ** 2), s ** 2 / (1 + s ** 2)
recon_star = -0.5 * np.log(2 * np.pi * s ** 2) - ((x - mu_star) ** 2 + var_star) / (2 * s ** 2)
elbo_star = recon_star - kl_closed(mu_star, np.log(var_star))
print(f"ELBO with the exact posterior: {elbo_star:.4f}  (equals log p(x): the bound is tight)")

closed form : reconstruction -2.2587   KL 0.3089   ELBO -2.5677
Monte Carlo : reconstruction -2.2541   KL 0.3088   ELBO -2.5631
exact log p(x) = -1.2556   (ELBO must be <= this; the gap is KL(q(z|x) || p(z|x)))
ELBO with the exact posterior: -1.2556  (equals log p(x): the bound is tight)

Let us follow the three pieces. (i) Closed form. For a Gaussian decoder, E_q[log N(x; z, s²)] = −½ log(2π s²) − ((x − μ)² + σ²) / (2 s²) (derived like the KL, using E_q[(x − z)²] = (x − μ)² + σ²). The code gives −2.2587; with the KL 0.3089 we obtain ELBO = −2.5677. (ii) Monte Carlo with reparametrisation. A million samples z = μ + σ ε, mean of log p(x | z) and of the estimated KL: ELBO ≈ −2.5631, equal to the closed form within the noise. (iii) The true value. In this case x is the sum of two independent Gaussians, so p(x) = N(0, 1 + s²) and log p(x) = −1.2556 is computed exactly: the ELBO (−2.57) is below −1.26, as it must be. The gap, 1.31 nat, is the KL between the encoder's q and the true posterior p(z | x).

The last line shows that the gap can be closed: for this toy model the true posterior is N(x/(1 + s²), s²/(1 + s²)) (mean 0.734, variance 0.0826). If the encoder produced exactly that q, the ELBO would be −1.2556, identical to log p(x): the bound is tight. In the real VAE the encoder gets it a bit wrong (it is a network with limited capacity and a diagonal Gaussian family), so the ELBO is always below log p(x); its error is exactly the gap that training tries to reduce.

The drawing summarises everything. The encoder does not produce a code but two vectors, μ(x) and log σ²(x) (the logarithm of the variance is used because it can be any real number, while the variance must be positive: σ² = exp(log σ²) guarantees it). The yellow node applies the reparametrisation with a noise ε drawn outside the computation graph (dashed line: no gradient must cross it). The decoder produces x̂. At the bottom, the loss to minimise is the negative of the ELBO: reconstruction + KL. At generation time the encoder is thrown away.

The loss to implement

Putting it together: for a batch of images, one minimises (the negative of the ELBO, averaged over the batch):

loss = E_ε [ ‖x − D(μ + σ ⊙ ε)‖² / (2 s²) ]  +  ½ Σ_j ( μ_j² + σ_j² − log σ_j² − 1 )

The expected value over ε is estimated with a single sample per image: the authors of the VAE report that one sample per point is enough if the mini-batch is large enough (for example 100). The ratio between the two terms is controlled by s²: choosing s amounts to weighting reconstruction and regularisation (we will see it, as β, in lesson 73).

Trade-offs

The diagonal Gaussian family is simple and makes the KL computable in closed form, but it limits what q can represent: if the true posterior has several peaks or correlations between components, the gap never closes. A single sample per image is cheap and, according to the authors, sufficient with large mini-batches; a tighter ELBO (importance weighting) exists but costs more. And the weight between reconstruction and regularisation is a choice: more weight on the KL gives an orderly, samplable latent space but worse reconstructions; less weight gives sharp reconstructions but a space that goes back to resembling that of an autoencoder.

Real pitfalls

1. Inconsistent sums and means. The KL is summed over the latent dimensions and the reconstruction over the pixels: if you use the sum in one and the mean in the other, the ratio of the two terms is wrong by a factor equal to the number of pixels.
2. Parametrising σ directly. It can become negative or zero; use log σ² (or softplus).
3. Sampling z without reparametrising. If z is not differentiable with respect to μ and σ, the gradients do not reach the encoder: no error, but the encoder does not learn.
4. Confusing ELBO and likelihood. A higher ELBO does not automatically imply a higher log p(x) in the same way for two different models: the gap can change.
5. Using the KL "by hand" with another convention. Some implementations use log σ instead of log σ², and the formula changes by a factor 2 in the logarithm term.

The five pitfalls are all balance errors: a scale factor, a sign or a convention changed without noticing silently shifts the ratio between reconstruction and KL. The remedy is the same as the toy case of this lesson: repeat the computation on an example where you know the exact answer.

To check yourself

Try answering before reading the calculation. KL(N(2, 1) ‖ N(0, 1)) is ½(4 + 1 − 0 − 1) = 2 nat, and KL(N(1, 0.25) ‖ N(0, 1)) is ½(1 + 0.25 + 1.386 − 1) = 0.818 nat. With μ = 2, σ = 0.5 and ε = −1, the code is z = 2 + 0.5·(−1) = 1.5. If the ELBO of an image is −100 and log p(x) is −90, the KL between q and the true posterior is 10 nat, because log p(x) = ELBO + KL.

If the encoder always produces μ = 0 and σ = 1 the KL is 0, the code carries no information about x and the decoder cannot reconstruct anything but the mean image: it is the extreme case that in lesson 73 we will call posterior collapse.

The score-function estimator has greater variance because it multiplies the value of the function by the derivative of the logarithm of the density, which can be large and variable, while the reparametrised one uses the derivative of the function, much more regular (24.9 times less variance in our example). The numerical calculations of this paragraph are done by hand, not with the program.

What you can do now

You can derive the ELBO in two ways, E_q[log p(x|z)] − KL(q(z|x) ‖ p(z)), and you know that its gap from log p(x) is exactly KL(q(z|x) ‖ p(z|x)). You know that the KL between a Gaussian and the standard normal is ½(μ² + σ² − log σ² − 1) per component (verified with Monte Carlo: 3.0249 against 3.0243). You know why reparametrisation is needed: it gives a correct gradient with a variance 24.9 times smaller than the score-function estimator (measured on a case with a known answer). And you saw, on a toy case, the ELBO −2.57 stay below log p(x) = −1.26 and match it when the encoder produces the true posterior. In the next lesson we turn all this into a few lines of PyTorch and generate new digits.

In the series, the previous lesson is Autoencoders: Compress, Reconstruct and Understand What a Latent Space Is.