Learn AI Create the Future — with Pranay
Main site

Clustering in Depth

Expectation-maximisation

EM solves a chicken-and-egg problem — you need group memberships to describe the groups, and group descriptions to assign memberships — by alternating guesses until they stop changing.

Read these first

On this page 5
  1. Why this matters
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

Expectation-maximisation, or EM, fits models whose group labels are hidden.

It alternates two easy steps until the answers settle: guess the memberships, then re-describe the groups.

Imagine a teacher who collected exam papers from two sections, and the pile got shuffled with no names. She knows section A is generally stronger, but not which paper came from where. Chicken and egg: to describe each section's level, she needs to know whose paper is whose. To guess whose paper is whose, she needs each section's level.

Her way out: start with a rough guess of the two levels. Sort papers by which level they fit better — keeping honest part-marks for papers in the middle. Then re-compute each section's level from the papers it received. The new levels re-sort a few borderline papers. Repeat until nothing moves.

That loop is EM.

Why this matters

The shuffled-papers problem is everywhere. Which customer segment produced this purchase? Which speaker said this word? The group label is hidden — the data never came with one. A hidden label like this is called a latent variable: a value the model believes exists but nobody recorded.

EM is the standard trick for fitting any model with latent variables. Fitting a Gaussian mixture model is its most famous job, but the same loop trains hidden Markov models and fills in missing data too.

How it works

guess group descriptions
        |
        v
E-step: give each point part-marks     <- "this paper looks 70% section A"
        for each group
        |
        v
M-step: re-describe each group from    <- "section A's new average: 78"
        the points it part-owns
        |
        v
repeat until nothing changes

The two names: the E-step (expectation) computes each point's membership shares under the current descriptions. The M-step (maximisation) refits each group's description, with every point counted in proportion to its share.

One guarantee makes EM trustworthy: every loop makes the model's overall fit to the data better, or leaves it equal. It never gets worse. It can, however, settle at a locally-best answer rather than the best possible one. That is the same trap as K-means starts. The same medicine applies: restart from several places.

A real example you have seen

Photo apps that group faces run this loop at heart. Guess the "typical look" of each person, softly assign each photo, re-compute each typical look, repeat. Nobody ever labelled the photos — the person identities are latent variables.

Remember this

  • EM fits models where the group label is hidden.
  • E-step shares points out; M-step re-describes groups. Alternate until stable.
  • Fit never worsens per loop — but restarts are still needed to dodge bad valleys.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install numpy

No scikit-learn this time — the whole point is to watch the loop with nothing hidden. Outputs verified with numpy 1.26.4; expect last-digit drift elsewhere.

EM for two height groups, by hand

Heights of 400 people: 200 drawn around 160 cm, 200 around 175 cm, shuffled together. EM must recover the two bumps from deliberately wrong starting guesses.

em_by_hand.py
import numpy as np

rng = np.random.default_rng(1)
heights = np.concatenate([rng.normal(160, 6, 200), rng.normal(175, 6, 200)])

mu = np.array([150.0, 172.0])          # wrong on purpose
sigma = np.array([4.0, 4.0])
weight = np.array([0.5, 0.5])

def bell(x, m, s):
    return np.exp(-0.5 * ((x - m) / s) ** 2) / (s * np.sqrt(2 * np.pi))

for step in range(1, 41):
    # E-step: how much does each bump claim each person?
    dens = weight * bell(heights[:, None], mu, sigma)
    avg_loglik = np.log(dens.sum(axis=1)).mean()
    r = dens / dens.sum(axis=1, keepdims=True)
    # M-step: refit each bump to the people it claimed
    n = r.sum(axis=0)
    mu = (r * heights[:, None]).sum(axis=0) / n
    sigma = np.sqrt((r * (heights[:, None] - mu) ** 2).sum(axis=0) / n)
    weight = n / len(heights)
    if step in (1, 2, 5, 10, 20, 40):
        print(f"step {step:2d}: means = {np.round(mu, 1)}, avg log-lik = {avg_loglik:.3f}")
Output
step  1: means = [156.2 171.6], avg log-lik = -4.384
step  2: means = [156.8 171.4], avg log-lik = -3.637
step  5: means = [157.5 171.5], avg log-lik = -3.619
step 10: means = [157.9 171.9], avg log-lik = -3.617
step 20: means = [158.5 172.7], avg log-lik = -3.614
step 40: means = [159.3 174. ], avg log-lik = -3.610

The walkthrough

r is the E-step's whole output — a 400×2 table of membership shares, one row per person, rows summing to 1. A 163 cm person might split 0.7 / 0.3. These shares are called responsibilities.

