The EM algorithm

What this chapter covers

k-means made a hard choice for every point: you belong to this cluster, full stop. That worked because the blobs were far apart and the choice was obvious. But most real groups overlap, and near the seam between two of them a hard choice is a guess dressed up as a fact. This chapter is about the algorithm that refuses to guess. It hands every point a probability of belonging to each group, fits the groups using those probabilities as weights, and repeats — softening every decision until the numbers stop moving.

That algorithm is Expectation-Maximization, EM, and it's one of the load-bearing ideas in all of statistics. The setup is always the same: you have data, and some piece of information you'd need to fit it cleanly is missing. Here the missing piece is which of two Gaussians produced each number. If you knew that, fitting the two Gaussians would be a one-liner. If you knew the two Gaussians, guessing the labels would be a one-liner. You know neither, so EM alternates: guess the labels softly from the current Gaussians (the E-step), refit the Gaussians from those soft labels (the M-step), and go around again. We build the whole loop by hand in NumPy on the cleanest possible case — a single column of numbers drawn from two overlapping bells — then check it against scikit-learn.

The thing to keep your eye on is the log-likelihood, one number that scores how well the current parameters explain the data. EM's central guarantee is that this number never goes down. Not "usually goes up" — never goes down, every single iteration, provably. That monotonic climb is the whole reason the method is trusted, and the animation makes you watch it happen. This chapter is about the general recipe; the fuller, 2-D story where each component is a tilted covariance ellipse gets its own chapter on Gaussian mixture models.

A bit of history

The name and the general form come from a single famous paper: Arthur Dempster, Nan Laird, and Donald Rubin, "Maximum Likelihood from Incomplete Data via the EM Algorithm," read to the Royal Statistical Society in 1977. They didn't invent the individual tricks — special cases had been floating around for years, from H.O. Hartley's 1958 work on counts to the Baum-Welch procedure for hidden Markov models in the late 1960s. What Dempster, Laird, and Rubin did was see that all of these were the same algorithm. Any problem you could phrase as "the data I have is an incomplete version of data I could fit easily" fit one master recipe, and they gave it a name, a general statement, and the argument for why it works.

That argument had a hole in it. The 1977 paper claimed the log-likelihood converged in a way that was slightly too strong, and in 1983 C.F. Jeff Wu published the correction, pinning down exactly what EM does and doesn't promise — it climbs to a stationary point of the likelihood, which need not be the global maximum. That caveat is not a footnote; it's the thing you feel every time you run EM in anger, and we come back to it. Nearly fifty years on, the paper is one of the most cited in the history of statistics, because the pattern it named is everywhere: mixtures, missing survey answers, hidden states, latent topics. Learn EM once and you keep recognizing it.

The intuition

Picture the numbers on a line. Somewhere on the left there's a pile, somewhere on the right there's a bigger pile, and in the middle they blur together so you can't say where one ends and the other starts. You believe two processes made this data, but the label that would tell you which process made each point was never recorded. This is the missing information, and it's the whole difficulty.

Here's the trick that dissolves it. Suppose someone whispered the true Gaussians in your ear — left one centered at -2, right one at 3, these spreads, this mix. Then for any point you could compute how likely each Gaussian was to have produced it, and the ratio gives you a soft label: 90% left, 10% right, say. Now suppose instead someone handed you those soft labels for every point. Then fitting the Gaussians is just a weighted average — each point contributes to the left Gaussian's mean in proportion to how much it looks left. Either half is easy given the other. You have neither, so you bootstrap: start with a guess at the Gaussians, however bad, compute the soft labels they imply, refit the Gaussians from those labels, and loop.

That loop is EM. The E-step is "given the current parameters, how much does each point belong to each component." The M-step is "given those soft memberships, what's the best-fit component." Each pass sharpens the other, and the beautiful part is that a single quantity — the log-likelihood of the data — goes up every time, so the process can't thrash or wander. Here's the data it's going to work on, as a density histogram, two bells smeared into one lumpy shape:

You can almost see two humps — a shorter one near -2, a broader one near 3 — but the middle is genuinely ambiguous, and no hard cutoff there is honest. EM's job is to pull the two bells back out of that lump.

The math

Write the data as $x_1, \dots, x_N$, one number each. We model it as a mixture of $K$ Gaussians (here $K = 2$), where component $k$ has mixing weight $\pi_k$ (with $\sum_k \pi_k = 1$), mean $\mu_k$, and variance $\sigma_k^2$. Collect all of these into $\theta$. The probability the model assigns to a point is the weighted sum over components:

The missing information is a hidden label $z_i \in \{1, \dots, K\}$ per point saying which component drew it. The E-step computes its posterior — the responsibility $\gamma_{ik}$, the probability that component $k$ produced point $i$ given the current parameters. It's Bayes' rule: the component's weighted density over the total.

Each row $\gamma_{i\cdot}$ sums to 1 — a full soft assignment. The M-step then treats these responsibilities as fractional memberships and refits each component by weighted averaging. Let $N_k = \sum_i \gamma_{ik}$ be the soft count of points in component $k$:

These are the ordinary maximum-likelihood formulas for a weight, a mean, and a variance, with $\gamma_{ik}$ standing in for "how much of point $i$ counts toward component $k$." The quantity both steps are quietly improving is the incomplete-data log-likelihood — the log-probability of the numbers alone, with the hidden labels summed out:

Why the climb is guaranteed takes one more idea. For any choice of soft labels $q(z_i)$, the log-likelihood splits exactly into a lower bound plus a gap:

The first term $\mathcal{L}(q, \theta)$ is the ELBO, the evidence lower bound; the second is a KL divergence, which is never negative. The E-step sets $q$ to the true posterior, which zeroes the KL gap and makes the bound tight. The M-step maximizes the bound over $\theta$. Raising a tight lower bound can only raise what sits above it, so:

That inequality is the entire promise of EM. It says nothing about reaching the best answer — only that you never get worse. Which is exactly the double-edged thing we build next.

What it's good at, what it isn't

EM's gift is that it makes intractable likelihoods tractable. Maximizing $\ell(\theta)$ directly means differentiating a log-of-a-sum, and the sum inside the log tangles all the components together into something with no closed-form solution. EM sidesteps that entirely: the E-step and M-step each have clean, closed forms, and alternating them optimizes the hard objective without ever touching it head-on. You get soft assignments for free, which is the honest output when groups overlap, and the monotonic climb means the loop is stable — no learning rate to tune, no step that can blow up. When your problem genuinely has hidden structure or missing values, EM is very often the first thing that works.

The catch is the one Wu had to fix in the proof: EM finds a local optimum, not the global one, and which local optimum depends entirely on where you start. A bad init can strand it in a bad fit that is perfectly stable — the log-likelihood plateaus, the algorithm reports success, and the answer is wrong. It can also do genuinely ugly things a mean never does: let a component collapse onto a single point, driving its variance toward zero and the likelihood toward infinity, a degenerate solution you have to guard against. And it can be slow, crawling toward the optimum in tiny steps once it's close. The standard defenses are all about the start: run it from several random inits and keep the best, or seed it with k-means, which is exactly what scikit-learn does.

The data

One column of 300 numbers. I built it by sampling from two Gaussians and throwing away the labels: 40% of the points come from $\mathcal{N}(-2, 1)$, the other 60% from $\mathcal{N}(3, 1.5)$ — different centers, different spreads, unequal mixing, and enough overlap in the middle that no clean line separates them. The generator records which Gaussian each point came from in a comp column, but that's a sealed envelope: EM sees only the numbers. We open the envelope once at the end, to grade how well the recovered soft clustering matches the truth. Two components in one dimension is the smallest mixture that still shows everything EM does, which is why it's the right place to watch the mechanism instead of the scenery.

Build it, one function at a time

Five functions. One Gaussian, the two steps, the score they optimize, and the loop that ties them together — that's the entire algorithm. Everything starts with the density of a single 1-D Gaussian, evaluated across all the points at once:

