EM Algorithm
Latent variables, E-step/M-step, GMM, convergence
You have 1000 customer purchase records but no segment labels. You believe there are K=3 segments: high-value, medium-value, and occasional. To fit a Gaussian Mixture Model, you need to know which segment each customer belongs to in order to compute per-segment means and variances. But you cannot know the segments until you have the parameters. Classic chicken-and-egg: you need labels to fit parameters and parameters to assign labels.
EM breaks the deadlock by replacing hard segment assignments with soft ones — probabilities. The E-step (Expectation): given current parameters θ^(t), compute P(segment k | customer i) for every i and k. These are soft memberships — each customer is distributed across all segments with weights summing to 1. The M-step (Maximization): given soft memberships, update parameters θ^(t+1) using weighted statistics. The mean of segment k is the weighted mean of all customers, with weights equal to P(segment k | customer). Repeat. Each iteration is guaranteed to increase the marginal log-likelihood — EM cannot decrease it, though it can converge to a local maximum.
EM generalizes far beyond Gaussian mixture models. K-means is a hard-assignment EM: the E-step assigns each point to its nearest centroid with probability 1, and the M-step updates centroids as unweighted means. Hidden Markov Models use EM under the name Baum-Welch. Probabilistic PCA uses EM. The unifying pattern: any model where the complete-data likelihood is tractable but the marginal likelihood (summing or integrating over hidden variables) is not — EM is the natural algorithm.
The theoretical guarantee: EM increases the marginal log-likelihood L(θ) = log P(X | θ) at every iteration. This follows from Jensen's inequality applied to the log-sum structure of the marginal likelihood. The E-step constructs a lower bound that is tight at the current θ. The M-step maximizes that lower bound. The next iteration starts from a point where the bound and the true objective coincide — so the objective has not decreased.
NOT this. EM is not an algorithm for mixture models. EM is a general framework for maximum likelihood estimation when data has missing or latent variables. The pattern is always: treat the missing data as if it were observed but uncertain (E-step fills in the expected complete data), then maximize the resulting expected complete-data log-likelihood (M-step). K-means, Baum-Welch for HMMs, and probabilistic PCA are all instances of this pattern.
Key points
- Use EM whenever you have latent variables and the complete-data log-likelihood has a closed-form maximizer. This is the pattern in GMMs, HMMs, probabilistic PCA, and factor models. When the M-step does not have a closed-form solution, replace it with gradient ascent on the expected complete-data log-likelihood — this is called Generalized EM and still monotonically increases the marginal likelihood.
- Trap: EM converges to local optima. Run EM with multiple random initializations — typically 5 to 10 — and take the solution with the highest final log-likelihood. K-means++ initialization (choosing initial centroids with probability proportional to their squared distance from already-chosen centroids) gives better starting points and reduces the number of restarts needed for reliable convergence.
- Diagnostic: plot log-likelihood per iteration. It should monotonically increase. If it ever decreases, there is a bug in the M-step — the expected complete-data log-likelihood is not being maximized correctly. A non-monotone EM log-likelihood is not a convergence issue, it is a correctness issue. The monotonicity guarantee is a theorem, not an approximation.
EM converts one intractable optimization — maximizing the marginal likelihood when variables are hidden — into a sequence of tractable steps by alternating between filling in hidden variable distributions and maximizing the resulting expected log-likelihood. Any model with latent variables and a tractable complete-data likelihood is a candidate for EM.
Recap
- EM breaks the chicken-and-egg of latent variables: need labels for params, need params for labels.
- E-step: given θ, compute soft responsibilities P(segment k | point i), summing to 1 per point.
- M-step: given responsibilities, update params with weighted statistics (weighted means/covariances).
- Monotonic guarantee: each iteration increases the marginal log-likelihood (Jensen on the log-sum), converging to a local max.
- K-means = hard-assignment EM; HMM Baum-Welch, probabilistic PCA are all EM instances.
- Use EM when complete-data likelihood is tractable but the marginal (over hidden vars) is not.
- Local optima trap: run 5–10 random inits (or k-means++) and keep the highest final log-likelihood; a decreasing log-likelihood = M-step bug.
Check your understanding
Q1. Which two of the following statements correctly describe the E-step and M-step of EM for Gaussian Mixture Models (GMMs)?
- A) The E-step computes the responsibility of each point xᵢ to cluster k as r_{ik} = P(z=k|xᵢ,θ) = πₖN(xᵢ;μₖ,Σₖ)/ΣⱼπⱼN(xᵢ;μⱼ,Σⱼ) — a soft posterior probability, so each point is 'partly' in every cluster rather than hard-assigned to just one.
- B) Given the responsibilities r_{ik}, the M-step updates each cluster in closed form as a responsibility-weighted statistic: πₖ = Σᵢr_{ik}/n, μₖ = Σᵢr_{ik}xᵢ/Σᵢr_{ik}, and Σₖ = Σᵢr_{ik}(xᵢ−μₖ)(xᵢ−μₖ)ᵀ/Σᵢr_{ik} — no gradient step is taken.
- C) The E-step assigns each point to its single nearest centroid by Euclidean distance and the M-step recomputes each mean as the unweighted centroid of its assigned points — EM for GMMs is mathematically identical to k-means for any covariance structure.
- D) The E-step directly maximises the marginal log-likelihood by computing exact analytic gradients with respect to μₖ, Σₖ, and πₖ, which makes EM just a special case of gradient ascent that happens to use a closed-form, optimally-chosen step size.
Q2. EM guarantees that the log-likelihood is non-decreasing at each step. Prove this using Jensen's inequality.
- A) The log-likelihood is non-decreasing because the M-step maximises the complete-data log-likelihood so that Q(θ,θ_old) ≥ Q(θ_old,θ_old), and by the data processing inequality, maximising Q can never decrease the marginal log-likelihood log P(X|θ) at the next step. Jensen's inequality separately ensures that log E[P(X,Z|θ)] ≥ E[log P(X,Z|θ)], which bounds the improvement from below at every iteration.
- B) The M-step finds θ_new = argmax_θ Q(θ,θ_old), so Q(θ_new,θ_old) ≥ Q(θ_old,θ_old). Since log P(X|θ) = Q(θ,θ_old) + H(θ,θ_old) where H is always constant and independent of θ by definition of the E-step, it follows directly that log P(X|θ_new) ≥ log P(X|θ_old). The proof does not use Jensen's inequality at all — it only requires the M-step to locate a non-decreasing point of Q.
- C) The non-decreasing property follows from the concavity of log: log P(X|θ) = log Σ_Z P(X,Z|θ). By Jensen's inequality applied to the concave log function: log Σ_Z P(X,Z|θ) ≥ Σ_Z log P(X,Z|θ) · P(Z|X,θ_old) = Q(θ,θ_old). This lower bound is tight only when P(Z|X,θ) = P(Z|X,θ_old) exactly, achieved at θ=θ_old. The M-step increases Q, which raises the lower bound, but this alone does not prove log P(X|θ_new) ≥ log P(X|θ_old) without separately accounting for how the bound's tightness changes at θ_new versus θ_old.
- D) Define Q(θ,θ_old) = E_{z|x,θ_old}[log P(x,z|θ)], the expected complete-data log-likelihood. Decompose: log P(x|θ) = Q(θ,θ_old) − H(θ,θ_old), where H(θ,θ_old) is a KL divergence and is always ≥ 0. By Jensen (log is concave): log P(x|θ) ≥ Σ_z q(z)log(P(x,z|θ)/q(z)) — the ELBO. The E-step sets q(z) = P(z|x,θ_old), making this bound tight at θ_old. The M-step sets θ_new = argmax_θ Q(θ,θ_old), so Q(θ_new,θ_old) ≥ Q(θ_old,θ_old) = log P(x|θ_old), hence log P(x|θ_new) ≥ log P(x|θ_old).
Q3. EM converges, but the solution is often a local optimum. What strategies help escape poor local optima in practice?
- A) EM converges to a local maximum of the log-likelihood, or sometimes a saddle point — not the global optimum; GMM landscapes have many local optima from different cluster assignments. Strategies: (1) multiple random restarts, keeping the highest final log-likelihood; (2) k-means++ initialisation, seeding centroids spread out by distance; (3) annealing — start with broad, overlapping covariances and shrink them gradually.
- B) EM only converges to local optima when the number of mixture components k is misspecified relative to the true generative process. The correct strategy is to first determine the true k using BIC or AIC computed across a range of candidate values, then run EM exactly once with that optimal k. With the correctly specified k, EM always converges to the unique global optimum, because the log-likelihood surface of a correctly-specified Gaussian mixture is provably log-concave in θ.
- C) The primary strategy is gradient-based correction: after EM converges, run a few additional steps of full gradient ascent directly on the marginal log-likelihood log P(X|θ) to escape the local optimum EM settled into. EM is fundamentally a lower-bound maximiser and can therefore stop at non-stationary points of the true objective; gradient ascent detects these by checking whether the gradient norm ‖∇log P(X|θ)‖ exceeds a small threshold ε, and continues climbing directly on the likelihood if so.
- D) Local optima in EM are unavoidable and not especially problematic in practice, because all local optima of the GMM log-likelihood correspond to valid, interpretable cluster solutions of roughly comparable quality. The global optimum is only preferable from a strict statistical standpoint — in applied settings, several different local solutions actually provide useful ensemble diversity by capturing different plausible clusterings of the same data.
Try it interactively
ML Systems Lab is a free interview-prep platform for ML engineers — work through the full interactive module, quizzes, and drills.
Open ML Systems Lab →