The ELBO: variational inference, VAEs and importance-weighted bounds

Before this lesson: lesson 03 (Jensen's inequality, KL divergence, Gibbs' inequality), lesson 04 (cross-entropy as a loss), lesson 06 (differential entropy, the Gaussian KL, forward and reverse KL). Bayes' rule and posteriors from probability lesson 01; importance sampling from the probability track; gradient estimators from calculus lesson 12. Time: about 80 minutes, plus the notebook.

---

Why this matters #

Bayesian inference needs the posterior $p(z \mid x) = p(x, z)/p(x)$. The numerator is easy: prior times likelihood. The denominator $p(x) = \int p(x, z)\,dz$ is an integral over all latent values, and in most models it has no closed form. The ELBO turns this integral into an optimization problem.

The same bound appears in many places:

• VAEs are trained by maximizing the ELBO (Kingma and Welling, 2013).
• Diffusion models: the DDPM loss is a reweighted ELBO over many noising steps (generative models track).
• EM for mixture models climbs the ELBO (ML fundamentals track, in preparation).
• Variational Bayesian neural networks put a distribution $q$ on the weights and maximize the ELBO.
• Science: fast approximate posteriors for inverse problems, such as seismic tomography, and latent-variable models of single-cell gene expression.

Two facts matter most in practice: how far the bound is from $\log p(x)$, and in which direction $q$ is wrong when its family is too simple. This lesson answers both.

In this lesson $\log$ is the natural logarithm, so all values are in nats; divide by $\ln 2$ to get bits. This is also what PyTorch and NumPy compute.

1. Predict first #

2. The picture #

The model. A latent-variable model is a joint distribution $p(x, z) = p(z)\,p(x \mid z)$. The prior $p(z)$ says which latent values are likely before we see data. The likelihood $p(x \mid z)$ says how the observation $x$ is produced from $z$. We observe one $x$ and keep it fixed.

• The evidence (or marginal likelihood) is $p(x) = \int p(x, z)\,dz$, a sum for discrete $z$.
• The posterior is $p(z \mid x) = p(x, z)/p(x)$, by Bayes' rule.
• A variational distribution $q(z)$ is any distribution over $z$ that we choose. It is our approximation of the posterior.

The simplest case: $z$ has two values. Suppose the model gives, for our observed $x$,

Then $p(x) = 0.02 + 0.18 = 0.2$, so $\log p(x) = \log 0.2 = -1.609$. The posterior is $p(z \mid x) = (0.02/0.2,\ 0.18/0.2) = (0.1,\ 0.9)$. A distribution $q$ on $\{0, 1\}$ is one number $t = q(z{=}1)$. Define

(with $0 \log 0 = 0$ at $t = 0$ and $t = 1$). Four values of $t$:

$t = q(z{=}1)$	0	0.5	0.9	1
$\mathrm{ELBO}(q)$	−3.912	−2.120	−1.609	−1.715
gap $\log p(x) - \mathrm{ELBO}(q)$	2.303	0.511	0	0.105

Interactive figure: loading…

Three things to see:

1. The ELBO is never above $\log p(x)$. It is a floor under a fixed ceiling: the ceiling does not depend on $q$.
2. It touches the ceiling at one $q$ only: $t = 0.9$, the posterior.
3. At $t = 1$, $q$ is a point mass on the most probable latent (the MAP latent). The ELBO is then $\log p(x, z{=}1) = \log 0.18 = -1.715$. That is below the maximum. The ELBO rewards $q$ for spreading over every $z$ that explains $x$, not only the best one.

The weights picture. Write

the importance weight of $z$ (as in importance sampling). If $q(z) > 0$ wherever $p(x, z) > 0$, then:

• $\mathbb{E}_q[w] = \sum_z q(z)\frac{p(x, z)}{q(z)} = \sum_z p(x, z) = p(x)$. The arithmetic mean of the weights is the evidence.
• $\mathrm{ELBO}(q) = \mathbb{E}_q[\log w]$. So $e^{\mathrm{ELBO}}$ is the geometric mean of the weights.

A geometric mean is never larger than the arithmetic mean, and the two are equal only when all values are equal. If $q$ is the posterior, then $w(z) = p(x, z)/p(z \mid x) = p(x)$ for every $z$: all weights are equal and the gap is 0. The more the weights spread, the larger the gap. At $t = 0.5$ the weights are $0.02/0.5 = 0.04$ and $0.18/0.5 = 0.36$. Their arithmetic mean is $0.2 = p(x)$; their geometric mean is $\sqrt{0.04 \cdot 0.36} = 0.12 = e^{-2.120}$.

3. Definitions and main results #