def gaussian(x, mu, var):
    """1-D Gaussian density N(x | mu, var), evaluated elementwise over x.

    var is the variance (sigma squared), not the standard deviation. This is
    the single building block the whole algorithm leans on — every step below
    is just weighted sums of this number.
    """
    return np.exp(-0.5 * (x - mu) ** 2 / var) / np.sqrt(2.0 * np.pi * var)

Note it takes the variance, not the standard deviation — that's the parameter the M-step actually updates, so carrying it directly keeps the code honest. On top of that one function sits the E-step. For every point, weight each component's density by its mixing weight and normalize across components; the result is the responsibility matrix, one row per point summing to 1:

def e_step(x, weights, means, variances):
    """E-step: responsibilities gamma[i, k] = P(component k produced x_i).

    For each point we weight every component's density by its mixing weight and
    normalize across components, so each row sums to 1. gamma[i, k] near 1 means
    point i almost certainly belongs to component k; 0.5 means it's on the fence.
    Returns an (N, K) matrix — the soft assignment the M-step averages over.
    """
    weighted = np.array([w * gaussian(x, m, v)
                         for w, m, v in zip(weights, means, variances)]).T  # (N, K)
    return weighted / weighted.sum(axis=1, keepdims=True)

That's Bayes' rule in three lines. A point sitting under the left bell gets most of its mass on component 0; a point in the murky middle splits close to 50/50. The M-step takes those soft memberships and refits each component by weighted averaging:

def m_step(x, gamma):
    """M-step: re-estimate weights, means, variances as weighted averages.

    N_k is the soft count of points in component k — the column sum of the
    responsibilities. Everything is the ordinary formula for a weighted mean and
    weighted variance, with gamma[:, k] as the weights: a point that half-belongs
    to a component contributes half of itself to that component's estimates.
    """
    Nk = gamma.sum(axis=0)                                   # (K,) soft counts
    weights = Nk / len(x)
    means = (gamma * x[:, None]).sum(axis=0) / Nk
    variances = (gamma * (x[:, None] - means) ** 2).sum(axis=0) / Nk
    return weights, means, np.maximum(variances, 1e-6)       # floor guards /0

Nk is the soft count — the number of points each component owns, fractional because ownership is shared. The weight is that count over N; the mean and variance are the responsibility-weighted mean and variance. The one guard is a floor on the variance, so a component can't collapse to a spike and send the likelihood to infinity. Now the score both steps are climbing, the incomplete-data log-likelihood:

def log_likelihood(x, weights, means, variances):
    """Incomplete-data log-likelihood: sum_i log sum_k pi_k N(x_i | mu_k, var_k).

    "Incomplete" because it's the likelihood of the x's alone, with the hidden
    component labels summed out. This is the single number EM optimizes, and the
    one guarantee is that it never goes down from one iteration to the next.
    """
    per = np.array([w * gaussian(x, m, v)
                    for w, m, v in zip(weights, means, variances)]).T   # (N, K)
    return float(np.log(per.sum(axis=1)).sum())

Straight from the math: for each point, the log of its total mixture density; sum over points. This is the number that must never fall. Watching it is how you know the loop is working and how you catch a bug — if it ever drops, something is wrong in the E-step or M-step. Last, the loop itself:

def em(x, init_weights, init_means, init_variances, max_iter=200, tol=1e-6):
    """EM loop: alternate E-step and M-step until the log-likelihood plateaus.

    Starts from whatever parameters you hand it — EM only finds a local optimum,
    so the init matters. Each pass computes responsibilities (E), re-estimates the
    components from them (M), and checks the log-likelihood; when it stops rising
    by more than tol, we've converged. Returns the fitted parameters, the final
    log-likelihood, and the full per-iteration history the chapter replays.
    """
    weights = np.array(init_weights, float)
    means = np.array(init_means, float)
    variances = np.array(init_variances, float)

    ll = log_likelihood(x, weights, means, variances)
    history = [_snapshot(x, weights, means, variances, ll)]
    for _ in range(max_iter):
        gamma = e_step(x, weights, means, variances)            # E-step
        weights, means, variances = m_step(x, gamma)            # M-step
        new_ll = log_likelihood(x, weights, means, variances)
        history.append(_snapshot(x, weights, means, variances, new_ll))
        if new_ll - ll < tol:                                   # plateaued
            ll = new_ll
            break
        ll = new_ll
    return weights, means, variances, ll, history

