Latent variables: EM and the ELBO

A cause you never see

Look at the scatter plot below before reading on. The points fall into loose clumps, and once you have seen a clump you cannot unsee it: there are a few groups, each with its own centre and its own spread. Now ask the fitting question. Fit a single Gaussian and you get one fat oval smeared across the middle, a model that calls the gaps between the clumps likely when nothing lives there. The failure is not in the parameters; it is in the shape of the model. One bell curve cannot represent a landscape with several hills.

So we write down a model with more than one hill, and immediately the accounting gets strange. A point on the left of the picture clearly belongs to the left clump, but nothing in the data tells us which clump generated it — the cluster label is a variable we never recorded. In the language of this volume, it is latent: real, part of the generative story, and hidden. Every observation comes with a missing companion, and the fitting problem has to reason about all the labels we did not see. That is what separates this part from plain maximum likelihood.

The hidden label is not an obstacle to be removed; it is the reason the algorithm is interesting. If we had the labels, fitting would collapse into a bookkeeping exercise: take the points labelled one, compute their mean and covariance, repeat for the others, and you are done in one line. The labels are exactly what we are missing, and the parameters are exactly what we want. Each one would be trivial to compute if we knew the other. That mutual dependence is the whole difficulty, and it is also the whole solution.

The strategy, once you see the shape of the problem, writes itself: guess the hidden labels, use the guess to fit the parameters, use the parameters to improve the guess, and repeat until the guesses stop moving. The surprise — and the reason this idea is everywhere in modern machine learning — is that this circular procedure is not a heuristic. It never makes the fit worse, and the reason why is a single inequality about the logarithm. By the end of this part you will be able to derive that guarantee, name the quantity it protects, and read both off a live plot.

One bell is not enough

The simplest model with hidden structure is the Gaussian mixture. Its generative story has two lines. First draw a hidden label z from a categorical distribution with weights $\pi_1,\dots,\pi_K$, so that component k is chosen with probability $\pi_k$. Then, having chosen the component, draw the observation from that component's Gaussian, $x\sim\mathcal{N}(\mu_k,\Sigma_k)$. Nothing in the data ever records which branch was taken; the label is discarded the moment the point is generated.

Sum over the hidden label and the density of a single observation falls out. It is a weighted average of the component bells, with the mixture weights as the averaging weights. The average of densities is not a density, so different components can overlap and trade mass; that flexibility is what lets a handful of ellipses hug a shape that no single one could.

Fitting means choosing all of it — the weights, the means, the covariances — by maximum likelihood. For the whole dataset, independence makes the likelihood a product over points, which becomes a sum once we take logarithms. And there the trouble starts: each term is the logarithm of a sum of K densities. Differentiate that and the derivative of the outer log brings down a factor whose denominator mixes every component; the parameters no longer decouple. Set the gradient to zero and the equations are coupled, with each component's optimal parameters depending on every other component's parameters through the responsibilities of every point. There is no closed form.

Compare that with the labelled problem, where each z_i is known. The log-likelihood then splits into K independent pieces, a sum over points grouped by label, and each group's mean and covariance are just weighted averages of its members. The unlabelled problem is one missing table of labels away from being trivial. The whole subject of this part is how to fill that table in without ever seeing it.

Two warnings belong here, because they are permanent features rather than numerical accidents. First, the mixture likelihood is not concave: it has many peaks, and EM drives toward whichever one the starting values happen to favour. Second, the components are not identifiable up to permutation. Swap the parameters of components one and two and you have written the same density with relabelled pieces, so the likelihood has a symmetry and the "best" solution is a set of K! equivalent points. Both facts are why the demo below needs a reset button and why real pipelines run the algorithm from several random starts.

Two steps, each easy

Since the labels are missing, replace certainty with belief. For each point x_i and each component k, ask how much that component wants the point, measured as the posterior probability of the label given the point and the current parameters. This is the responsibility, and because it is a posterior it is a ratio of weighted densities that sums to one over k. It is the soft version of the label we never saw, and it is computed by Bayes' rule in one pass over the data.

Now pretend the responsibilities are the labels and do the easy fit. Each component gets a soft count of its members, $N_k=\sum_i\gamma_{ik}$, and uses those counts as weights. The weight is the soft share of the data the component claims, the mean is the responsibility-weighted average of the points, and the covariance is the responsibility-weighted scatter around that mean. Nothing is solved globally; each line is a weighted average, computable in closed form.

Alternate the two lines and you have expectation–maximisation. The E-step holds the parameters fixed and recomputes the responsibilities — the posterior beliefs about the hidden labels. The M-step holds the responsibilities fixed and recomputes the parameters. Start from any sensible initialisation, iterate to a fixed point, and read off the answer. On the canvas, each press of EM step performs one E-step and one M-step: watch the colours of the points sharpen as responsibilities harden toward zero and one, and the ellipses slide and tilt to cover their clusters.

