Variational Autoencoders
From the problem up — built piece by piece, with the tools introduced only when the problem demands them.
1 The problem
We have a collection of high-dimensional examples — images of faces, sentences in English, recordings of speech, and (further afield) things like robot trajectories. Call any single example $x$. We have lots of them: $x_1, x_2, \ldots, x_N$.
We want to build a machine that can do two things.
-
Score plausibility. Given a new $x$, tell us how likely it is to be the same kind of thing as our examples. A clear photo of a face should score high; random pixel noise should score low.
-
Generate. Produce new examples that look like they came from the same source — new faces, new sentences, new trajectories — that we've never seen before but that fit in.
Why this is hard. There is no rule we can write down for "looks like a face" or "sounds like English." These concepts live in the data, not in any formula we have access to. So we cannot hand-engineer the machine; it has to learn what "plausible" means from the examples themselves.
These two jobs — scoring and generating — feel different but they aren't. They are two queries against the same underlying object: a description of which $x$'s are likely and which aren't. If we had that description, scoring is just looking up a number, and generating is drawing samples from it.
So our real problem is: learn that description from data. Everything else — autoencoders, neural networks, KL divergence, the reparameterization trick — is machinery built to do this one thing.
This is exactly the language probability was designed for, which is where we go next.
2 The probabilistic reframe
We just said our two jobs — scoring and generating — are queries against one underlying object: a description of which $x$'s are likely. Probability theory has a name for that object: a distribution.
So we make our first move:
Assume our examples $x_1, \ldots, x_N$ were drawn from some unknown distribution over high-dimensional space. Call it $p_{\text{data}}(x)$.
aside This is a choice, not a fact. We're viewing the world this way because it gives us tools — but every modeling assumption is a load-bearing decision we live with downstream.
If we had $p_{\text{data}}$, both jobs would be one-liners.
Given a new $x$, evaluate $p_{\text{data}}(x)$. Big number → plausible. Small number → unlikely.
Draw $x \sim p_{\text{data}}$. The sample is, by construction, "the same kind of thing" as our examples.
But we don't have $p_{\text{data}}$. We have samples from it. So we build our own $p_\theta(x)$ — a family of distributions controlled by parameters $\theta$ — and try to make it match.
$$\text{Make } p_\theta \approx p_{\text{data}}.$$
Scoring and generating both fall out of this. Everything we build from here on serves this single objective.
-
"Match" needs a definition. Two distributions can be close or far apart, but we haven't said how to measure that.
-
$p_\theta$ needs a form. We have to pick what kind of distribution $p_\theta$ is — a Gaussian, a mixture, something more elaborate. Different choices are different tradeoffs between flexibility (can it match $p_{\text{data}}$?) and tractability (can we actually compute with it?).
-
$\theta$ needs to be chosen from data. Even after picking the form, we need a principle that says which $\theta$ is the right one given our examples.
Each of those questions has a long history. The rest of the document walks through the answers in roughly the order they were discovered: how to choose $\theta$ (maximum likelihood, Bayesian inference), why simple forms for $p_\theta$ aren't enough (latent variables), how to compute with latent variables (EM, then variational inference), and only then — once the framework is built — what changes when we let neural networks into it.
3 A tour of distributions
Before we get principled, let's get our hands on the actual objects. We'll meet four distributions and, for each one, do the two jobs from Section 1: score a point, and generate one.
Both jobs cost something, and the cost depends on the distribution. Scoring is sometimes a formula evaluation. Generating is almost never just one operation. By the end you'll see that "how hard is sampling?" is an axis the whole field of generative modeling is organized along.
3.1 Bernoulli — the warmup
$x = 1 \Rightarrow$ return $\theta$. $x = 0 \Rightarrow$ return $1 - \theta$.
Draw $u \sim \text{Uniform}(0, 1)$. Return $1$ if $u < \theta$, else $0$. The uniform-to-Bernoulli mapping is exact — no approximation.
Used everywhere: binary pixel models (MNIST), binary classification, per-pixel decoders in VAEs. Cheap on both ends — we won't be this lucky again.
3.2 Gaussian — the workhorse
Substitute $x, \mu, \sigma^2$ into the formula. One exponential, one division. Pointwise — no integration needed.
The Gaussian CDF has no closed-form inverse, so the usual "draw $u$ and invert" trick doesn't apply. Box–Muller takes two uniforms in and gives two independent standard Gaussians out:
Already not as primitive as Bernoulli — but still constant-cost. The trick exploits a structural property of the Gaussian that most distributions don't have.
Why Box–Muller works (polar decomposition)
If $X_1, X_2$ are independent $\mathcal{N}(0,1)$, their joint density is $p(x_1, x_2) = \tfrac{1}{2\pi} \exp(-\tfrac{x_1^2 + x_2^2}{2})$. Switch to polar coordinates; the Jacobian is $r$:
$$p(r,\theta) = \tfrac{r}{2\pi} \exp(-r^2/2).$$This factorizes — $\theta$ is uniform on $[0, 2\pi)$ (easy), and $R^2$ is exponential with mean 2 (also easy via inverse-CDF, $r = \sqrt{-2 \ln u_1}$). Sample each marginal, convert back to Cartesian, done.
The most common output distribution in generative modeling. A VAE decoder predicting pixel intensities outputs exactly this — one Gaussian per pixel, with $\mu, \sigma$ from the network.
3.3 Mixture of Gaussians — the bridge
A single Gaussian is one bump. Real data isn't.
Evaluate each component's density at $x$, weight by $\pi_k$, add. Closed form, cheap.
You can't "evaluate the formula" to sample — you have to commit to a component first, then draw from it.
The intermediate $k$ is generated by the process but never observed in the data. That asymmetry is the entire latent-variable idea — see the box below.
$k$ is a latent variable. We never observe which component a point came from — we just see $x$. But the generative process used $k$ to produce $x$. This is the entire idea behind latent variable models, which we'll formalize later.
Components fixed at $\mathcal{N}(-2, 0.5^2)$, $\mathcal{N}(0, 0.7^2)$, $\mathcal{N}(2.2, 0.4^2)$. $\pi_3$ adjusts so the weights sum to 1. Sampled points are colored by which component generated them — information you wouldn't have in a real dataset.
Cheap on both ends — if you know the components. Inferring which component a point came from given just $x$ is harder, and that posterior inference problem is what VAEs eventually generalize.
3.4 A 2D warning
Move to two dimensions and a non-Gaussian shape and the cracks show.
Pick a moon, an angle $t \in [0, \pi]$, add Gaussian noise. Trivial — the distribution is defined by this procedure.
To write $p(x_1, x_2)$ for an arbitrary 2D point, you'd have to integrate over all $(t, \text{moon})$ that could plausibly have produced it. No clean formula. The procedure that generates the data is not the same as the formula that evaluates it.
Now flip the asymmetry.
The same object — "a distribution" — pays its cost in different currencies depending on which way you defined it.
3.5 The axis that organizes the whole field
At the level of what we want, scoring and generating are the same job. At the level of computation, they pull in different directions — and every generative model family makes a different bargain.
| Family | Score $p_\theta(x)$ | Sample $x \sim p_\theta$ | The bargain |
|---|---|---|---|
| Bernoulli / Gaussian | cheap formula | cheap | Too rigid for real high-dim data |
| Mixture model | cheap sum | pick + sample | Components don't scale to images |
| VAE | intractable | one forward pass | Sample fast, evaluate via a bound |
| Autoregressive (PixelCNN, GPT) | factorizes | sequential | Exact likelihood, slow generation |
| Normalizing flow | change of vars | invert | Both cheap, model form constrained |
| Diffusion | via objective | many steps | Best samples today, slow inference |
| GAN | not modeled | one forward pass | Sample only — no $p(x)$ at all |
We're about to build, slowly and carefully, the row labeled VAE. Every design choice — encoder, decoder, KL term, reparameterization trick — exists to manage a particular asymmetry: we want generation to be a single forward pass, and we're willing to pay for that by not being able to compute $p_\theta(x)$ exactly, only bound it.
4 Choosing $\theta$ from data
Three things to settle were left from Section 2. The form of $p_\theta$ we picked in Section 3. Now: which $\theta$ given the data, and what does "match" mean.
4.1 The setup
We have data $\mathcal{D} = \{x_1, \ldots, x_N\}$ and a chosen family $p_\theta(x)$. We need a principle that picks one $\theta$ from the family.
-
Maximum likelihood (MLE). Pick the $\theta$ that makes the data we actually saw look most probable. No prior beliefs about $\theta$ — the data does all the talking.
-
Maximum a posteriori (MAP). Same idea, but with a prior $p(\theta)$ that biases the answer toward "reasonable" $\theta$'s before seeing data. MLE plus a regularizer, essentially.
4.2 Likelihood — the principle
a trap to name $L(\theta)$ is not a probability over $\theta$. It does not integrate to 1 over $\theta$, and it doesn't have to. It is $p_\theta(\mathcal{D})$ — a function of $\theta$ that happens to use the data-distribution machinery to define its values.
4.3 MLE worked — Gaussian
Differentiate w.r.t. $\mu$:
$$\frac{\partial \ell}{\partial \mu} = \frac{1}{\sigma^2}\sum_i (x_i - \mu) = 0 \;\Longrightarrow\; \sum_i x_i = N\mu.$$Differentiate w.r.t. $\sigma^2$:
$$\frac{\partial \ell}{\partial \sigma^2} = -\frac{N}{2\sigma^2} + \frac{1}{2\sigma^4}\sum_i(x_i - \mu)^2 = 0.$$Add points and watch the MLE Gaussian (indigo) reshape live. The dashed vertical line is $\hat\mu$; the shaded band is $\hat\mu \pm \hat\sigma$. The fit only sees these specific points — try clustering them on one side to see the mean shift.
4.4 MLE worked — Bernoulli
Let $k = \sum_i x_i$ (the number of 1's in $N$ Bernoulli trials):
$$\ell(\theta) = k \log \theta + (N - k) \log(1 - \theta).$$Differentiate, set to zero:
$$\frac{k}{\theta} - \frac{N-k}{1-\theta} = 0 \;\Longrightarrow\; k(1-\theta) = (N-k)\theta.$$4.5 MLE worked — Mixture (and where the wall is)
Take $\partial \ell / \partial \mu_k$:
$$\frac{\partial \ell}{\partial \mu_k} = \sum_i \underbrace{\frac{\pi_k \mathcal{N}(x_i \mid \mu_k, \sigma_k^2)}{\sum_j \pi_j \mathcal{N}(x_i \mid \mu_j, \sigma_j^2)}}_{=\, \gamma_{ik}} \cdot \frac{x_i - \mu_k}{\sigma_k^2} = 0.$$That fraction $\gamma_{ik}$ is the posterior probability that point $i$ came from component $k$. Solve for $\mu_k$:
$$\mu_k = \frac{\sum_i \gamma_{ik}\, x_i}{\sum_i \gamma_{ik}}.$$But $\gamma_{ik}$ depends on $\mu_k$ (and all the other parameters). The equation is implicit. There's no algebraic solution — the unknown is on both sides.
Understand this deeply — what's a "latent variable," why does it cause this, and where else does this same problem live?
A The setup — one coin, two dice
Forget mixtures for a moment. Here's the smallest latent variable model in the world:
I flip a coin. If it's heads, I roll die A (fair, 1–6). If tails, I roll die B (loaded toward sixes). You see the die number. You don't see the coin.
latent z the coin flip — $z \in \{H, T\}$. Unobserved.
observed x the die outcome — $x \in \{1, \ldots, 6\}$. This is your data.
parameters $\pi = p(z = H)$, and the two die distributions $p_A(x), p_B(x)$.
You do this 1000 times. You see $x_1, \ldots, x_{1000}$. Your job: MLE for $\pi, p_A, p_B$.
B Imagine you saw the coin too — the problem decouples
Say you had the full pairs $(z_i, x_i)$. The log-likelihood splits:
$$\log p(x, z) = \sum_i \log \pi(z_i) + \sum_i \log p_{z_i}(x_i).$$That gives you three independent estimation problems:
- $\pi$: count the fraction of $z_i = H$.
- $p_A$: take the $x_i$'s where $z_i = H$, compute their empirical distribution.
- $p_B$: take the $x_i$'s where $z_i = T$, compute their empirical distribution.
Each is a one-line MLE. Closed form. Trivial. When the latent is observed, the dataset partitions into independent sub-problems. This is the world latent variable models wish they lived in.
C Now hide the coin — the log-of-sum appears
You only see the dice rolls. The probability of one observed roll is now
$$p(x_i) = p(x_i, z = H) + p(x_i, z = T) = \pi \cdot p_A(x_i) + (1 - \pi) \cdot p_B(x_i).$$That sum is the cost of not seeing $z$ — you had to sum over all the values $z$ could have been. The log-likelihood becomes
$$\log p(x_1, \ldots, x_N) = \sum_i \log\!\left[\pi \cdot p_A(x_i) + (1 - \pi) \cdot p_B(x_i)\right].$$The log is now outside a sum. And $\log(a + b) \neq \log a + \log b$ — there's no algebraic move that pulls the $p_A$ and $p_B$ apart. The parameters of die A and die B are now entangled inside every single term. You cannot estimate one without simultaneously thinking about the other.
D The exact problem, named — chicken and egg
Suppose you see the roll $x_i = 6$. What does this tell you about die A vs die B?
It depends on what you already think they are.
- If you already believe die B is loaded toward sixes, this roll is evidence "$z$ was T."
- If you already believe die A is loaded toward sixes, the same roll is evidence "$z$ was H."
The roll itself is ambiguous evidence until you've already committed to a model. So:
To estimate $p_A$ and $p_B$, you need to know which rolls came from which die. To know which rolls came from which die, you need to know $p_A$ and $p_B$. Neither side closes. The dataset can't be partitioned the way it could in part B, because the partitioning is itself unknown and depends on the answer.
This is exactly what the gradient equations in 4.5 are telling us, in algebraic form. The implicit $\gamma_{ik}$ that depends on $\mu_k$ which depends on $\gamma_{ik}$ — calculus's way of saying "you need the labels to find the parameters, and you need the parameters to find the labels."
E The general fact — every latent variable model has this shape
This wasn't a quirk of coins and dice. It's a universal structural fact about latent variable models:
| case | likelihood per datapoint | closed form? |
|---|---|---|
| Latents observed: data is $(z_i, x_i)$ | $\log p(z_i, x_i)$ — log of a single term | Yes — problem decouples |
| Latents hidden: data is $x_i$ only | $\log \sum_z p(x_i, z)$ — log of a sum | No — parameters entangled |
The mixture's $\sum_k$ in its density is structurally identical to the coin-and-dice's $\sum_z$. The mixture is a latent variable model precisely because that sum is hiding a marginalized latent. Same shape, same breakage.
F Where this same problem appears in the real world
Once you've seen the pattern, you start to see it everywhere. Each row below is a real problem with the same chicken-and-egg structure — you observe one thing, but to model it you'd need a hidden thing you don't get to see.
| problem | observed $x$ | latent $z$ |
|---|---|---|
| 3D vision from 2D images | pixels of an image | the 3D scene — geometry, materials, lighting, camera. Infinitely many scenes can produce the same pixels. |
| Speech recognition (HMMs) | audio waveform | the phoneme / word sequence that was spoken. To estimate the sound-to-phoneme mapping you'd want the alignment; to align you'd need the mapping. |
| Topic modeling (LDA) | words in a document | which topic each word came from. Topic distributions need word labels; word labels need topic distributions. |
| SLAM (robot localization & mapping) | noisy sensor readings over time | the robot's true pose and the map of the world. Pose depends on map, map depends on pose — that's literally what "simultaneous" means in the name. |
| Cocktail party / source separation | mixed audio at one microphone | the individual speakers' voices being mixed together. You'd need the mixing pattern to recover voices; you'd need voices to estimate the mixing. |
| Genetics — ancestry inference | DNA from individuals | which ancestral populations contributed, and in what proportion. Population allele frequencies need ancestries; ancestries need frequencies. |
Same shape in every row: observe a marginal, latent has been integrated out, parameters entangled inside a log of a sum. This is one problem with many faces, not many unrelated problems.
G What this predicts about the rest of the document
If "log of a marginal sum" is the universal obstruction, then everything that comes next is a different strategy for living with it:
- EM (Section 5): don't fight the sum. Alternate. Pretend you know $z$, do the easy estimate, then re-estimate $z$ given your new parameters, repeat.
- Variational inference (Section 6): when even computing $p(z \mid x)$ for the alternation is intractable, approximate it with a tractable family.
- VAE (Section 7): when $z$ is continuous and high-dimensional (the sum becomes an integral over $\mathbb{R}^d$), use neural networks for the approximating family, and reparameterize so gradients flow.
Three sections, one obstruction, three different ways of dodging it.
For the Gaussian and the Bernoulli, $\nabla \ell = 0$ closed — we got an answer. For the mixture, it doesn't. The reason: the mixture is a latent variable model, and $p_\theta(x) = \sum_k \pi_k \mathcal{N}(x \mid \mu_k, \sigma_k^2)$ is really a marginal: $p_\theta(x) = \sum_k p_\theta(x, k)$. The latent $k$ got summed out, and the log of that sum is what breaks the algebra. Every algorithm we build from here on — EM, variational inference, the VAE itself — is a different answer to "what do you do when the latent variable makes the gradient implicit?"
a preview The fixed-point structure of the mixture equation ($\mu_k$ depends on $\gamma_{ik}$, which depends on $\mu_k$) is exactly what the EM algorithm exploits: alternate between computing $\gamma_{ik}$ given the current $\mu_k$ (E-step) and updating $\mu_k$ given $\gamma_{ik}$ (M-step). EM is what comes next.
4.6 MLE ↔ KL — what "match" actually means
Back in Section 2 we left "$p_\theta \approx p_{\text{data}}$" without saying what "$\approx$" meant. MLE secretly answers it.
As $N \to \infty$, the average log-likelihood approaches an expectation under $p_{\text{data}}$:
$$\frac{1}{N}\ell(\theta) = \frac{1}{N}\sum_i \log p_\theta(x_i) \;\longrightarrow\; \mathbb{E}_{x \sim p_{\text{data}}}[\log p_\theta(x)].$$Now expand the KL between $p_{\text{data}}$ and $p_\theta$:
$$\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)].$$The first term does not depend on $\theta$. So minimizing KL is equivalent to maximizing the second term — which is exactly the (asymptotic) log-likelihood.
4.7 MAP — adding a prior
MLE has no opinion about $\theta$ before seeing the data. MAP gives it one.
Data: $x_1, \ldots, x_N \sim \mathcal{N}(\mu, \sigma^2)$, $\sigma^2$ known. Prior: $\mu \sim \mathcal{N}(\mu_0, \tau^2)$.
Setting the gradient of $\log p(\mathcal{D} \mid \mu) + \log p(\mu)$ to zero:
$$\hat\mu_{\text{MAP}} = \frac{\tfrac{N}{\sigma^2}\, \bar x \,+\, \tfrac{1}{\tau^2}\, \mu_0}{\tfrac{N}{\sigma^2} + \tfrac{1}{\tau^2}}.$$4.8 The wall — and what's next
The Gaussian and Bernoulli MLEs were one-line closed forms. The mixture wasn't — and the obstruction was specific.
A latent variable model defines a joint $p_\theta(x, z)$. The data likelihood is the marginal:
$$p_\theta(x) = \int p_\theta(x, z)\, dz \;=\; \int p_\theta(x \mid z)\, p_\theta(z)\, dz.$$5 EM — the classical answer
The mixture left us with an equation where the unknown was on both sides. EM is the elegant way to live with that.
5.1 The fixed-point insight
Recall the mixture's MLE equation from 4.5:
Stop trying to solve it algebraically. Alternate. Pick a starting $\theta$, compute $\gamma$ from it, use that $\gamma$ to update $\theta$, repeat. This converges to a (local) maximum of the likelihood — and the reason it works is deeper than it looks. We'll see in 5.4.
5.2 EM — the two-step recipe
For each datapoint $i$ and each possible latent value $k$:
$$\gamma_{ik} \;=\; p(z = k \mid x_i, \theta^{\text{old}}) \;=\; \frac{\pi_k\, \mathcal{N}(x_i \mid \mu_k, \sigma_k^2)}{\sum_j \pi_j\, \mathcal{N}(x_i \mid \mu_j, \sigma_j^2)}.$$Using the responsibilities as fixed weights, the maximization closes:
$$\mu_k = \frac{\sum_i \gamma_{ik}\, x_i}{\sum_i \gamma_{ik}}, \quad \sigma_k^2 = \frac{\sum_i \gamma_{ik}(x_i - \mu_k)^2}{\sum_i \gamma_{ik}}, \quad \pi_k = \frac{1}{N}\sum_i \gamma_{ik}.$$5.3 EM on the coin and dice — the chicken-and-egg, solved
The mixture's chicken-and-egg in 4.5 had the simplest possible form: one coin and two dice. Now we have the algorithm. Let's run it on that exact setup and watch the deadlock dissolve.
recall from §4.5 Coin $z \in \{H, T\}$ with $p(z = H) = \pi$. If H, roll die A with distribution $p_A$; if T, roll die B with distribution $p_B$. You see the dice rolls $x_1, \ldots, x_N \in \{1, \ldots, 6\}$. You don't see the coin. Want to recover $\pi, p_A, p_B$.
Don't escape the chicken-and-egg. Iterate it. Pick any starting $\pi, p_A, p_B$. Use them to guess the coin flips. Use those guesses to refit the parameters. Repeat.
For each observed roll $x_i$, given the current $\pi, p_A, p_B$, compute the probability the coin was H:
$$\gamma_i \;=\; p(z = H \mid x_i) \;=\; \frac{\pi \cdot p_A(x_i)}{\pi \cdot p_A(x_i) + (1-\pi) \cdot p_B(x_i)}.$$With the $\gamma_i$ in hand, refit each parameter as a weighted MLE — exactly Part B of the §4.5 explainer (the world where the coin was observed), but with fractional counts instead of hard counts.
$$\pi^{\text{new}} \;=\; \frac{1}{N}\sum_i \gamma_i, \quad\quad p_A^{\text{new}}(k) \;=\; \frac{\sum_i \gamma_i \cdot \mathbb{1}[x_i = k]}{\sum_i \gamma_i}, \quad\quad p_B^{\text{new}}(k) \;=\; \frac{\sum_i (1-\gamma_i) \cdot \mathbb{1}[x_i = k]}{\sum_i (1-\gamma_i)}.$$Data: $x = [6, 6, 6, 3, 1]$.
Initialization: $\pi = 0.5$, $p_A$ uniform ($= 1/6$ each), $p_B = [0.1, 0.1, 0.1, 0.1, 0.1, 0.5]$.
E-step. For each $x_i$, compute $\gamma_i$:
| $x_i$ | $\pi \cdot p_A(x_i)$ | $(1-\pi) \cdot p_B(x_i)$ | $\gamma_i = p(H \mid x_i)$ |
|---|---|---|---|
| 6 | $0.5 \cdot 0.167 = 0.083$ | $0.5 \cdot 0.5 = 0.250$ | $\approx 0.25$ |
| 6 | 0.083 | 0.250 | ≈ 0.25 |
| 6 | 0.083 | 0.250 | ≈ 0.25 |
| 3 | $0.5 \cdot 0.167 = 0.083$ | $0.5 \cdot 0.1 = 0.050$ | ≈ 0.625 |
| 1 | 0.083 | 0.050 | ≈ 0.625 |
Notice: the 6s get assigned mostly to die B (low $\gamma$); the 1 and 3 get assigned mostly to die A (high $\gamma$). Bayes is already using the structure $p_B$ has (a fat tail at 6) to credit the 6s to die B.
M-step. $\sum_i \gamma_i \approx 1.99$, so
$$\pi^{\text{new}} = 1.99 / 5 \approx 0.40.$$For $p_A^{\text{new}}$, the "1-count" is $\gamma_5 \approx 0.625$, "3-count" is $\gamma_4 \approx 0.625$, "6-count" is $3 \cdot 0.25 = 0.75$. Normalize by $\sum_i \gamma_i \approx 1.99$:
$$p_A^{\text{new}} \approx [0.31,\; 0,\; 0.31,\; 0,\; 0,\; 0.38].$$For $p_B^{\text{new}}$, swap $\gamma_i \to (1 - \gamma_i)$; the 6-count dominates:
$$p_B^{\text{new}} \approx [0.12,\; 0,\; 0.12,\; 0,\; 0,\; 0.75].$$One iteration in, and both dice have already moved sharply toward their roles — die A is broadening to fit the 1/3, die B is concentrating on 6. The "ambiguous" 6s ($\gamma \approx 0.25$ vs $0.75$) become less ambiguous in the next iteration, because the model now genuinely believes die B loves 6s. The chicken-and-egg dissolves through self-reinforcement.
Top strip: each square is one observed roll, labeled with its die value. Rose = "model thinks this came from die A" ($\gamma \to 1$); indigo = "die B" ($\gamma \to 0$); gray = ambiguous. Middle: the current $p_A, p_B$ as bar charts. Bottom: log-likelihood over iterations — should be monotonically non-decreasing.
Try Re-init several times: roughly half the time, "die A" learns the loaded die and "die B" learns the uniform one — same likelihood, swapped labels. That's the label-switching non-identifiability §5.6 will name.
why this converges Each iteration provably does not decrease the data log-likelihood. The full proof — that EM is coordinate ascent on the ELBO $\mathcal{F}(q, \theta)$, which lower-bounds $\log p_\theta(x)$ — is in 5.5. The visceral version: watch the bottom plot in the demo go up monotonically and never come back down.
-
Survives → all the way to the VAE. The alternation structure (E-step / M-step), responsibilities as Bayes' rule, and the weighted-MLE form of the M-step. The VAE keeps every one of these, just with neural networks in place of the closed-form updates.
-
Breaks when the latent is continuous. The E-step here is one line of Bayes' rule because $z \in \{H, T\}$ — two terms in the denominator. When $z \in \mathbb{R}^d$, that denominator becomes the intractable integral all over again. That cliff is exactly what §5.6 names and §6 climbs.
5.4 Watch EM in action on a 2D GMM
Points are colored by their responsibility — gray when ambiguous, sharp colors when one component dominates. Ellipses are the current $\mathcal{N}(\mu_k, \Sigma_k)$ at $1\sigma$. Try re-initializing several times: you'll sometimes get a clean fit, sometimes a stuck local optimum.
5.5 Why EM works — the lower bound view
EM looks like a heuristic. It isn't. The whole algorithm is coordinate ascent on a lower bound of the log-likelihood, and recognizing this is what allows EM to generalize to variational inference.
For fixed $\theta$, the ELBO is largest when KL is zero — which happens when
$$q^\star(z) = p(z \mid x, \theta).$$So the E-step is: set $q$ to the true posterior. After this step, the bound is tight: $\mathcal{F}(q^\star, \theta) = \log p_\theta(x)$.
With $q$ held at $q^\star$, the ELBO becomes (dropping the $\theta$-independent $-\mathbb{E}_{q^\star}[\log q^\star]$):
$$\theta^{\text{new}} = \arg\max_\theta\; \mathbb{E}_{q^\star(z)}\!\left[\log p_\theta(x, z)\right].$$This is just a weighted MLE — exactly the M-step updates from 5.2.
EM alternately maximizes the ELBO $\mathcal{F}(q, \theta)$ over $q$ (E-step) and $\theta$ (M-step). Each step never decreases the ELBO, and the ELBO is always ≤ $\log p_\theta(x)$, so the likelihood is monotonically non-decreasing. Convergence to a stationary point is guaranteed.
5.6 Where EM breaks
The E-step requires
$$p(z \mid x, \theta) = \frac{p_\theta(x, z)}{\int p_\theta(x, z')\, dz'}.$$$\log p_\theta(x)$ for a mixture is non-convex. EM converges to a local maximum that depends on initialization. Try the demo: re-init a few times on the same data — you'll get visibly different fits.
5.7 Bridge — what to do when the E-step is impossible
EM's recipe is "maximize the ELBO, alternating between $q$ and $\theta$." That recipe is good. The specific E-step — "set $q$ to the true posterior" — is what breaks.
Keep the ELBO. Keep coordinate ascent. Stop demanding that $q$ equal the true posterior — instead, restrict $q$ to a family we can compute with, and find the best $q$ in that family by optimization. That generalization is variational inference.
6 Variational inference
EM was coordinate ascent on the ELBO with $q$ = the true posterior. Drop that last constraint and you get VI.
6.1 The move — approximate, don't compute
-
Pick a tractable family $\mathcal{Q}$. E.g., diagonal Gaussians $q(z) = \mathcal{N}(z \mid \mu, \mathrm{diag}(\sigma^2))$. Cheap to evaluate, cheap to sample.
-
Optimize $q$ over the family. Instead of solving $q = p(z \mid x, \theta)$ in closed form (which we can't), search for the $q \in \mathcal{Q}$ that maximizes the ELBO.
-
Accept the gap. The best $q \in \mathcal{Q}$ generally won't equal the true posterior — there's a residual KL. We've turned an integration problem into an optimization problem.
trade you've just made Inference is no longer exact — but it's always tractable, regardless of how complicated $p_\theta(x \mid z)$ is. That trade is what lets variational inference scale to models EM cannot touch.
6.2 The ELBO, properly
Same identity as 5.4, written with a parametric variational family $q_\phi(z \mid x)$.
Maximizing $\mathcal{L}(\theta, \phi; x)$ over $\theta$ and $\phi$ jointly does two things at once:
- Pushes up $\log p_\theta(x)$ — improves the generative model.
- Pushes down $\mathrm{KL}(q_\phi \,\|\, p(z \mid x))$ — improves the inference network.
6.3 The ELBO, rearranged — reconstruction minus KL
The form above is theoretically clean. There's an equivalent form that becomes the VAE training objective.
Start with the ELBO and expand $\log p_\theta(x, z) = \log p_\theta(x \mid z) + \log p(z)$:
$$\mathcal{L} = \mathbb{E}_{q_\phi(z|x)}\!\left[\log p_\theta(x \mid z) + \log p(z) - \log q_\phi(z \mid x)\right].$$Pull the data-dependent term out and group the rest:
$$\mathcal{L} = \mathbb{E}_{q_\phi(z|x)}\!\left[\log p_\theta(x \mid z)\right] + \mathbb{E}_{q_\phi(z|x)}\!\left[\log \frac{p(z)}{q_\phi(z \mid x)}\right].$$The second expectation is $-\mathrm{KL}(q_\phi(z \mid x) \,\|\, p(z))$.
What you'll see in any VAE training code — reconstruction loss − KL term (or, equivalently, negative reconstruction + KL, to minimize). It is not a separate, heuristic combination. It is the ELBO, rearranged.
6.4 Mean-field VI — the classical instance
VI predates neural nets. The original instance: pick the family $\mathcal{Q}$ to be fully factorized distributions.
If the true posterior has correlated dimensions, the best factorized $q$ "shrinks" — it fits inside the high-density region rather than wrapping around it. The classical picture:
True posterior (gray) is a correlated 2D Gaussian. The mean-field approximation (indigo) is the best factorized Gaussian — it cannot tilt. Notice it underestimates the variance along the principal direction. This bias is intrinsic to the mean-field assumption.
6.5 Amortized inference — one network, all datapoints
Classical VI fits a separate $q_i(z)$ for each datapoint $i$, by optimizing local parameters. With $N$ in the millions, that's prohibitive.
Amortization adds a second source of suboptimality on top of the family restriction:
- The best $q \in \mathcal{Q}$ for a particular $x$ may not equal the true posterior — approximation gap.
- The best $q_\phi(z \mid x)$ for that same $x$, with $\phi$ shared across all data, may not equal the best $q \in \mathcal{Q}$ — amortization gap.
Toy model: $p(z) = \mathcal{N}(0, 1)$, $p(x \mid z) = \mathcal{N}(z, 0.5^2)$, observation $x = 1.5$. The true posterior $p(z \mid x)$ is itself Gaussian (gray). Slide $\mu_q, \sigma_q$ to move the variational $q(z)$ (indigo). The ELBO equals $\log p(x)$ (orange line, fixed) only when $q$ matches the true posterior. Otherwise, ELBO $<$ $\log p(x)$, and the gap is KL.
6.6 What still blocks us
We have everything we need conceptually: a model $p_\theta(x, z)$, an amortized variational posterior $q_\phi(z \mid x)$, a single objective $\mathcal{L}(\theta, \phi)$. We want to maximize it by SGD over $(\theta, \phi)$. One last obstruction:
The reconstruction term is
$$\nabla_\phi \mathbb{E}_{q_\phi(z|x)}\!\left[\log p_\theta(x \mid z)\right].$$We need a gradient with respect to $\phi$ — but $\phi$ controls the distribution we're sampling from. We can't just sample $z$ and differentiate; the sampling operation breaks the gradient chain.
If you've seen RL, this is the same problem. Both VI and policy-gradient RL face one identity:
$$\nabla_\phi \mathbb{E}_{q_\phi(z)}[f(z)] \;=\; \mathbb{E}_{q_\phi(z)}\!\left[f(z) \cdot \nabla_\phi \log q_\phi(z)\right].$$Same estimator, two communities, two names.
| field | $q_\phi(z)$ | $f(z)$ | known as |
|---|---|---|---|
| VI / VAE | posterior $q_\phi(z \mid x)$ | ELBO integrand | score-function estimator |
| RL | policy $\pi_\phi(a \mid s)$ | return $R$ | REINFORCE / likelihood ratio |
6.7 Bridge — everything is in place except one trick
Generative model $p_\theta(x \mid z) p(z)$. Amortized approximate posterior $q_\phi(z \mid x)$. Objective: ELBO = reconstruction − KL. Optimize jointly by SGD. The only thing left is to make that "sample $z$, then backprop" step actually differentiable.