E-step, M-step, check the score, stop when it stops rising by more than a hair. It records a snapshot every iteration — the parameters, the log-likelihood, and each point's current responsibility — which is what the animation replays. Convergence is just the log-likelihood going flat, because a flat monotone sequence has nowhere left to go.

Watch it work

This is the payoff for using one dimension. Below is a real run of the em function above, one frame per iteration, and I started it from a deliberately terrible init on purpose: both component means bunched near the center at -0.5 and 0.5, equal weights, unit variance. From there you get to watch EM pry the two Gaussians apart.

Three panels, all live. The top overlays the two current component curves — each Gaussian scaled by its mixing weight — on the fixed data histogram; watch the cyan curve slide left and the orange one slide right and fatten. The middle strip is the 300 points, each colored by its current responsibility: pure cyan means "certainly the left component," pure orange means "certainly the right," and a muddy blend in between means the point is still on the fence. Watch the middle harden from mush into two committed colors. The bottom panel traces the log-likelihood, the one number that only ever climbs.

Press play and the first frame does almost all the work. The log-likelihood leaps from -1,386.5 at the bad init to -703.9 after a single E/M pass — the two means snap apart from their bunched start toward -2 and 3 in one step. From there it's refinement: -689.7, then -682.8, then a long slow creep as the curves settle into the histogram and the fence-sitters in the middle pick a side. By iteration 8 the score is already -681.5 and the picture has stopped visibly changing; the loop runs to 20 iterations only because it keeps squeezing out improvements smaller than you can see, and stops when they fall below the tolerance. It converges to a log-likelihood of -681.48.

Now look at the bottom panel across the whole run and find a single frame where the line goes down. There isn't one. That's not luck of this seed — it's the guarantee from the math made visible, and it's the reason you can trust EM to be climbing even when, as here, it starts from a fit that's frankly ridiculous. Here's that same climb on its own, every iteration:

A cliff at the first step, then a knee, then flat. Monotone the whole way. That shape — huge early gains, then a long tail of diminishing returns — is the signature of EM, and it's why people cap the iterations and set a tolerance rather than wait for the numbers to stop moving completely.

The full implementation

The whole file, no library, top to bottom. This is exactly what the animation ran:

"""Expectation-Maximization for a 1-D Gaussian mixture, built from scratch.

The data is one column of numbers that came from two overlapping Gaussians, and
nobody wrote down which point came from which. EM recovers the two Gaussians —
their means, their spreads, and how much of the data each one produced — by
alternating two steps: an E-step that computes, for every point, how strongly it
belongs to each component (soft assignment), and an M-step that re-estimates the
components as responsibility-weighted averages. Pure NumPy — no ML library in
this file.

Every function below appears in the chapter one step at a time (the
`# region:` markers are what the book's include directives pull in).
"""

import numpy as np
import pandas as pd


# region: gaussian
def gaussian(x, mu, var):
    """1-D Gaussian density N(x | mu, var), evaluated elementwise over x.

    var is the variance (sigma squared), not the standard deviation. This is
    the single building block the whole algorithm leans on — every step below
    is just weighted sums of this number.
    """
    return np.exp(-0.5 * (x - mu) ** 2 / var) / np.sqrt(2.0 * np.pi * var)
# endregion