Notation. $\mathrm{KL}(q\,\|\,p) = \mathbb{E}_q[\log(q/p)] \ge 0$, with equality only when $q = p$ (Gibbs' inequality, lesson 03). $H(q) = -\mathbb{E}_q[\log q(z)]$ is the entropy of $q$; for a continuous $z$ it is the differential entropy of lesson 06.

Definition. For a fixed observation $x$ and a distribution $q$ over $z$, the evidence lower bound is

Physicists call $-\mathrm{ELBO}$ the variational free energy.

Theorem 1 (the evidence decomposition) #

Let $q$ be a distribution with $q(z) = 0$ wherever $p(z \mid x) = 0$. Then

Hence $\mathrm{ELBO}(q) \le \log p(x)$, with equality if and only if $q(z) = p(z \mid x)$.

The condition is needed for a finite KL. If $q$ puts mass where the posterior has none, then $\log w = -\infty$ there: the ELBO is $-\infty$ and the KL is $+\infty$.

Proof.

1. $\log p(x)$ does not depend on $z$, so $\log p(x) = \mathbb{E}_q[\log p(x)]$.
2. Bayes' rule: $p(x) = p(x, z)/p(z \mid x)$ for every $z$ with $p(z \mid x) > 0$. This includes every $z$ that $q$ can draw. So $\log p(x) = \mathbb{E}_q\Big[\log \frac{p(x, z)}{p(z \mid x)}\Big]$.
3. Multiply and divide by $q(z)$ inside the log: $\log\frac{p(x, z)}{p(z \mid x)} = \log\frac{p(x, z)}{q(z)} + \log\frac{q(z)}{p(z \mid x)}$.
4. Take $\mathbb{E}_q$ of both terms. The first gives $\mathrm{ELBO}(q)$. The second gives $\mathrm{KL}(q\,\|\,p(z \mid x))$.
5. $\mathrm{KL} \ge 0$, with equality if and only if $q = p(z \mid x)$ (Gibbs). So $\mathrm{ELBO}(q) \le \log p(x)$, with the same equality case. $\blacksquare$

Check with the toy. At $t = 0.5$ the gap is $0.511$, and $\mathrm{KL}((0.5, 0.5)\,\|\,(0.1, 0.9)) = 0.511$ (the first quiz above).

The Jensen route (lesson 03). If $q(z) > 0$ wherever $p(x, z) > 0$, then $p(x) = \mathbb{E}_q[w]$. The log is concave, so Jensen's inequality gives

This route gives the bound in one line, but not the size of the gap. Theorem 1 says the gap is exactly a KL divergence.

A consequence for Monte Carlo. Step 2 says $\log w(z) = \log p(x)$ for every $z$ when $q$ is the posterior. So the Monte Carlo estimate $\frac1S \sum_{s=1}^S \log w(z_s)$ of the ELBO has zero variance at the posterior, for any number of samples $S$.

The two standard forms #

Split $\log p(x, z)$ in two ways. Each split gives a form of the ELBO that is used in practice.

Form 1: energy + entropy.

The first term (the negative energy) pulls $q$ toward latents with a high joint density, as MAP does. The entropy term pays $q$ for being spread out. Without it, the best $q$ would be a point mass at the MAP latent. In the toy at $t = 0.5$: $\mathbb{E}_q[\log p(x, z)] = \frac12(\log 0.02 + \log 0.18) = -2.813$ and $H(q) = \log 2 = 0.693$. The sum is $-2.120$, as in the table.

Form 2: reconstruction − KL to the prior. Use $\log p(x, z) = \log p(x \mid z) + \log p(z)$:

The reconstruction term rewards latents that explain $x$ well. The KL to the prior keeps $q$ close to the prior: moving $q$ away from $p(z)$, so that $z$ can store information about this particular $x$, costs nats. This is the VAE form. Lesson 10 reads it as a trade-off between distortion and rate.

The three expressions, $\log p(x) - \mathrm{KL}(q\,\|\,p(z \mid x))$, energy + entropy, and reconstruction − KL, are the same number. Use the one whose terms you can compute.

A continuous example. Let $z \sim \mathcal{N}(0, 1)$ and $x \mid z \sim \mathcal{N}(z, 1)$, and observe $x = 2$. Section 4 computes everything by hand: $\log p(x) = -2.2655$ and the posterior is $\mathcal{N}(1, \frac12)$. For $q = \mathcal{N}(m, s^2)$ both forms have closed forms.

Interactive figure: loading…

What maximizing the ELBO does #

Over $q$, with the model fixed. $\log p(x)$ does not depend on $q$. So by Theorem 1,

This is the reverse KL of lesson 06. It is the direction we can work with. It needs only expectations under $q$ (we can sample from $q$) and values of $p(x, z)$; the unknown $\log p(x)$ is a constant. The forward KL $\mathrm{KL}(p(z \mid x)\,\|\,q)$ needs expectations under the posterior, which is the thing we cannot compute. Fitting $q$ in this way is called variational inference (VI).

The reverse KL $\mathbb{E}_q[\log(q/p(z \mid x))]$ is large when $q$ puts mass where the posterior is small. So $q$ avoids low-posterior regions. Two consequences:

• Mode seeking. For a posterior with well-separated modes, a single Gaussian $q$ locks onto one mode. Each mode gives a local maximum of the ELBO; the starting point decides which one you find.
• Variances that are too small. Let the posterior be $\mathcal{N}(\mu, \Sigma)$ and let $q$ be a factorized (mean-field) Gaussian $\prod_i \mathcal{N}(m_i, s_i^2)$. The reverse-KL optimum is $m = \mu$ and $s_i^2 = 1/(\Sigma^{-1})_{ii}$. This is the variance of $z_i$ given all other coordinates, and it is never larger than the marginal variance $\Sigma_{ii}$. With unit variances and correlation $\rho = 0.9$,

so $(\Sigma^{-1})_{ii} = 1/(1 - \rho^2)$ and $s_i^2 = 1 - \rho^2 = 0.19$ instead of 1. (The forward KL would give $s_i^2 = \Sigma_{ii} = 1$.)

Over the model parameters $\theta$. Write the model as $p_\theta$. Theorem 1 gives

A step in $\theta$ that raises the ELBO raises the log-likelihood, or moves the model's posterior toward $q$, or both. So it is approximate maximum likelihood with a bias: it prefers models whose posteriors are easy to fit with the chosen family. For a dataset, the sum of the per-example ELBOs is a lower bound on the log-likelihood of the data.

If we set $q$ to the exact posterior before each $\theta$ step, the gap is 0 at the current $\theta$. This is EM (expectation–maximization): the E-step sets $q = p_{\theta}(z \mid x)$, the M-step maximizes the ELBO in $\theta$. Then $\log p_\theta(x)$ never decreases.

Amortized inference, two gaps, and posterior collapse #

A dataset has many observations $x_1, \dots, x_n$, and each has its own posterior. Fitting a separate $q$ for each $x$ is slow. A VAE uses one network, the encoder, that outputs $q_\phi(z \mid x)$ for every $x$. This is amortized inference: the cost of learning to infer is shared by all $x$.

Let $q^*_x$ be the best member of the family for this $x$ (for example, the best diagonal Gaussian). The gap splits in two (Cremer, Li and Duvenaud, 2018):

• The approximation gap is $\mathrm{KL}(q^*_x\,\|\,p(z \mid x))$: the family is too simple.
• The amortization gap is the price of one shared encoder instead of a separate optimization for each $x$.

A different failure is posterior collapse. The decoder models $x$ well without using $z$. Then the best encoder outputs $q_\phi(z \mid x) = p(z)$ for every $x$, and the KL term is 0. The ELBO can still be good, but $z$ carries no information about $x$.

Posterior collapse can happen at the optimum, not only because of bad training (Lucas et al., 2019). The smallest example is a linear VAE. Data $x \sim \mathcal{N}(0, 1)$. Decoder $p(x \mid z) = \mathcal{N}(a z, \sigma^2)$ with prior $z \sim \mathcal{N}(0, 1)$. Then the model gives $x \sim \mathcal{N}(0, a^2 + \sigma^2)$. Keep $\sigma^2$ fixed and learn $a$: the best fit makes $a^2 + \sigma^2 = 1$ when it can. If the decoder noise $\sigma^2$ is 1 or more, the noise alone already explains the data, and the best $a$ is 0.

Interactive figure: loading…

In this example the decoder noise $\sigma^2$ acts as a weight. A Gaussian decoder gives $\log p(x \mid z) = -\frac{\|x - \hat x(z)\|^2}{2\sigma^2} + \text{const}$, so a large $\sigma^2$ makes the reconstruction term cheap compared with the KL term. Common remedies for collapse in real VAEs are KL annealing (the weight of the KL term grows from 0 to 1 during training) and "free bits" (no KL penalty below a few nats per latent group); see Kingma and Welling (2019), §2.8.

The VAE loss #

A VAE has an encoder $q_\phi(z \mid x)$ and a decoder $p_\theta(x \mid z)$, both neural networks, and a prior $p(z) = \mathcal{N}(0, I)$. It maximizes the average ELBO in Form 2 over both $\phi$ and $\theta$. The generative-models track covers the architecture and training practice. Here are the three pieces of the loss.

1. Gaussian encoder. The encoder outputs a mean vector $\mu_\phi(x)$ and a log-variance vector $\log \sigma^2_\phi(x)$, both of length $d$ (the latent dimension): $q_\phi(z \mid x) = \mathcal{N}\big(\mu_\phi(x), \operatorname{diag}(\sigma_\phi^2(x))\big)$.

2. The KL term in closed form. For one dimension,

and for a diagonal Gaussian the KL is the sum of this over the $d$ dimensions.

Derivation of the Gaussian KL to $\mathcal{N}(0, 1)$

1. $\log q(z) = -\frac12\log(2\pi\sigma^2) - \frac{(z - \mu)^2}{2\sigma^2}$ and $\log p(z) = -\frac12\log(2\pi) - \frac{z^2}{2}$.
2. Under $q$: $\mathbb{E}_q[(z - \mu)^2] = \sigma^2$ and $\mathbb{E}_q[z^2] = \mu^2 + \sigma^2$.
3. So $\mathbb{E}_q[\log q] = -\frac12\log(2\pi\sigma^2) - \frac12$ and $\mathbb{E}_q[\log p] = -\frac12\log(2\pi) - \frac12(\mu^2 + \sigma^2)$.
4. Subtract: $\mathrm{KL} = \mathbb{E}_q[\log q] - \mathbb{E}_q[\log p] = -\frac12\log\sigma^2 - \frac12 + \frac12(\mu^2 + \sigma^2)$. $\blacksquare$

The KL is 0 only at $\mu = 0$, $\sigma^2 = 1$: the function $\sigma^2 - 1 - \log\sigma^2$ is $\ge 0$ with its minimum at $\sigma^2 = 1$.

3. The reconstruction term with the reparameterization trick. The term $\mathbb{E}_{q_\phi}[\log p_\theta(x \mid z)]$ has no closed form for a neural decoder, so we estimate it with one (or a few) samples of $z$. But we also need its gradient with respect to $\phi$, and $\phi$ is inside the distribution we sample from. The reparameterization trick writes the sample as a deterministic function of $\phi$ and a fixed noise:

Now $\varepsilon$ does not depend on $\phi$, and backpropagation goes through $z$ to $\mu_\phi$ and $\sigma_\phi$. This gradient estimator has much lower variance than the score-function (REINFORCE) estimator (calculus lesson 12).

What $\log p_\theta(x \mid z)$ is. It depends on the decoder's output distribution:

• Bernoulli decoder (binary or $[0, 1]$ pixels): the decoder outputs logits $\ell$, and

where $\sigma$ is the sigmoid. This is minus the binary cross-entropy (lesson 04).

• Gaussian decoder with a fixed variance $\sigma_x^2$: $\log p_\theta(x \mid z) = -\frac{\|x - \hat x\|^2}{2\sigma_x^2} - \frac{D}{2}\log(2\pi\sigma_x^2)$, where $\hat x$ is the decoder output and $D$ the data dimension. Up to a constant, this is a scaled squared error.

In PyTorch the loss for a batch is a few lines (the encoder and decoder are any networks):

import torch
import torch.nn.functional as F

def vae_loss(x, encoder, decoder):
    mu, logvar = encoder(x)                     # q(z | x) = N(mu, diag(exp(logvar)))
    eps = torch.randn_like(mu)
    z = mu + torch.exp(0.5 * logvar) * eps      # reparameterization
    logits = decoder(z)                         # Bernoulli decoder p(x | z)
    rec = -F.binary_cross_entropy_with_logits(logits, x, reduction="none").sum(dim=1)
    kl = 0.5 * (mu**2 + logvar.exp() - 1 - logvar).sum(dim=1)
    elbo = rec - kl                             # one-sample ELBO estimate per example
    return -elbo.mean()                         # loss = -ELBO, in nats per example

Two details matter. Sum over the data dimensions and the latent dimensions (not a mean), or the two terms get the wrong relative weight. And the loss is in nats per example; divide by $D \ln 2$ to report bits per dimension, the usual unit for images.

Theorem 2 (the importance-weighted bound; Burda, Grosse and Salakhutdinov, 2015) #

The ELBO is the log of a geometric mean of weights. If we average several weights before taking the log, the average spreads less, and Jensen's inequality loses less. This is the IWAE (importance weighted autoencoder) bound.

Let $q(z) > 0$ wherever $p(x, z) > 0$. Draw $z_1, \dots, z_K$ independently from $q$, write $w_k = w(z_k) = p(x, z_k)/q(z_k)$, and let $\hat p_K = \frac1K \sum_{k=1}^K w_k$. Define

Then:

1. $\mathcal{L}_1 = \mathrm{ELBO}(q)$;
2. $\mathcal{L}_K \le \mathcal{L}_{K+1} \le \log p(x)$ for every $K$;
3. if $w$ is bounded and the ELBO is finite, $\mathcal{L}_K \to \log p(x)$ as $K \to \infty$.

(Burda et al. state part 3 with bounded $w$ only; the finite ELBO is the extra assumption that the short proof below uses.)

Toy check. In the two-value example with uniform $q$, the weights are $0.04$ and $0.36$, each with probability $\frac12$. For $K = 2$: both draws are $0.04$ with probability $\frac14$, both $0.36$ with probability $\frac14$, and one of each with probability $\frac12$. So

between $\mathcal{L}_1 = -2.120$ and $\log p(x) = -1.609$. The gap falls from 0.511 to 0.255.

Proof of Theorem 2

1. Part 1. For $K = 1$, $\hat p_1 = w_1$, and $\mathbb{E}[\log w_1]$ is the ELBO.
2. Upper bound. Each weight has mean $\mathbb{E}_q[w] = \int p(x, z)\,dz = p(x)$. So $\mathbb{E}[\hat p_K] = p(x)$: $\hat p_K$ is an unbiased estimate of $p(x)$. Jensen for the concave log: $\mathcal{L}_K = \mathbb{E}[\log \hat p_K] \le \log \mathbb{E}[\hat p_K] = \log p(x)$.
3. Monotonicity, setup. Take $K + 1$ samples. Let $I$ be a uniformly random subset of $\{1, \dots, K + 1\}$ of size $K$, independent of the samples. Each index is in $I$ with probability $K/(K + 1)$. So $\mathbb{E}_I\Big[\frac1K\sum_{k \in I} w_k\Big] = \frac1K \cdot \frac{K}{K+1}\sum_{k=1}^{K+1} w_k = \hat p_{K+1}$.
4. Jensen over $I$, with the samples fixed: $\log \hat p_{K+1} = \log \mathbb{E}_I\Big[\frac1K\sum_{k \in I} w_k\Big] \ge \mathbb{E}_I\Big[\log \frac1K\sum_{k \in I} w_k\Big]$.
5. Average over the samples. For each fixed $I$, $\frac1K\sum_{k \in I} w_k$ is an average of $K$ independent weights, so its expected log is $\mathcal{L}_K$. Hence $\mathcal{L}_{K+1} \ge \mathcal{L}_K$.
6. Limit. By the strong law of large numbers, $\hat p_K \to p(x)$ with probability 1. If $w \le B$, then $\frac1K\sum_k \log w_k \le \log \hat p_K \le \log B$ (the left inequality is "geometric mean ≤ arithmetic mean"). The left side is an average of independent integrable variables when the ELBO is finite. This sandwich lets the limit move inside the expectation (uniform integrability), so $\mathcal{L}_K \to \log p(x)$. $\blacksquare$

Why the gap shrinks like $1/K$. Both the ELBO and $\mathcal{L}_K$ have the form $\mathbb{E}[\log(\text{an unbiased estimate of } p(x))]$. The Jensen gap of the log grows with the spread of the estimate. The average of $K$ weights has variance $\operatorname{Var}_q(w)/K$. For large $K$, a second-order Taylor expansion of $\log$ around $p(x)$ (the delta method) gives

So the gap shrinks like $1/K$, and it is small when $q$ makes the weights nearly equal.

Interactive figure: loading…

The price.

• Compute. $K$ times the work per example.
• High dimensions. Importance weights degrade when $z$ has many dimensions: one weight tends to dominate the sum.
• Training signal. A tighter bound is not always better for training. Average $M$ independent estimates of the gradient of $\mathcal{L}_K$. The signal-to-noise ratio of the encoder's gradient is of order $\sqrt{M/K}$, while the decoder's is of order $\sqrt{MK}$ (Rainforth et al., 2018). So a large $K$ helps the decoder and hurts the encoder.
• Evaluation. For evaluation, a large $K$ is standard: Burda et al. report test log-likelihoods with $K = 5000$.

Another view: IWAE is an ELBO with a richer approximation

Draw $K$ samples from $q$ and keep one of them with probability proportional to its weight. The kept sample has an implicit distribution $q_{\mathrm{EW}}$, which is closer to the posterior than $q$. One can show $\mathrm{ELBO}(q) \le \mathcal{L}_K \le \mathrm{ELBO}(q_{\mathrm{EW}}) \le \log p(x)$ (Cremer, Morris and Duvenaud, 2017; Kingma and Welling, 2019, §2.6). So IWAE does not change the model; it fits a better approximate posterior that is built from $K$ samples of the simple one.

Answers to section 1. (1) The maximum over all $q$ is $\log p(x)$, reached only by the posterior $q = p(z \mid x)$, where the KL gap is 0 (Theorem 1). (2) Tighter, or equal: $\mathcal{L}_{10} \ge \mathcal{L}_1$. And $\mathcal{L}_K \le \log p(x)$ always, as an expectation (Theorem 2). A single estimate can land above $\log p(x)$ by noise. (3) On one mode, with a variance that is too small for the whole posterior: the reverse KL is mode seeking. (If the modes overlap strongly, $q$ can sit between them.)

4. Worked example #

Model. $z \sim \mathcal{N}(0, 1)$, $x \mid z \sim \mathcal{N}(z, 1)$, observed $x = 2$. Everything is in closed form, so every claim can be checked.

(a) Exact answers. $x = z + \text{noise}$ with independent unit-variance noise, so $x \sim \mathcal{N}(0, 2)$ and

Posterior: $p(z \mid x) \propto \exp\big(-\frac{z^2}{2} - \frac{(x - z)^2}{2}\big)$. The exponent is a quadratic in $z$ with coefficient $-\frac12 \cdot 2$ on $z^2$, so the precision is $1 + 1 = 2$. The linear coefficient is $x$, so the mean is $x/2 = 1$. Hence $p(z \mid x) = \mathcal{N}(1, \frac12)$.

(b) The ELBO for $q = \mathcal{N}(m, s^2)$, in the reconstruction − KL form:

• Reconstruction: $\log p(x \mid z) = -\frac12\log 2\pi - \frac12(x - z)^2$, and $\mathbb{E}_q[(x - z)^2] = (x - m)^2 + s^2$ (squared bias plus variance). So

• KL to the prior (the VAE formula above):

$q$	reconstruction	KL to prior	ELBO	$\log p(x) - \mathrm{ELBO}$	$\mathrm{KL}(q\,\|\,\mathcal{N}(1, \frac12))$
$\mathcal{N}(1, 1)$	−1.9189	0.5000	−2.4189	0.1534	0.1534
$\mathcal{N}(2, \frac12)$	−1.1689	2.0966	−3.2655	1.0000	1.0000
$\mathcal{N}(1, \frac12)$	−1.6689	0.5966	−2.2655	0	0

Read the rows step by step.

• $\mathcal{N}(2, \frac12)$ explains the data best (the highest reconstruction term), but pays 2.10 nats to move away from the prior.
• $\mathcal{N}(1, 1)$ has the right mean but is too wide.
• The last column is computed separately, with the Gaussian KL formula of lesson 06:

It equals the gap in every row, as Theorem 1 says. For $\mathcal{N}(2, \frac12)$: $\frac12(0 + \frac{0.5 + 1}{0.5} - 1) = \frac12(3 - 1) = 1$ exactly.

• Check of Form 1 for $\mathcal{N}(1, 1)$: $\mathbb{E}_q[\log p(x, z)] = -3.8379$ and $H(q) = \frac12\log(2\pi e) = 1.4189$. The sum is $-2.4189$. ✓

(c) Maximize over $q$. Write the ELBO as a function of $m$ and $s^2$:

• $\partial/\partial m$: $(2 - m) - m = 0$, so $m = 1$.
• $\partial/\partial s^2$: $-\frac12 - \frac12 + \frac{1}{2s^2} = 0$, so $s^2 = \frac12$.

The optimum is exactly the posterior. The Gaussian family contains it, so the approximation gap is 0 and the bound is tight.

(d) IWAE for $q = \mathcal{N}(2, \frac12)$. Both $q$ and the posterior have variance $\frac12$, so the normalizing constants cancel and

Under $q$, $z \sim \mathcal{N}(2, \frac12)$, so $3 - 2z \sim \mathcal{N}(3 - 4,\ 4 \cdot \frac12) = \mathcal{N}(-1, 2)$. So $\log w - \log p(x)$ is Gaussian with mean $-1$ and variance 2. The ELBO gap is $-\mathbb{E}[\log w - \log p(x)] = 1$, as in the table. Section 5 estimates the IWAE gaps by sampling:

$K$	1	10	100
$\log p(x) - \mathcal{L}_K$	1.00	0.19	0.03

Compare with the large-$K$ formula. Here $w/p(x) = e^{Y}$ with $Y \sim \mathcal{N}(-1, 2)$, and $\mathbb{E}[e^{2Y}] = e^{2(-1) + \frac12 \cdot 4 \cdot 2} = e^{2}$ (the mean of a log-normal variable), so

Here $\mathbb{E}[e^{Y}] = e^{-1 + \frac12 \cdot 2} = 1$, as it must be, because $\mathbb{E}_q[w] = p(x)$. The formula gives $6.39/(2K) = 0.032$ at $K = 100$ (close to 0.028) and $0.32$ at $K = 10$ (too large: $K = 10$ is not yet "large").

5. In code #

NumPy and SciPy check the numbers of section 4 by other routes: numerical integration for the ELBO, the law of the log-weights for IWAE, and a numerical optimizer for the mean-field example. The notebook has these cells and the exercises.

import numpy as np
from scipy.stats import norm
from scipy.special import logsumexp
from scipy.optimize import minimize

# Model: z ~ N(0, 1), x | z ~ N(z, 1). Observed x = 2, so x ~ N(0, 2) and p(z | x) = N(1, 1/2).
x = 2.0
log_px = norm(0, np.sqrt(2)).logpdf(x)
print(round(log_px, 4))                                  # -2.2655

def kl_gauss(m1, v1, m2, v2):                            # KL(N(m1, v1) || N(m2, v2)), 1-D
    return 0.5 * (np.log(v2 / v1) + (v1 + (m1 - m2) ** 2) / v2 - 1)

# (a) ELBO by numerical integration on a grid; the gap equals KL(q || posterior).
z = np.linspace(-10, 12, 200_001)
log_joint = norm(0, 1).logpdf(z) + norm(z, 1).logpdf(x)  # log p(z) + log p(x | z)
for m, s2 in [(1.0, 1.0), (2.0, 0.5), (1.0, 0.5)]:
    q = norm(m, np.sqrt(s2))
    elbo = np.trapezoid(q.pdf(z) * (log_joint - q.logpdf(z)), z)
    print(m, s2, round(elbo, 4), round(log_px - elbo, 4), round(kl_gauss(m, s2, 1.0, 0.5), 4))
# 1.0 1.0 -2.4189 0.1534 0.1534
# 2.0 0.5 -3.2655 1.0 1.0
# 1.0 0.5 -2.2655 0.0 0.0

# (b) IWAE gap for q = N(2, 1/2): log w - log p(x) = 3 - 2z ~ N(-1, 2) under q.
#     Never average the weights directly: use logsumexp(log w) - log K.
rng = np.random.default_rng(1)
for K, n in [(1, 1_000_000), (10, 200_000), (100, 40_000)]:
    Y = rng.normal(-1.0, np.sqrt(2.0), size=(n, K))
    print(K, round(-np.mean(logsumexp(Y, axis=1) - np.log(K)), 2))
# 1 1.0
# 10 0.19
# 100 0.03

# (c) Reverse KL with a factorized q underestimates the marginal variances.
rho = 0.9
S = np.array([[1.0, rho], [rho, 1.0]])
P = np.linalg.inv(S)
def reverse_kl(log_var):                                 # KL(N(0, diag) || N(0, S))
    D = np.diag(np.exp(log_var))
    return 0.5 * (np.trace(P @ D) - 2 + np.log(np.linalg.det(S) / np.linalg.det(D)))
res = minimize(reverse_kl, x0=np.zeros(2), method="Nelder-Mead", options={"xatol": 1e-8, "fatol": 1e-12})
print(np.exp(res.x).round(3))      # [0.19 0.19] = 1 - rho^2; the true marginal variances are 1

The grid integral in (a) is a route that does not use the closed forms of section 4. In (b) the log of an average of weights must be computed in log space: the weights themselves overflow or underflow for real models.

6. Where it shows up #

• VAEs. The loss is $-\mathrm{ELBO}$ with an encoder $q_\phi(z \mid x)$ and a decoder $p_\theta(x \mid z)$, trained jointly by stochastic gradients with the reparameterization trick. The IWAE objective replaces $\mathcal{L}_1$ by $\mathcal{L}_K$. The generative-models track covers architectures, β-VAEs and VQ-VAEs.
• Diffusion models. The DDPM loss is a reweighted ELBO in which the "encoder" is the fixed noising process (generative models track). The first stage of latent diffusion (Stable Diffusion) is an autoencoder trained with a reconstruction loss (plus perceptual and adversarial terms) and a very small KL-to-prior penalty: Form 2 with a strongly down-weighted KL term.
• EM and variational Bayes. EM is coordinate ascent on the ELBO with an exact E-step. Variational Bayesian neural networks ("Bayes by Backprop") put a Gaussian $q$ on the weights and maximize the ELBO; the KL to the prior acts as a regularizer.
• Topic models and probabilistic programming. Latent Dirichlet allocation was fitted with mean-field VI. Pyro and NumPyro fit user-written models by stochastic VI ("SVI") on the ELBO.
• ML for science: inverse problems. Variational inference gives approximate posteriors for seismic tomography at a fraction of the cost of Monte Carlo sampling (Zhang and Curtis, 2020). The mean-field result of section 3 is the warning: a factorized $q$ reports too little uncertainty when the parameters are correlated, as neighbouring velocity cells are.
• ML for science: single-cell biology. scVI (Lopez et al., 2018) is a VAE for single-cell RNA counts; its latent space is used for clustering, batch correction and differential expression.
• Simulation-based inference: the other KL. If you can simulate pairs $(z, x)$ from the model, you can train $q_\phi(z \mid x)$ by maximizing $\mathbb{E}_{p(z, x)}[\log q_\phi(z \mid x)]$ (neural posterior estimation). This minimizes the forward KL from the posterior, averaged over $x$, so $q_\phi$ tends to cover all modes.

7. Common mistakes #

• Reading the ELBO as $\log p(x)$. It is $\log p(x)$ minus a gap. A model can win on the ELBO by having a smaller gap, not a higher likelihood. Compare models with $\mathcal{L}_K$ for a large $K$.
• The wrong KL direction. Maximizing the ELBO minimizes $\mathrm{KL}(q\,\|\,p(z \mid x))$, not $\mathrm{KL}(p(z \mid x)\,\|\,q)$. This is why VI underestimates uncertainty.
• Confusing the two KL terms. The gap is the KL to the posterior. The term inside Form 2 is the KL to the prior. Making the KL to the prior small does not make the bound tight.
• Dropping the entropy term. Without $H(q)$ (or the $-\log q$ part of a Monte Carlo estimate), the optimum collapses to a point mass at the MAP latent.
• Mean instead of sum in the VAE loss. Averaging the reconstruction over pixels but summing the KL over latents changes the weight of the two terms (like a large β). The ELBO sums both.
• "Larger $K$ is always better." The bound is tighter, but the cost grows like $K$ and the encoder's gradient becomes noisier.
• Treating one estimate as a guaranteed bound. $\mathcal{L}_K \le \log p(x)$ holds in expectation; a single Monte Carlo estimate can land above $\log p(x)$.

8. Check yourself #

1. Two latent values with $p(x, z{=}0) = 0.1$ and $p(x, z{=}1) = 0.3$, and $q$ uniform. Compute $\log p(x)$, the ELBO, the gap and $\mathrm{KL}(q\,\|\,p(z \mid x))$.
   Answer

   $p(x) = 0.4$, so $\log p(x) = -0.916$. The weights are $w = (0.1/0.5,\ 0.3/0.5) = (0.2, 0.6)$, so ELBO $= \frac12(\log 0.2 + \log 0.6) = \frac12\log 0.12 = -1.060$. Gap $= -0.916 + 1.060 = 0.144$. The posterior is $(0.25, 0.75)$, so $\mathrm{KL} = \frac12\log\frac{0.5}{0.25} + \frac12\log\frac{0.5}{0.75} = \frac12\log\frac43 = 0.144$. ✓
2. Why does variational inference minimize $\mathrm{KL}(q\,\|\,p(z \mid x))$ and not $\mathrm{KL}(p(z \mid x)\,\|\,q)$?
   Answer

   The reverse KL needs only expectations under $q$ (which we can sample) and values of $\log p(x, z)$. The unknown $\log p(x)$ is a constant, so maximizing the ELBO minimizes the reverse KL. The forward KL needs expectations under the posterior, which we cannot compute or sample.
3. In the worked example, restrict $q$ to the family $\mathcal{N}(m, 1)$. What is the best $m$, and what is the remaining gap called?
   Answer

   $\partial/\partial m$: $(2 - m) - m = 0$, so $m = 1$. ELBO $= -2.4189$, gap $= 0.1534 = \mathrm{KL}(\mathcal{N}(1, 1)\,\|\,\mathcal{N}(1, \frac12))$. It is the approximation gap of this family: no member does better.
4. The posterior is a 2-D Gaussian with unit variances and correlation $\rho = 0.5$. What variances does a mean-field $q$ get from the reverse KL? From the forward KL?
   Answer

   Reverse KL: $1/(\Sigma^{-1})_{ii} = 1 - \rho^2 = 0.75$. Forward KL: the marginal variances, 1.
5. If $q$ is the exact posterior, what is $\mathcal{L}_K$ for each $K$, and what is the variance of its Monte Carlo estimate?
   Answer

   Every weight equals $p(x)$, so $\log \hat p_K = \log p(x)$ for every draw. Hence $\mathcal{L}_K = \log p(x)$ for all $K$, and the estimate has zero variance.
6. A VAE has a 3-D latent. For every input $x$ the encoder outputs $\mu = (0, 0, 0)$ and $\sigma^2 = (1, 1, 1)$. What is the KL term, and what does this tell you about the model?
   Answer

   The KL term is $0$: $q_\phi(z \mid x) = \mathcal{N}(0, I) = p(z)$. Since this holds for every $x$, the encoder ignores the input and $z$ carries no information about $x$. This is posterior collapse: the decoder models $x$ without $z$.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Write the evidence decomposition for a latent-variable model.

   Show answer

2. Derive the ELBO with Jensen's inequality in one line.

   Show answer

3. Write the ELBO as reconstruction minus a KL. Which KL is it?

   Show answer

4. Write the ELBO as energy plus entropy.

   Show answer

5. Maximizing the ELBO over $q$ minimizes which divergence, and what follows?

   Show answer

6. Give $\mathrm{KL}(\mathcal{N}(\mu, \sigma^2)\,\|\,\mathcal{N}(0, 1))$.

   Show answer

7. What is the reparameterization trick, and why is it used?

   Show answer

8. Define the IWAE bound $\mathcal{L}_K$.

   Show answer

9. How does the IWAE gap $\log p(x) - \mathcal{L}_K$ behave for large $K$?

   Show answer

10. What are the approximation gap and the amortization gap?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Derive the ELBO." — Two ways: insert $q(z)/q(z)$ into $\log\frac{p(x, z)}{p(z \mid x)}$ to get ELBO + KL (this also gives the gap); or Jensen on $\log \mathbb{E}_q[p(x, z)/q(z)]$.
• "What does each term of the VAE loss do?" — The reconstruction term rewards latents that explain $x$. The KL to the prior limits the information in $z$ and keeps $q$ close to the prior, so that samples from the prior decode to realistic data.
• "Why can't we just maximize $\log p(x)$?" — It needs $\int p(x \mid z)\,p(z)\,dz$, which has no closed form for a neural decoder, and naive Monte Carlo with samples from the prior has a huge variance.
• "Why is VI mode seeking? Why does it underestimate variance?" — The reverse KL punishes $q$ for mass where the posterior is small. A mean-field Gaussian gets the conditional variances $1/(\Sigma^{-1})_{ii}$.
• "Why the reparameterization trick?" — To backpropagate through sampling: $z = \mu + \sigma\varepsilon$ moves the randomness into $\varepsilon$, which does not depend on the parameters.
• "Why is IWAE tighter, and why not train with $K = 1000$?" — The average of $K$ weights spreads less, so the Jensen gap shrinks like $1/K$. But the cost is $K$ times larger, and the encoder's gradient signal-to-noise ratio falls like $1/\sqrt{K}$.
• "What is posterior collapse? How do you detect and reduce it?" — The KL term goes to 0 and the encoder outputs the prior for every $x$. Watch the KL per latent dimension; use KL annealing, free bits, or a weaker decoder.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution:

1. The ELBO of the two-value example as a function of $t$; check that the gap equals the KL to the posterior.
2. The closed-form ELBO of the linear-Gaussian model against a Monte Carlo estimate with reparameterized samples.
3. iwae_bound with logsumexp, for $K = 1, 10, 100$, and the $1/K$ behaviour of the gap.
4. Fit $q = \mathcal{N}(m, s^2)$ by gradient ascent on the ELBO with reparameterization gradients, and recover the exact posterior $\mathcal{N}(1, \frac12)$.
5. A mean-field $q$ for a correlated 2-D Gaussian posterior: the variances $1 - \rho^2$ of section 3.

11. Watch and read #

Evidence Lower Bound (ELBO) - CLEARLY EXPLAINED!Kapil Sachdeva · 12 min

A step-by-step derivation of the ELBO and its KL gap, as in Theorem 1.

Variational Inference | Evidence Lower Bound (ELBO) | Intuition & VisualizationMachine Learning & Simulation · 25 min

The ELBO under the fixed log-evidence, with an interactive plot like the figures of section 2.

Variational AutoencodersArxiv Insights · 15 min

The VAE architecture, the reparameterization trick and β-VAEs.

Probabilistic ML - Lecture 24 - Variational InferenceTübingen Machine Learning · Philipp Hennig · 88 min

A full university lecture: evidence lower bounds, mean-field factorization, and VI for a Gaussian mixture.

• Kingma and Welling, An Introduction to Variational Autoencoders (2019): §2.2 (the decomposition and its gap), §2.4 (reparameterization), §2.6 (importance sampling and IWAE), §2.8 (KL annealing, free bits).
• Kingma and Welling, Auto-Encoding Variational Bayes (2013), the original VAE paper.
• Burda, Grosse and Salakhutdinov, Importance Weighted Autoencoders (2015): §3 (the bound and Theorem 1), Appendix A (the proof).
• Blei, Kucukelbir and McAuliffe, Variational Inference: A Review for Statisticians (2017): §2.2 (the ELBO), §2.3 (mean-field; Figure 1 shows the too-small variances).
• Cremer, Li and Duvenaud, Inference Suboptimality in Variational Autoencoders (2018): the approximation and amortization gaps.
• Rainforth et al., Tighter Variational Bounds are Not Necessarily Better (2018): the signal-to-noise argument against large $K$ in training.
• Lucas, Tucker, Grosse and Norouzi, Don't Blame the ELBO! A Linear VAE Perspective on Posterior Collapse (2019): the linear VAE of the third figure.
• Kevin Murphy, Probabilistic Machine Learning: Advanced Topics, free PDF: chapter 10 (variational inference) and chapter 21 (VAEs).
• Pyro, SVI Part I: the "guide" ($q$) and the ELBO in a probabilistic programming language.