What makes this more than a plausible loop is a theorem: each full iteration leaves the log-likelihood no smaller than it found it. The E-step cannot reduce it, because it changes the responsibilities to match the model. The M-step cannot reduce it either, for a reason worth stating precisely, because the same reason is the guarantee behind the ELBO. In fact the two steps are doing opposite jobs on the same quantity — the E-step closes a gap and the M-step lifts a floor — and the next section makes both statements exact.

The practical caveats follow from the non-concavity. EM converges to a stationary point of the likelihood, not necessarily the global maximum, and it can be slow when the components overlap: the responsibilities stay mushy for many iterations while the ellipses creep. Degenerate solutions also lurk, in which one component collapses onto a single point and its covariance shrinks toward zero while its likelihood shoots to infinity; a small ridge on the covariance, exactly what the demo adds, removes the pathology. And because the order of the components is arbitrary, you should never read meaning into which ellipse is labelled one.

Jensen's inequality does the work

Introduce any distribution q(z) over the hidden label. Multiply and divide inside the log by q(z), so the likelihood becomes an expectation under q of a ratio. Jensen's inequality now applies, because the logarithm is concave: the log of an average is at least the average of the logs. Lowering the log through the expectation produces a quantity that cannot exceed the likelihood, and that is the whole derivation.

The inequality is Jensen's, and the direction is set by the curvature of the log: concavity means the function lies above its chords, so the average of the outputs is never more than the output of the average. A flat-footed but useful sanity check is the one-line statement that for any positive Y, $\log\mathbb{E}[Y]\ge\mathbb{E}[\log Y]$. The gap we have thrown away is exact and has a name. It is the Kullback–Leibler divergence from q to the true posterior, and it is never negative.

Read that identity carefully, because it explains the whole algorithm in one line. For a fixed model, the log-likelihood is a constant, and the split assigns all of it between the ELBO and the KL gap. To make the bound as tight as possible you have only one lever, q: choose the distribution that minimises the divergence. The minimiser is the posterior itself, $q(z)=p(z\mid x)$, and at that choice the divergence is zero and the ELBO equals the likelihood exactly. That is the E-step. It is not a heuristic guess at the labels; it is the exact solution of a well-posed optimisation problem.

The M-step is the mirror image. Fix q — freeze the responsibilities — and the KL term no longer depends on the parameters that matter, so maximising the log-likelihood is the same as maximising the ELBO over $\pi$, $\mu$ and $\Sigma$. The frozen q turns the awkward log of a sum into a sum of logs, and the updates in the previous section are exactly the maximisers. The same quantity is thus optimised in both steps: the E-step tightens the bound from below by killing the gap, and the M-step raises the bound itself. The likelihood, pinned between them, can only climb. That is the monotonicity theorem, and it is why the upper line in the chart never dips.

For the mixture specifically, the ELBO has a compact form: a responsibility-weighted log of weighted density, plus the entropy of the responsibilities. The first term rewards components for explaining points; the entropy term is largest when the responsibilities are spread evenly and vanishes when they harden to zeros and ones. Fitting a mixture is therefore a tug-of-war between covering the data and committing to a partition, and the balance is set automatically by the same bound. The variational view — pick a family of q, maximise the ELBO over both the family and the model — generalises EM beyond the exact posterior, which is the doorway to the variational autoencoder and to the approximate inference that trains it.

One caution about what the bound is and is not. The ELBO is a lower bound on the marginal likelihood of the observed data, so maximising it is a principled substitute for maximum likelihood whenever the marginal is intractable. It does not bound the test-set performance, and because EM climbs to a local maximum of the bound, a better ELBO at convergence just means a better local optimum, not the global one. In the demo, reset and run again and the final value will differ: the data has several peaks, and the starting point chooses which one you reach.

Hidden structure, every time

Every formula in one place

Where to go deeper

• A. P. Dempster, N. M. Laird and D. B. Rubin, Maximum likelihood from incomplete data via the EM algorithm, Journal of the Royal Statistical Society B, 1977 — the paper that named the algorithm and proved the monotone ascent.
• Christopher Bishop, Pattern Recognition and Machine Learning, chapter 9 — the mixture model and EM in the notation this part borrows, with the ELBO derived and then generalised.
• Kevin Murphy, Probabilistic Machine Learning: An Introduction, chapter 11 — a modern treatment that keeps the latent-variable view in front of you throughout.
• Radford Neal and Geoffrey Hinton, A view of the EM algorithm that justifies incremental, sparse, and other variants, 1998 — EM as coordinate ascent on the ELBO, which is the perspective that makes the bound indispensable.
• David Blei, Alp Kucukelbir and Jon McAuliffe, Variational inference: a review for statisticians, JASA, 2017 — what happens when the E-step cannot be solved exactly and you optimise the ELBO anyway.