# region: e_step
def e_step(x, weights, means, variances):
    """E-step: responsibilities gamma[i, k] = P(component k produced x_i).

    For each point we weight every component's density by its mixing weight and
    normalize across components, so each row sums to 1. gamma[i, k] near 1 means
    point i almost certainly belongs to component k; 0.5 means it's on the fence.
    Returns an (N, K) matrix — the soft assignment the M-step averages over.
    """
    weighted = np.array([w * gaussian(x, m, v)
                         for w, m, v in zip(weights, means, variances)]).T  # (N, K)
    return weighted / weighted.sum(axis=1, keepdims=True)
# endregion


# region: m_step
def m_step(x, gamma):
    """M-step: re-estimate weights, means, variances as weighted averages.

    N_k is the soft count of points in component k — the column sum of the
    responsibilities. Everything is the ordinary formula for a weighted mean and
    weighted variance, with gamma[:, k] as the weights: a point that half-belongs
    to a component contributes half of itself to that component's estimates.
    """
    Nk = gamma.sum(axis=0)                                   # (K,) soft counts
    weights = Nk / len(x)
    means = (gamma * x[:, None]).sum(axis=0) / Nk
    variances = (gamma * (x[:, None] - means) ** 2).sum(axis=0) / Nk
    return weights, means, np.maximum(variances, 1e-6)       # floor guards /0
# endregion


# region: log_likelihood
def log_likelihood(x, weights, means, variances):
    """Incomplete-data log-likelihood: sum_i log sum_k pi_k N(x_i | mu_k, var_k).

    "Incomplete" because it's the likelihood of the x's alone, with the hidden
    component labels summed out. This is the single number EM optimizes, and the
    one guarantee is that it never goes down from one iteration to the next.
    """
    per = np.array([w * gaussian(x, m, v)
                    for w, m, v in zip(weights, means, variances)]).T   # (N, K)
    return float(np.log(per.sum(axis=1)).sum())
# endregion


# region: em
def em(x, init_weights, init_means, init_variances, max_iter=200, tol=1e-6):
    """EM loop: alternate E-step and M-step until the log-likelihood plateaus.

    Starts from whatever parameters you hand it — EM only finds a local optimum,
    so the init matters. Each pass computes responsibilities (E), re-estimates the
    components from them (M), and checks the log-likelihood; when it stops rising
    by more than tol, we've converged. Returns the fitted parameters, the final
    log-likelihood, and the full per-iteration history the chapter replays.
    """
    weights = np.array(init_weights, float)
    means = np.array(init_means, float)
    variances = np.array(init_variances, float)

    ll = log_likelihood(x, weights, means, variances)
    history = [_snapshot(x, weights, means, variances, ll)]
    for _ in range(max_iter):
        gamma = e_step(x, weights, means, variances)            # E-step
        weights, means, variances = m_step(x, gamma)            # M-step
        new_ll = log_likelihood(x, weights, means, variances)
        history.append(_snapshot(x, weights, means, variances, new_ll))
        if new_ll - ll < tol:                                   # plateaued
            ll = new_ll
            break
        ll = new_ll
    return weights, means, variances, ll, history
# endregion


def _snapshot(x, weights, means, variances, ll):
    """One frame of the run: the current parameters, the incomplete-data
    log-likelihood, and each point's responsibility to component 1 (the value
    the animation colors points by so you watch the soft assignment harden)."""
    resp1 = e_step(x, weights, means, variances)[:, 1]
    return {
        "weights": weights.copy(),
        "means": means.copy(),
        "variances": variances.copy(),
        "ll": ll,
        "resp1": resp1.copy(),
    }


def load_data(path="../data/mixture.csv"):
    """The 1-D mixture snapshot: one column x of numbers, plus the true source
    component of each. EM sees only x; the component column is ground truth held
    back to grade the recovered clustering at the end."""
    return pd.read_csv(path)

The library version

Nobody hand-rolls EM for a Gaussian mixture in production, and once you've built it you don't need to. scikit-learn's GaussianMixture is the same E-step/M-step loop with the sharp edges filed off — a k-means initialization instead of a random one, several restarts so a single bad seed can't decide the answer, and covariance handling that generalizes past our single variance. On 1-D data with two components it optimizes the identical log-likelihood, so it should land where we did:

def sklearn_gmm(x, k=2, seed=0):
    """Fit a k-component 1-D Gaussian mixture with scikit-learn's EM.

    Returns the same four things our em() returns: mixing weights, means,
    variances, and the total incomplete-data log-likelihood. score() reports the
    mean log-likelihood per sample, so we multiply by N to match our total.
    """
    X = np.asarray(x, float).reshape(-1, 1)
    gm = GaussianMixture(n_components=k, covariance_type="full",
                         max_iter=200, tol=1e-6, random_state=seed)
    gm.fit(X)
    weights = gm.weights_
    means = gm.means_.ravel()
    variances = gm.covariances_.ravel()
    ll = float(gm.score(X) * len(X))
    return weights, means, variances, ll

The one wrinkle is bookkeeping: score returns the mean log-likelihood per point, so we multiply by N to recover the total our em reports. Everything else lines up. Here are both fits laid over the data — the two fitted component curves, scaled by their weights, on the same histogram from the top:

The two bells sit right where your eye guessed they were, and the mixture of them traces the lumpy histogram. The recovered parameters are close to the truth we built with: means -1.92 and 3.09 against the true -2 and 3, weights 0.43 and 0.57 against the true 0.40 and 0.60. Not exact, because 300 samples is a finite draw, but that's the fit, not the algorithm — it found the best Gaussians for this particular data.

Scratch versus library

Same data, same two components, both run to convergence — our from-scratch EM against sklearn's GaussianMixture. Every fitted parameter, side by side:

The pairs are indistinguishable. Both reach an incomplete-data log-likelihood of -681.48; the weights agree to three decimals (0.431 and 0.569), the means to two (-1.92 and 3.09), the variances to two (1.13 and 2.03). That's not a coincidence and it's not a tie I rounded into being — on a clean two-Gaussian mixture there's one dominant optimum, and any correct EM that climbs to it lands on the same parameters, whether it started from k-means like sklearn or from my deliberately bad init. The one bookkeeping note: EM is free to call the left Gaussian "component 0" or "component 1," so I sort both fits by mean before comparing. Line them up and they're the same model.

Now open the sealed envelope. We never fit on the true component labels, but we can grade against them: take each point's hard assignment (whichever component has the higher responsibility) and score it with the adjusted Rand index, which measures agreement between two labelings after correcting for chance. Both our clustering and sklearn's score 0.87 — not a perfect 1.0, and they shouldn't, because the two Gaussians genuinely overlap and the points in the middle really are ambiguous. EM recovered the structure about as well as it can be recovered; the 13% it "misses" is honest uncertainty the data itself contains, not a failure of the fit.

Takeaways

EM is the tool you reach for when the thing that would make your problem easy is the thing you don't have. Missing labels, missing measurements, a hidden state, a latent cause — phrase the gap as "incomplete data" and EM gives you a loop that fills it in softly and refits, over and over, with a guarantee that you're never moving backward. That guarantee is the whole reason to know it cold: no learning rate, no divergence, just a monotone climb to a stationary point. On this toy mixture the climb was dramatic and the ending was right, but keep Wu's correction in your pocket — the ending is a local optimum, and on a harder problem a bad start lands you in a bad one that looks just as converged. Run it from several inits, or seed it with k-means, and keep the best log-likelihood. That habit is not optional.

The reason EM earns a chapter of its own, apart from any one model, is that it underlies a whole shelf of methods you'll meet again. Fit this exact recipe in two or more dimensions with full covariance matrices and you get Gaussian mixture models, where each component is a tilted, stretched ellipse and k-means turns out to be the hard-assignment limit of the same loop. Swap the Gaussians for multinomials over words and you get topic models, EM discovering themes in documents nobody labeled. Run it over sequences with a transition structure and it's Baum-Welch, the training algorithm for hidden Markov models. They look like different algorithms until you notice they're all the same E-step and M-step, climbing the same kind of bound. Build the loop by hand once, on two bells in one dimension, and you've built the engine for all of them.