The EM Algorithm: GMMs, the k-Means Connection, and Why You Can Train Hidden Variable Models

The EM algorithm answers a question that seems impossible: how do you do MLE when you have missing data or unobserved latent variables? It turns out the answer is elegant — alternate between filling in the missing variables (E-step) and maximising the likelihood (M-step) — and it is why we can train Gaussian Mixture Models, HMMs, and topic models at all.

Maximum likelihood estimation is straightforward when all variables are observed: write down the log-likelihood, take the gradient, set it to zero. But many powerful models have latent variables — variables that are never observed. Gaussian Mixture Models have cluster assignments. Hidden Markov Models have hidden state sequences. Topic models (LDA) have topic assignments per word. You cannot take the gradient of the log-likelihood directly because summing over all possible latent configurations is intractable. The EM algorithm solves this.

The core idea: complete data vs incomplete data

Define complete data as (x, z) where x is observed and z is the latent variable. If we knew z, the log-likelihood log P(x, z | θ) — the complete-data log-likelihood — might be easy to maximise. We do not know z. The EM algorithm iterates: E-step: compute Q(θ | θ_old) = E_{z|x,θ_old}[log P(x, z | θ)] — the expected complete-data log-likelihood, where the expectation is over the posterior of z given current parameters. M-step: θ_new = argmax_θ Q(θ | θ_old) — maximise the expected complete-data log-likelihood. Repeat until convergence.

Why EM is guaranteed to converge

EM converges because each iteration non-decreasingly improves the observed-data log-likelihood log P(x | θ). The proof uses Jensen's inequality applied to the log-concavity of the expectation: log P(x|θ) ≥ Q(θ|θ_old) + H(posterior). The M-step increases Q, which means the bound increases, which means log P(x|θ) cannot decrease. Convergence is to a local maximum or saddle point — EM does not guarantee the global MLE. Initialisation matters enormously.

Gaussian Mixture Models (GMMs)

Model: a mixture of K Gaussians. The generative process: (1) sample cluster assignment z ~ Categorical(π₁, ..., πK). (2) Sample x ~ N(μ_z, Σ_z). Observed: x. Latent: z. E-step: compute soft cluster assignments (responsibilities): r_{nk} = P(z_n = k | x_n, θ) = π_k N(x_n; μ_k, Σ_k) / Σ_j π_j N(x_n; μ_j, Σ_j). Each point x_n gets a soft assignment to each cluster — the r_{nk} sum to 1 over k. M-step: update parameters using the weighted data. N_k = Σ_n r_{nk} (effective number of points in cluster k). π_k = N_k / n. μ_k = Σ_n r_{nk} x_n / N_k (weighted mean). Σ_k = Σ_n r_{nk} (x_n - μ_k)(x_n - μ_k)ᵀ / N_k (weighted covariance). These updates are closed-form — this is the M-step maximisation in action.

k-Means as hard-assignment EM

k-Means is EM on a GMM with hard assignment (no soft probabilities) and spherical, equal-variance, equal-weight Gaussians. E-step → assignment step: each point is assigned to its nearest centroid (argmax_k r_{nk} = 1, all others 0). M-step → update step: centroids are updated as the mean of assigned points. k-Means minimises Σ_n min_k ||x_n - μ_k||² — the total within-cluster sum of squares. This is not the log-likelihood of any natural generative model; it is a hard approximation that works well empirically when clusters are roughly spherical and equal-sized.

Hidden Markov Models (HMMs)

A sequence model: hidden states z₁, z₂, ..., zT with Markov transitions P(zₜ|zₜ₋₁ = A (transition matrix), observations xₜ ~ P(xₜ|zₜ) = B (emission model). We observe x₁:T; we do not observe z₁:T. E-step — forward-backward algorithm: efficiently computes the posterior over hidden states γₜ(k) = P(zₜ = k | x₁:T, θ) and state transition posteriors ξₜ(j,k) = P(zₜ₋₁=j, zₜ=k | x₁:T, θ). The forward-backward algorithm exploits the Markov structure to do this in O(TK²) instead of O(KT). M-step — Baum-Welch: update A, B, π using the posterior counts, exactly as in GMMs. The Baum-Welch algorithm IS EM for HMMs.

Latent Dirichlet Allocation (LDA)

LDA is a topic model: documents are mixtures of topics, topics are distributions over words. Latent variable: topic assignments for each word. Because the Dirichlet-Multinomial conjugacy makes the full posterior tractable, LDA can be trained with variational EM (approximate E-step using a variational distribution) or collapsed Gibbs sampling. The core structure — latent variables, E-step to compute posteriors, M-step to update parameters — is still EM.

Practical considerations

Initialisation: k-Means++ initialisation (choose initial centroids with probability proportional to distance from existing centroids) dramatically improves convergence. For GMMs, initialise with k-Means output. Multiple restarts: run EM from many random starting points and pick the highest likelihood solution. Convergence criterion: stop when log-likelihood improvement drops below a threshold (e.g., 1e-6). Degenerate solutions: a Gaussian can collapse on a single data point with σ → 0 and likelihood → ∞. Fix by adding a minimum covariance floor.

Interview questions on this topic

"Explain the EM algorithm in one minute." — We have observed data x and latent variables z. Direct MLE is hard because we'd need to marginalise over all z. EM alternates: E-step computes the expected log-likelihood under the current posterior over z; M-step maximises that expectation. Each iteration is guaranteed not to decrease the observed log-likelihood. Converges to a local maximum.

"How is k-Means a special case of EM for GMMs?" — k-Means uses hard cluster assignments (each point belongs to exactly one cluster) and assumes equal, spherical Gaussians. GMMs use soft assignments (each point belongs partially to each cluster) and learn cluster shapes. In the limit of zero variance, the soft GMM assignments approach hard k-Means assignments.

"Why does EM not always find the global maximum? What do you do about it?" — EM climbs the likelihood from its starting point and converges to a local max or saddle. With multiple modes, different initialisations converge to different solutions. In practice: multiple restarts, choosing the solution with the highest final likelihood. For k-Means specifically, k-Means++ gives a better-than-random initialisation with provable 8 ln k approximation guarantee.

"When would you use GMMs vs k-Means in practice?" — GMMs: when you need soft cluster assignments (uncertainty about which cluster a point belongs to), when clusters are not spherical (full covariance GMM), when you want a proper generative model (can sample from it, compute likelihoods). k-Means: faster, when hard assignments suffice, when clusters are roughly spherical and balanced in size.

Try on Colab: implement GMM from scratch using the EM algorithm on the Old Faithful eruption dataset (2D). Initialise with k=2 components. Plot the decision boundaries and ellipses at each EM step to visualise convergence. Compare log-likelihood trajectories for 10 random starts. Show that k-Means initialisation (warm start) converges faster and more reliably than random initialisation.