The M-step is weighted averaging, nothing fancier. Each bump's new mean is the average of all 400 heights, weighted by that bump's responsibility column. Same trick for the spread and the group share. Soft membership means every point teaches every group, in proportion.

Watch the log-likelihood column: it only ever rises. From -4.384 toward -3.610, never a single decrease. That monotone climb is EM's signature guarantee, and checking it is the best debugging tool you have. A drop means a bug, every time.

Notice the crawl. Means move from (156, 172) toward the truth of (160, 175), and step 40 is still travelling. When two bumps overlap, responsibilities near the middle stay soft, and EM converges slowly. This is normal, not broken. Library implementations stop on a likelihood-change tolerance, not a step count.

Common mistakes

Stopping when the means look stable at one decimal. Steps 2 through 10 look nearly frozen while the likelihood still climbs. Track the likelihood, not the parameters.

Initialising both groups identically. Identical bumps earn identical responsibilities, and the M-step then rebuilds identical bumps — a perfect tie that never breaks. Starts must differ. Libraries break the symmetry with K-means or random responsibilities.

Letting a bump shrink onto one point. If a bump's spread heads toward zero, its likelihood explodes and the arithmetic overflows. Real implementations put a floor under the variance — that is exactly reg_covar in scikit-learn's GaussianMixture.

Computing responsibilities without logs at scale. Products of many tiny densities underflow float64. The standard fix is the log-sum-exp trick; hand-rolled EMs on real data should work in log space throughout.

Try it yourself

Start with mu = np.array([164.0, 168.0]) and sigma = np.array([15.0, 15.0]). Watch the means barely separate even after 40 steps — wide, nearly-identical starts produce nearly-uniform responsibilities. Then explain why the likelihood still never falls.

What to learn next

Researcher — Mathematics and papers.

The general formulation

Observed data $x$, latent variables $z$, parameters $\theta$. The target is the marginal log-likelihood $\log p_\theta(x) = \log \sum_z p_\theta(x, z)$, awkward because the sum sits inside the log. For any distribution $q(z)$:

$$ \log p_\theta(x) = \underbrace{\mathbb{E}{q(z)}!\left[\log \frac{p\theta(x, z)}{q(z)}\right]}{\mathcal{L}(q, \theta)} + \mathrm{KL}!\left(q(z) \,|\, p\theta(z \mid x)\right) $$

Where:

  • $\mathcal{L}(q, \theta)$ — the evidence lower bound (ELBO), since KL divergence is non-negative.
  • $\mathrm{KL}$ — Kullback–Leibler divergence, zero only when its arguments match.

E-step: set $q(z) = p_{\theta^{(t)}}(z \mid x)$, the current posterior — this closes the KL gap, making the bound tight. For a GMM this is the responsibility computation. M-step: maximise $\mathcal{L}$ over $\theta$ with $q$ fixed, which is a complete-data likelihood — usually closed form.

Monotonicity follows in two lines: the E-step lifts the bound to touch the objective; the M-step raises the bound; the objective sits above the bound. Hence $\log p_{\theta^{(t+1)}}(x) \ge \log p_{\theta^{(t)}}(x)$. The canonical reference is Dempster, Laird and Rubin (1977), Maximum likelihood from incomplete data via the EM algorithm; Wu (1983) supplies the convergence conditions the original paper got wrong.

Convergence behaviour

EM converges to a stationary point of the likelihood — a local maximum or saddle. The rate is linear, with factor governed by the fraction of missing information: when components overlap heavily, responsibilities carry little information and convergence crawls (the exact behaviour visible in the developer block). Near-separated components give near-Newton speed.

Salakhutdinov, Roweis and Ghahramani (2003) analyse when EM's slowness bites and when direct gradient ascent on the likelihood wins. Modern practice often replaces the M-step with a gradient step (generalised EM), which retains monotonicity provided the step improves the bound.

The wider family

  • Hard EM / classification EM: replace responsibilities with their argmax. For spherical equal-weight Gaussians this is K-means (Celeux and Govaert, 1992).
  • Variational EM: when the posterior $p(z \mid x)$ is intractable, restrict $q$ to a tractable family and accept a loose bound — the root of variational inference, and the ancestor of the VAE objective.
  • Stochastic EM / online EM (Cappé and Moulines, 2009): E-steps on minibatches with running sufficient statistics, for streaming-scale data.
  • MAP-EM: add priors; the M-step becomes penalised maximum likelihood — one principled route around collapsing components.
  • Baum–Welch is EM specialised to hidden Markov models — the E-step is the forward-backward algorithm.

Complexity per iteration for a $K$-component GMM in $d$ dimensions over $n$ points: $O(nKd^2)$ with full covariances, dominated by Mahalanobis terms; $O(nKd)$ diagonal.

What to learn next