Generative models, latent variables and the ELBO

Before this lesson: the KL divergence and Jensen's inequality (information theory, lesson 03); forward and reverse KL (information theory, lesson 06); Bayes' rule, maximum likelihood and the Gaussian density (the probability and statistics track). Gaussian mixtures and EM appear in the ML fundamentals track (not yet published), but this lesson explains what it needs. The ELBO is also derived in information theory, lesson 09; here we use it to train generative models. Time: about 75 minutes, plus the notebook.

---

Why this matters #

Generative models produce images (Stable Diffusion), text (GPT), audio, protein structures and weather forecasts (GenCast). They differ in architecture, but they share a few ideas, and this lesson gives the first ones:

• Maximum likelihood as a KL divergence. It explains why likelihood-trained models cover all the data but may put mass between modes, and why GANs behave differently (lesson 05).
• Latent variables. Gaussian mixtures, factor analysis, VAEs (lesson 02) and diffusion models (lesson 07, with latents $x_1, \dots, x_T$) all build a complex $p_\theta(x)$ from simple parts.
• The ELBO. It is the training loss of VAEs, and the DDPM loss is a reweighted ELBO. The IWAE bound is the standard way to report the test log-likelihood of a VAE.
• The language of the track. "Exact likelihood", "a bound", "no likelihood": every later lesson places its model in the family table of section 2.

1. Predict first #

2. The picture #

What a generative model is #

We have data $x_1, \dots, x_n$, drawn from an unknown data distribution $p_{\text{data}}$. A generative model is a distribution $p_\theta(x)$ that should be close to $p_{\text{data}}$. We use it in three ways:

1. Sample: draw new $x \sim p_\theta$ (new images, new weather fields).
2. Score: evaluate the density $p_\theta(x)$ of a given $x$ (to detect unusual inputs, to compare models, to compress).
3. Represent and condition: map $x$ to a code $z$, or sample $x$ given a condition $y$ (text-to-image, a forecast given today's weather).

Not every family can do all three. The table below shows which can.

Maximum likelihood fits the forward KL #

The standard way to fit $p_\theta$ is maximum likelihood estimation (MLE): choose the $\theta$ that makes the data most probable, $\max_\theta \frac1n \sum_i \log p_\theta(x_i)$. Section 3 proves that this is the same as minimizing the forward KL divergence $\mathrm{KL}(p_{\text{data}} \,\|\, p_\theta)$. The other order, the reverse KL $\mathrm{KL}(p_\theta \,\|\, p_{\text{data}})$, behaves very differently. The picture first: fit one Gaussian to two bumps.

Interactive figure: loading…

Why the difference? Write both divergences as averages:

• Forward: $\mathrm{KL}(p \,\|\, q) = \mathbb E_{x \sim p}\big[\log p(x) - \log q(x)\big]$. The average is over the data. Wherever the data has mass and $q$ is tiny, $-\log q(x)$ is huge. So $q$ must put mass on every region that has data. This is mode covering (also called mass covering).
• Reverse: $\mathrm{KL}(q \,\|\, p) = \mathbb E_{x \sim q}\big[\log q(x) - \log p(x)\big]$. The average is over $q$. Regions that $q$ ignores cost nothing, but mass in places where $p$ is tiny costs a lot. So $q$ stays inside one mode. This is mode seeking.

The numbers for the two bumps at $\pm 3$: the forward fit is $\mathcal N(0, 10)$, which is the mean and the variance of $p$ (moment matching, proved in section 3). The reverse fit is close to $\mathcal N(2.98, 1.02^2)$ (or its mirror image), with reverse KL $0.689 \approx \log 2$ nats: it captures one of two equal halves. The track's 1D toy, "two Gaussians" $\tfrac12\mathcal N(-1, 1) + \tfrac12\mathcal N(1, 1)$, has its means only $2$ standard deviations apart, so its density has a single flat-topped mode; there the two fits are nearly equal ($\sigma = 1.41$ and $\sigma = 1.39$).

The families of generative models #

The families of this track differ in what they give you. All of them can sample, but only some can evaluate $p_\theta(x)$:

Family	Density $p_\theta(x)$	Sampling cost	Latent space	Lesson
Autoregressive (PixelCNN, GPT)	exact	$D$ sequential steps for $D$ dimensions	none	03
Normalizing flow (RealNVP, Glow)	exact	one pass	$z$, same size as $x$	04
VAE	lower bound (ELBO)	one pass	$z$, usually smaller than $x$	01–02
GAN	none	one pass	$z$	05
Energy-based model	unnormalized: $e^{-E_\theta(x)}/Z_\theta$, $Z_\theta$ unknown	MCMC, many steps	none	06
Diffusion and score models	a bound (ELBO); exact via an ODE, but costly	many denoising steps	noisy $x_t$, same size as $x$	07–10

Interactive figure: loading…

The track uses a few shared toy datasets, so that you can compare the families on the same data: "two Gaussians" in 1D (this lesson), a 1D Gaussian $\mathcal N(0, \sigma_{\text{data}}^2)$ where every diffusion formula has a closed form, "two points" $x_0 \in \{-1, +1\}$, "8 Gaussians" in 2D (the figure above, standard deviation $0.15$ per bump), and "two moons" from scikit-learn, used where a network is trained.

Latent variables: a sampling recipe #

A latent variable model has a hidden variable $z$ and two simple parts:

• a prior $p(z)$, a simple distribution over the latent;
• a likelihood (or decoder) $p_\theta(x \mid z)$, a simple distribution over $x$ for each $z$.

To sample, draw $z$ first, then $x$ given $z$. This is ancestral sampling. The model's density is the marginal, also called the evidence:

Running example: two Gaussians. $z \in \{1, 2\}$ with $p(z = 1) = p(z = 2) = \tfrac12$, and $x \mid z = k \sim \mathcal N(\mu_k, 1)$ with $\mu_1 = -1$, $\mu_2 = +1$. So

The posterior. After we see $x$, Bayes' rule gives the probability of each cause: $p(z = k \mid x) = p(z = k)\,p(x \mid z = k)/p(x)$. In our example the log-odds are simple:

so $p(z = 2 \mid x) = \sigma(2x)$, where $\sigma(u) = 1/(1 + e^{-u})$ is the logistic function.

Numbers at $x = 0.5$. $p(x, z{=}1) = \tfrac12\mathcal N(0.5; -1, 1) = 0.0648$ and $p(x, z{=}2) = \tfrac12\mathcal N(0.5; 1, 1) = 0.1760$. So $p(x) = 0.2408$, $\log p(x) = -1.4238$, and the posterior is $(0.2689, 0.7311) = (\sigma(-1), \sigma(1))$.

A floor under the ceiling #

Choose any distribution $q(z \mid x)$ over the latent; it is our guess for the posterior and is called the variational distribution. Compute

Section 3 proves that this is a floor under the fixed ceiling $\log p(x)$, and that the distance between them is $\mathrm{KL}(q(z \mid x) \,\|\, p(z \mid x))$. Try it before you read the proof.

Interactive figure: loading…

3. Definitions and main results #

Notation. $p_{\text{data}}$ is the data distribution; $p_\theta(x)$ the model. In a latent variable model, $p_\theta(x, z) = p(z)\,p_\theta(x \mid z)$. The KL divergence is $\mathrm{KL}(q \,\|\, p) = \mathbb E_q[\log(q/p)] \ge 0$, with equality only when $q = p$ (Gibbs' inequality, information theory, lesson 03). $H(q) = -\mathbb E_q[\log q]$ is the entropy. All logarithms are natural, so the unit is the nat.

Maximum likelihood is forward KL #

Theorem (MLE = forward KL). For any model $p_\theta$,

$H(p_{\text{data}})$ does not depend on $\theta$. So maximizing the expected log-likelihood is the same as minimizing $\mathrm{KL}(p_{\text{data}} \,\|\, p_\theta)$.

Proof.

1. $\mathrm{KL}(p_{\text{data}} \,\|\, p_\theta) = \mathbb E_{p_{\text{data}}}[\log p_{\text{data}}(x)] - \mathbb E_{p_{\text{data}}}[\log p_\theta(x)]$.
2. The first term is $-H(p_{\text{data}})$ by the definition of entropy.
3. Move it to the other side. $\blacksquare$

With a finite training set, replace $p_{\text{data}}$ by the empirical distribution (equal weight $1/n$ on each $x_i$): the average negative log-likelihood $-\frac1n\sum_i \log p_\theta(x_i)$ is then the cross-entropy between the empirical distribution and the model. For continuous data, $H$ is a differential entropy and can be negative, but it is still a constant in $\theta$.

Corollary (a Gaussian fitted by MLE matches the moments). If $q = \mathcal N(\mu, \sigma^2)$ and $p$ has mean $m$ and variance $v$, the forward KL $\mathrm{KL}(p \,\|\, q)$ is smallest at $\mu = m$, $\sigma^2 = v$.

Proof of the corollary

1. $\mathbb E_p[-\log q(x)] = \tfrac12\log(2\pi\sigma^2) + \mathbb E_p[(x - \mu)^2]/(2\sigma^2)$.
2. Write $x - \mu = (x - m) + (m - \mu)$. The cross term has mean 0, so $\mathbb E_p[(x - \mu)^2] = v + (m - \mu)^2$.
3. The term $(m - \mu)^2$ is smallest at $\mu = m$.
4. With $\mu = m$, the function $\tfrac12\log\sigma^2 + v/(2\sigma^2)$ has derivative $\frac{1}{2\sigma^2} - \frac{v}{2\sigma^4}$ in $\sigma^2$, which is 0 at $\sigma^2 = v$; it is negative before and positive after, so this is the minimum.
5. $\mathrm{KL}(p \,\|\, q) = \mathbb E_p[-\log q] - H(p)$, and $H(p)$ does not depend on $q$. $\blacksquare$

For the two bumps at $\pm 3$: $m = 0$ and $v = 1 + 3^2 = 10$, so the forward fit is $\mathcal N(0, 10)$, as in the figure. The reverse KL has no such formula; it has a local minimum on each bump.

Why a continuous latent is hard #

With $z \in \mathbb R^d$ and a neural network decoder, $p_\theta(x) = \int p(z)\,p_\theta(x \mid z)\,dz$ has no closed form. A grid with $n$ points per axis needs $n^d$ decoder calls. The posterior $p_\theta(z \mid x) = p_\theta(x, z)/p_\theta(x)$ is unknown too, because it needs $p_\theta(x)$. Sampling is still easy (ancestral sampling). So evaluating and training are hard, while sampling is easy. The ELBO gives a quantity we can compute and train.

Theorem 1: the evidence decomposition #

Theorem 1. Let $p_\theta(x) > 0$. Let $q(z \mid x)$ be any distribution that puts mass only where $p_\theta(z \mid x) > 0$ (it may put zero mass where the posterior is positive, as a hard assignment does). Then

So $\mathrm{ELBO}(q) \le \log p_\theta(x)$, with equality exactly when $q(z \mid x) = p_\theta(z \mid x)$.

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$ that $q$ can draw.
3. Multiply and divide by $q$ inside the log: $\log\frac{p(x, z)}{p(z \mid x)} = \log\frac{p(x, z)}{q(z \mid x)} + \log\frac{q(z \mid x)}{p(z \mid x)}$.
4. Take $\mathbb E_q$ of both sides: the first term gives the ELBO, the second gives the KL.
5. $\mathrm{KL} \ge 0$, with equality only when $q$ is the posterior. $\blacksquare$

The Jensen route. Define the importance weight $w(z) = p(x, z)/q(z \mid x)$. Then $\mathbb E_q[w] = \sum_z q\,\frac{p(x, z)}{q} = p(x)$. The logarithm is concave, so Jensen's inequality gives $\log p(x) = \log \mathbb E_q[w] \ge \mathbb E_q[\log w] = \mathrm{ELBO}(q)$. This route is one line, but it does not give the size of the gap. Theorem 1 does: the gap is exactly a KL divergence.

What each term rewards. Split $\log p(x, z) = \log p_\theta(x \mid z) + \log p(z)$:

The reconstruction term rewards latents that explain $x$ well. The rate term charges nats for moving $q$ away from the prior, that is, for storing information about $x$ in $z$. A second form is $\mathrm{ELBO}(q) = \mathbb E_q[\log p(x, z)] + H(q)$: fit plus spread. Without the entropy $H(q)$, the best $q$ would be a point mass on the most probable latent.

What maximizing does. Over $q$, the term $\log p(x)$ is a constant. So maximizing the ELBO over $q$ is the same as minimizing $\mathrm{KL}(q \,\|\, p(z \mid x))$: a reverse KL, so the best $q$ in a restricted family is mode seeking. Over $\theta$, a step can raise $\log p_\theta(x)$ or shrink the gap. So the optimum is biased toward models whose posteriors are easy to fit with $q$.

Theorem 2: EM never decreases the log-likelihood #

When the exact posterior is computable (mixtures, hidden Markov models, factor analysis), we can make the gap zero. Take a dataset $x_1, \dots, x_n$ with one distribution $q_i$ per data point, and $\mathcal L(\theta, q) = \sum_i \mathrm{ELBO}(x_i; \theta, q_i)$. The EM algorithm (expectation–maximization) alternates two steps:

1. E-step: $q_i \leftarrow p_{\theta_t}(z \mid x_i)$ for every $i$. By Theorem 1, every gap is 0.
2. M-step: $\theta_{t+1} = \arg\max_\theta \mathcal L(\theta, q)$. For a Gaussian mixture this has a closed form: weighted means, weighted variances and average responsibilities.

Theorem 2. $\log p_{\theta_{t+1}}(x_{1:n}) \ge \log p_{\theta_t}(x_{1:n})$.

Proof of Theorem 2

1. After the E-step, $\mathcal L(\theta_t, q) = \log p_{\theta_t}(x_{1:n})$ (all gaps are 0).
2. The M-step gives $\mathcal L(\theta_{t+1}, q) \ge \mathcal L(\theta_t, q)$.
3. Theorem 1 at $\theta_{t+1}$ gives $\log p_{\theta_{t+1}}(x_{1:n}) \ge \mathcal L(\theta_{t+1}, q)$.
4. Chain the three: $\log p_{\theta_{t+1}}(x_{1:n}) \ge \mathcal L(\theta_{t+1}, q) \ge \mathcal L(\theta_t, q) = \log p_{\theta_t}(x_{1:n})$. $\blacksquare$

EM is coordinate ascent on $\mathcal L$: the E-step maximizes over $q$, the M-step over $\theta$. It needs the exact posterior in the E-step. For a neural decoder this posterior is unknown, so a VAE cannot use EM.

Amortized inference and the two gaps #

A VAE chooses a family $\mathcal Q$ (for example, Gaussians with diagonal covariance) and one encoder network $q_\phi(z \mid x)$ with parameters $\phi$. For any $x$, the encoder outputs a member of $\mathcal Q$. This is amortized inference: one forward pass replaces a separate optimization for each data point. The model and the encoder are trained together by stochastic gradient ascent on $\sum_i \mathcal L(\theta, \phi; x_i)$, where $\mathcal L(\theta, \phi; x)$ is the ELBO with $q = q_\phi(\cdot \mid x)$. The gap is not 0, so a step that raises the ELBO can lower $\log p_\theta(x)$. There is no monotone guarantee as in EM.

	EM	per-point VI	amortized VI (VAE)
$q$ for each $x$	exact posterior	best member of $\mathcal Q$, optimized for this $x$	encoder output $q_\phi(\cdot \mid x)$
gap	0	approximation gap	approximation gap + amortization gap
cost for a new $x$	one E-step (if tractable)	an inner optimization loop	one forward pass

The decomposition (Cremer, Li and Duvenaud, 2018). Let $q^*_x$ be the best member of $\mathcal Q$ for this $x$, the one with the largest ELBO. Then

Both parts are $\ge 0$.

• The approximation gap equals $\mathrm{KL}(q^*_x \,\|\, p(z \mid x))$: the family $\mathcal Q$ is too simple. Fixes: a richer family (full covariance, normalizing flows, hierarchical latents) or a tighter bound (IWAE, below).
• The amortization gap means that the encoder does not output the best member for this $x$. Fixes: a larger encoder, a few gradient steps on each $q$ after the encoder (semi-amortized inference), or more encoder updates per decoder update.

Theorem 3: the IWAE bound #

Draw $K$ samples from the encoder instead of one, and average the weights inside the logarithm.

Theorem 3 (Burda, Grosse and Salakhutdinov 2016). Let $q(z \mid x) > 0$ wherever $p(x, z) > 0$. Draw $z_1, \dots, z_K$ independently from $q(z \mid x)$, let $w_k = p(x, z_k)/q(z_k \mid x)$ and $\hat p_K = \frac1K\sum_{k=1}^K w_k$. Define the IWAE bound (importance-weighted autoencoder bound)

Then (i) $\mathcal L_1 = \mathrm{ELBO}$; (ii) $\mathcal L_K \le \log p(x)$; (iii) $\mathcal L_K \le \mathcal L_{K+1}$; (iv) if the weights are bounded, $\mathcal L_K \to \log p(x)$ as $K \to \infty$.

Why this is true, in words. $\hat p_K$ is an unbiased estimate of $p(x)$. The log of an unbiased estimate is biased downward (Jensen). Averaging more weights makes the estimate less noisy, so the downward bias shrinks.

Proof of Theorem 3

1. (i) For $K = 1$, $\mathbb E[\log w_1] = \mathbb E_q[\log p(x, z) - \log q(z \mid x)]$, which is the ELBO.
2. (ii) $\mathbb E[w_k] = \int q(z \mid x)\,\frac{p(x, z)}{q(z \mid x)}\,dz = p(x)$, so $\mathbb E[\hat p_K] = p(x)$. Jensen for the concave log: $\mathbb E[\log \hat p_K] \le \log \mathbb E[\hat p_K] = \log p(x)$.
3. (iii) Take $K + 1$ samples and a subset $I$ of size $K$, chosen uniformly at random. 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$, the $K$ weights in $I$ are independent draws, so the right side has expectation $\mathcal L_K$. Hence $\mathcal L_{K+1} \ge \mathcal L_K$.
6. (iv) By the law of large numbers, $\hat p_K \to p(x)$ with probability 1. If the weights are bounded, $\log \hat p_K$ is bounded too (between the log of the smallest and the largest weight), so the limit can be taken inside the expectation. $\blacksquare$

Interactive figure: loading…

How fast the IWAE gap closes, and its cost for training

For large $K$, expand $\log \hat p_K$ around $p(x)$ to second order. With $\mathbb E[\hat p_K] = p(x)$ and $\operatorname{Var}(\hat p_K) = \operatorname{Var}_q(w)/K$:

$$\log p(x) - \mathcal L_K \approx \frac{\operatorname{Var}_q(w)}{2K\,p(x)^2} .$$

The gap shrinks like $1/K$, and it is small when $q$ makes the weights nearly equal. The dashed grey curve in the figure is this prediction. When $q$ is much narrower than the posterior, $\operatorname{Var}_q(w) = \infty$ and the formula fails: the bound still rises, but slowly. (For the Gaussian example above, this happens when $q$'s variance is at most half of the posterior's variance.)

The price is $K$ times the compute. For training, a second price appears: the encoder's gradient gets noisier relative to its size as $K$ grows; its signal-to-noise ratio falls like $1/\sqrt K$ (Rainforth et al., 2018). For evaluation, large $K$ is standard: Burda et al. report test log-likelihoods with $K = 5000$.

Answers to section 1. (1) It covers both bumps. MLE minimizes the forward KL, which is mode covering: for bumps at $\pm 3$ the fit is $\mathcal N(0, 10)$. (2) About 0.3 nats. A point mass has $H(q) = 0$, so its ELBO is $\log p(x, z{=}2)$, and the gap is $-\log p(z{=}2 \mid x) = -\log\sigma(1) = 0.313$ nats. (3) The ELBO can go up: the refinement removes part of the amortization gap. $\log p_\theta(x)$ cannot change, because it does not depend on $q$.

4. Worked example #

Same mixture, same $x = 0.5$. The log-joints are $\log p(x, z{=}1) = -2.7371$ and $\log p(x, z{=}2) = -1.7371$. They differ by exactly the log-odds $2x = 1$.

(a) The ELBO for three choices of $q = (q_1, q_2)$. Use $\mathrm{ELBO} = \mathbb E_q[\log p(x, z)] + H(q)$.

$q$	$\mathbb E_q[\log p(x, z)]$	$H(q)$	ELBO	gap $\log p(x) - \mathrm{ELBO}$
uniform $(0.5, 0.5)$	$-2.2371$	$0.6931$	$-1.5439$	$0.1201$
hard $(0, 1)$	$-1.7371$	$0$	$-1.7371$	$0.3133$
posterior $(0.2689, 0.7311)$	$-2.0060$	$0.5822$	$-1.4238$	$0$

Step by step for the first row: $\mathbb E_q[\log p(x, z)] = 0.5 \cdot (-2.7371) + 0.5 \cdot (-1.7371) = -2.2371$; $H(q) = -2 \cdot 0.5 \log 0.5 = \log 2 = 0.6931$; the sum is $-1.5439$. Check it with Theorem 1: $\mathrm{KL}(q \,\|\, p(z \mid x)) = \tfrac12\log\tfrac{0.5}{0.2689} + \tfrac12\log\tfrac{0.5}{0.7311} = 0.3101 - 0.1899 = 0.1201$. ✓

(b) The two gaps. If $\mathcal Q$ = hard assignments (as in hard EM or k-means), the best member is $z = 2$, so the approximation gap at $x = 0.5$ is $0.3133$ nats, whatever the encoder. If $\mathcal Q$ = all distributions on $\{1, 2\}$, the approximation gap is 0, and only the encoder can be wrong. Take the encoder $q_\phi(z{=}2 \mid x) = \sigma(a + bx)$. The exact posterior is $\sigma(2x)$, so $(a, b) = (0, 2)$ gives zero gap for every $x$. A half-trained encoder $(a, b) = (0, 1)$ outputs $q_2 = \sigma(0.5) = 0.6225$ at $x = 0.5$, and its gap is $\mathrm{KL}((0.3775, 0.6225) \,\|\, (0.2689, 0.7311)) = 0.0280$ nats: a pure amortization gap.

A case where even the best encoder of this form has a gap: change the model to $\tfrac12\mathcal N(0, 1) + \tfrac12\mathcal N(0, 4)$ (same mean, different variances). Its posterior logit is $\log\frac{p(z = 2 \mid x)}{p(z = 1 \mid x)} = -\log 2 + \tfrac38 x^2$, which is symmetric in $x$. A linear encoder $\sigma(a + bx)$ cannot follow it; its best version is almost constant ($a = 0.244$, $b = 0$) and has an average gap of $0.181$ nats over the data. One more input feature, $x^2$, removes the gap: the family "all distributions on $\{1, 2\}$" was never the problem, the encoder was.

(c) IWAE with $K = 2$, by exact enumeration. Take $q$ uniform. The weights are $w(1) = 0.0648/0.5 = 0.1295$ and $w(2) = 0.1760/0.5 = 0.3521$. The pair $(z_1, z_2)$ takes four equally likely values:

• $(1, 1)$: $\log w(1) = -2.0439$;
• $(2, 2)$: $\log w(2) = -1.0439$;
• $(1, 2)$ or $(2, 1)$, probability $\tfrac12$: $\log\frac{0.1295 + 0.3521}{2} = \log 0.2408 = -1.4238 = \log p(x)$.

So $\mathcal L_2 = \tfrac14(-2.0439) + \tfrac14(-1.0439) + \tfrac12(-1.4238) = -1.4839$, and

The gap falls from $0.1201$ to $0.0601$, about half. The large-$K$ formula, with $\operatorname{Var}_q(w)/p(x)^2 = 0.2136$, predicts $0.2136/(2 \cdot 2) = 0.0534$ for $K = 2$. It becomes accurate by $K = 10$: predicted $0.0107$, exact $0.0110$.

5. In code #

NumPy and SciPy check the numbers of section 4. The IWAE bound for a uniform $q$ is computed exactly, without Monte Carlo noise: the number of draws with $z = 2$ is binomial.

import numpy as np
from scipy.stats import norm, binom
from scipy.special import logsumexp, expit

# The two-Gaussians model: p(z = k) = 1/2, x | z = k ~ N(mu_k, 1), mu = (-1, +1)
mu, x = np.array([-1.0, 1.0]), 0.5
log_joint = np.log(0.5) + norm.logpdf(x, mu, 1.0)      # log p(x, z = k)
log_px = logsumexp(log_joint)                          # log p(x)
post = np.exp(log_joint - log_px)                      # p(z | x)
print(log_joint.round(4), round(log_px, 4), post.round(4), round(expit(2 * x), 4))
# [-2.7371 -1.7371] -1.4238 [0.2689 0.7311] 0.7311

def elbo(q):                                           # E_q[log p(x, z)] + H(q)
    q = np.asarray(q, float); m = q > 0
    return np.sum(q[m] * log_joint[m]) - np.sum(q[m] * np.log(q[m]))

def kl(q, p):                                          # KL(q || p) on {1, 2}
    q = np.asarray(q, float); m = q > 0
    return np.sum(q[m] * np.log(q[m] / p[m]))

for q in ([0.5, 0.5], [0.0, 1.0], post):
    print(round(elbo(q), 4), round(log_px - elbo(q), 4), round(kl(q, post), 4))
# -1.5439 0.1201 0.1201      uniform q: gap = KL to the posterior
# -1.7371 0.3133 0.3133      hard assignment
# -1.4238 0.0 0.0            q = posterior: no gap

# Exact IWAE bound for q uniform: the number m of draws with z = 2 is Binomial(K, 1/2)
w = np.exp(log_joint) / 0.5                            # importance weights p(x, z)/q(z)
for K in [1, 2, 10, 100, 1000]:
    m = np.arange(K + 1)
    L_K = np.sum(binom.pmf(m, K, 0.5) * np.log(((K - m) * w[0] + m * w[1]) / K))
    print(K, round(log_px - L_K, 5))
# 1 0.12011 | 2 0.06006 | 10 0.01102 | 100 0.00107 | 1000 0.00011

The gap halves from $K = 1$ to $K = 2$ and then falls like $1/K$. Numerics: in high dimensions the raw weights $w_k$ underflow or overflow. Always compute $\log \hat p_K = \operatorname{logsumexp}(\log w) - \log K$.

6. Where it shows up #

• VAEs (lesson 02). The loss is $-\mathrm{ELBO}$ with a Gaussian encoder and a neural decoder. Test log-likelihoods are reported as $\mathcal L_K$ with large $K$.
• Diffusion models (lesson 07). DDPM is a latent variable model with latents $x_1, \dots, x_T$ and a fixed "encoder", the noising process. Its loss is a reweighted ELBO.
• Latent diffusion (lesson 13). The first stage is an autoencoder trained with a reconstruction term and, in the KL-regularized version, a very small rate term. Latent diffusion models of weather fields and of seismic ground motion use the same two stages.
• Mixtures in climate science. Gaussian mixtures fitted by EM to Northern Hemisphere geopotential height anomalies found three robust circulation regimes (Smyth, Ide and Ghil, J. Atmos. Sci., 1999). The posterior gives each day's regime probabilities.
• Amortized inference in science. One encoder that outputs a posterior for any new observation replaces one optimization per seismogram or per field (simulation-based inference uses the same idea). The amortization gap is the accuracy you pay for that speed.
• MLE vs other objectives. Likelihood-trained models (autoregressive, flows, VAEs, diffusion) cover all modes but may put mass where there is no data; GANs (lesson 05) do not minimize the forward KL and can drop modes ("mode collapse").

7. Common mistakes #

• Mixing up forward and reverse KL. MLE minimizes $\mathrm{KL}(p_{\text{data}} \,\|\, p_\theta)$, the forward KL (mode covering). Variational inference over $q$ minimizes $\mathrm{KL}(q \,\|\, p(z \mid x))$, a reverse KL (mode seeking).
• Comparing models by the ELBO as if it were the likelihood. A model can win by having a smaller gap, not a higher $\log p(x)$. Compare with $\mathcal L_K$ for large $K$.
• Dropping $-\log q$ in a Monte Carlo ELBO. Without the entropy, $q$ collapses to a point mass at the most probable latent.
• Mixing up the two gaps. A bigger or better-trained encoder cannot fix the approximation gap: it is set by the family $\mathcal Q$. A richer family lowers the approximation gap, but does not remove the amortization gap by itself.
• Expecting EM's guarantee from a VAE. EM's monotone guarantee comes from the exact posterior in every E-step (gap 0) and from coordinate ascent. A VAE has a restricted encoder and takes joint gradient steps in $(\theta, \phi)$, so $\log p_\theta(x)$ can go down.
• Ignoring the support condition of IWAE. If $q(z \mid x) = 0$ where $p(x, z) > 0$, then $\hat p_K$ is no longer unbiased and $\mathcal L_K$ does not reach $\log p(x)$ (check yourself 5).

8. Check yourself #

1. Mixture with weights $\pi = (0.2, 0.8)$, means $(-1, 1)$, unit variances. What is $p(z{=}2 \mid x{=}0)$?
   Answer

   The log-odds are $\log\frac{0.8}{0.2} + 2 \cdot 0 = \log 4$, so $p(z{=}2 \mid x{=}0) = 4/5 = 0.8$. At $x = 0$ both likelihoods are equal, so the posterior equals the prior.
2. In the lesson's mixture, what is the gap of a hard assignment at $x = 0$?
   Answer

   The posterior is $(\frac12, \frac12)$, so the gap is $-\log\frac12 = \log 2 = 0.693$ nats. With two components, this is the largest gap that a hard assignment to the more likely component can have.
3. Why does the EM E-step make the bound tight, and why can a VAE not do this step?
   Answer

   Setting $q_i = p_{\theta_t}(z \mid x_i)$ makes $\mathrm{KL}(q_i \,\|\, p(z \mid x_i)) = 0$, so the ELBO equals $\log p_{\theta_t}(x_i)$ (Theorem 1). For a neural decoder, $p_\theta(z \mid x)$ needs $p_\theta(x) = \int p(z)p_\theta(x \mid z)\,dz$, which has no closed form. So the VAE uses a parametric encoder and keeps a gap.
4. A VAE gives test values $\mathcal L_1 = -95.2$, $\mathcal L_{10} = -93.8$, $\mathcal L_{1000} = -93.1$ nats. What do you know about $\log p(x)$ and about the ELBO gap?
   Answer

   $\log p(x) \ge \mathcal L_{1000} \approx -93.1$ (up to Monte Carlo noise). The ELBO gap $\log p(x) - \mathcal L_1$ is at least $95.2 - 93.1 = 2.1$ nats. There is no upper limit from these numbers: $\log p(x)$ can be above $-93.1$.
5. In the lesson's mixture at $x = 0.5$, use the hard $q = (0, 1)$ as the IWAE proposal. What is $\mathcal L_K$ for every $K$, and why does it not approach $\log p(x)$?
   Answer

   Every sample is $z = 2$ with weight $p(x, z{=}2)/1 = 0.1760$, so $\mathcal L_K = \log 0.1760 = -1.7371$ for all $K$. The support condition fails: $q(z{=}1) = 0$ while $p(x, z{=}1) > 0$, so $\mathbb E[\hat p_K] = p(x, z{=}2) < p(x)$.
6. You fit $q = \mathcal N(\mu, \sigma^2)$ to $p = \tfrac12\mathcal N(-2, 1) + \tfrac12\mathcal N(2, 1)$ by forward KL. What are $\mu$ and $\sigma^2$?
   Answer

   Moment matching: $\mu = 0$ and $\sigma^2 = \operatorname{Var}_p(x) = 1 + 2^2 = 5$. (The variance of a mixture is the average within-component variance, 1, plus the variance of the means, 4.)

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. Which KL divergence does maximum likelihood minimize?

   Show answer

2. Why is the forward KL mode covering?

   Show answer

3. Define the ELBO for a latent variable model.

   Show answer

4. What is the gap $\log p(x) - \mathrm{ELBO}(q)$ exactly?

   Show answer

5. In the mixture $\tfrac12\mathcal N(-1, 1) + \tfrac12\mathcal N(1, 1)$, what is $p(z = 2 \mid x)$?

   Show answer

6. Why does EM never decrease $\log p(x)$?

   Show answer

7. What is the amortization gap?

   Show answer

8. Write the IWAE bound.

   Show answer

9. Why is $\mathcal L_K \le \log p(x)$?

   Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Why do we train with maximum likelihood, and what does it optimize?" — It minimizes the forward KL to the data distribution. It is consistent and needs only samples of the data, but it is mode covering: a model that is too simple puts mass between the modes.
• "Derive the ELBO. Which derivation tells you the gap?" — Jensen on $\log\mathbb E_q[w]$ gives the bound; inserting $q/q$ into $\log\frac{p(x, z)}{p(z \mid x)}$ gives ELBO + KL, so the gap is $\mathrm{KL}(q \,\|\, p(z \mid x))$.
• "Why not maximize $\log p(x)$ of a VAE directly?" — The integral over $z$ has no closed form, and the log of a Monte Carlo average is biased downward. The ELBO is a bound we can estimate without bias and differentiate.
• "How would you measure the amortization gap?" — Compare the ELBO of the encoder output, the ELBO after optimizing $q$ for each test point (same family), and $\mathcal L_K$ with large $K$ as an estimate of $\log p(x)$ (Cremer et al., 2018).
• "Is a tighter bound always better for training?" — No. IWAE costs $K$ times more, and the encoder's gradient signal-to-noise ratio falls like $1/\sqrt K$.

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 marginal and the posterior of a 1D Gaussian mixture; check Bayes' rule numerically.
2. The ELBO as a function of $q(z = 2)$; find its maximum and check that the gap equals the KL.
3. EM for a 2D Gaussian mixture on 8-Gaussians data; check that $\log p$ never decreases.
4. The IWAE bound $\mathcal L_K$ for $K = 1, \dots, 100$ on a linear-Gaussian model with a closed-form $p(x)$; check that it rises with $K$ and stays below $\log p(x)$.
5. Forward vs reverse KL fits by grid search; compare with the figure in section 2.

11. Watch and read #

Stanford CS236: Deep Generative Models I 2023 I Lecture 1 - IntroductionStanford Online · Stefano Ermon · 57 min

What a generative model is and the map of model families used in this track.

Maximum Likelihood, clearly explained!!!StatQuest with Josh Starmer · 6 min

A short refresher on maximum likelihood before the KL view of section 3.

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

The ELBO as a floor under the log-evidence, with plots of the gap.

EM algorithm: how it worksVictor Lavrenko · 8 min

The E-step and M-step for a Gaussian mixture, step by step.

• Diederik P. Kingma and Max Welling, An Introduction to Variational Autoencoders (2019): §2.1 (encoder and amortized inference), §2.2 (the ELBO and its gap), §2.6 (importance sampling of $\log p(x)$ and IWAE).
• Yuri Burda, Roger Grosse and Ruslan Salakhutdinov, Importance Weighted Autoencoders (2016): §3 (the bound and Theorem 1) and Appendix A (the proof).
• Chris Cremer, Xuechen Li and David Duvenaud, Inference Suboptimality in Variational Autoencoders (2018): §3.1 (approximation and amortization gaps).
• Stanford CS236: Deep Generative Models, course notes on maximum likelihood learning and variational inference.
• Kevin Murphy, Probabilistic Machine Learning: Advanced Topics (free PDF): ch. 20 (generative models overview), §10.1 (variational inference and the ELBO).