🔔 Gaussian Mixture Models with Expectation Maximization

By Better Wise

k-means gave every point exactly one label and assumed round, equally sized clusters. You already know the loop: assign points to the nearest centroid, recompute centroids. A Gaussian mixture keeps that alternation but replaces the hard assignment with a posterior and the centroid with a weighted maximum-likelihood fit, which buys you soft labels, elliptical clusters and unequal mixing weights. Before you touch NumPy, derive the two steps on paper.

By the end of this lesson you can:

  • write the GMM likelihood and say how many parameters it has;
  • derive the E-step and M-step from the latent-variable lower bound;
  • compute responsibilities with log-sum-exp so nothing underflows;
  • update means, covariances and weights in closed form, with regularization;
  • read a log-likelihood curve and diagnose the two classic failures.

A latent-variable view of your data

Assume each point was produced in two stages: draw a component index $z \in {1,\dots,k}$ with probability $\pi_k$, then draw $x$ from $\mathcal{N}(\mu_k, \Sigma_k)$. Marginalizing $z$ gives the density you actually fit:

$$ p(x) = \sum_{k=1}^{K} \pi_k , \mathcal{N}(x \mid \mu_k, \Sigma_k) $$

Parameters: $(k-1)$ free weights, $k d$ mean entries, $k,d(d+1)/2$ covariance entries. In $d=50$ with $k=10$ full covariances that is over 12,000 numbers, which is why full-covariance GMMs live in low dimension or on PCA-reduced features.

The objective is the log-likelihood, and the sum inside the log is the whole problem:

$$ \mathcal{L}(\theta) = \sum_{n=1}^{N} \log \sum_{k=1}^{K} \pi_k , \mathcal{N}(x_n \mid \mu_k, \Sigma_k) $$

Without the latent labels, the gradient of that inner log couples every component to every other one. You cannot solve for $\mu_k$ in closed form.

The EM lower bound, in four lines on paper

Introduce $q(z_n)$, any distribution over the latent label of point $n$. Then

$$ \log p(x_n \mid \theta) = \log \sum_k q_{nk} \frac{p(x_n, z_n{=}k \mid \theta)}{q_{nk}} ;\ge; \sum_k q_{nk} \log \frac{p(x_n, z_n{=}k \mid \theta)}{q_{nk}} $$

by Jensen, since $\log$ is concave. The right-hand side is a lower bound you can push up in two cheap moves:

  1. Fix $\theta$, maximize over $q$: the bound is tight at the posterior, which gives the E-step.
  2. Fix $q$, maximize over $\theta$: the log is separable per component, so each parameter has a closed-form solution, which gives the M-step.

Because step 1 tightens the bound and step 2 raises it, the observed log-likelihood can never decrease. That is the only correctness guarantee you need to remember.

E-step: responsibilities, computed in log space

The optimal $q_{nk}$ is the posterior over components:

$$ r_{nk} = \frac{\pi_k , \mathcal{N}(x_n \mid \mu_k, \Sigma_k)}{\sum_j \pi_j , \mathcal{N}(x_n \mid \mu_j, \Sigma_j)} $$

In $d=30$ the density is around $10^{-20}$ per point, and a product over components underflows to zero. So never compute densities directly. Work with log joints, subtract the row maximum, then exponentiate:

import numpy as np

def log_gaussian(X, mu, cov):
    d = X.shape[1]
    L = np.linalg.cholesky(cov)              # one factorization, two uses
    sol = np.linalg.solve(L, (X - mu).T)     # (d, n)
    maha = np.sum(sol ** 2, axis=0)          # squared Mahalanobis distance
    logdet = 2.0 * np.sum(np.log(np.diag(L)))
    return -0.5 * (d * np.log(2 * np.pi) + logdet + maha)

def lse(A, axis=1):
    m = A.max(axis=axis, keepdims=True)
    return (m + np.log(np.exp(A - m).sum(axis=axis, keepdims=True))).squeeze(axis)

log_joint = np.log(pi)[None, :] + np.stack(
    [log_gaussian(X, mu[j], cov[j]) for j in range(K)], axis=1)   # (n, K)
row_ll = lse(log_joint, axis=1)                # per-sample log-likelihood
log_resp = log_joint - row_ll[:, None]
resp = np.exp(log_resp)
total_ll = row_ll.sum()

Cholesky gives you the log-determinant and the Mahalanobis term from one factorization. Keep log_resp around: the monitored objective falls out of the same computation for free.

Note: row_ll is $\log p(x_n)$. Summing it gives the exact objective you are maximizing. If you ever compute the log-likelihood in a separate function, your two numbers will disagree and you will not know which one is wrong.

Exercise (project thread). With $\pi = (0.5, 0.5)$, $\mu_1 = 0$, $\mu_2 = 4$, $\Sigma_1 = \Sigma_2 = 1$ in 1-D, compute $r_{n1}$ and $r_{n2}$ for $x = 1$ without running code.

Answer

The log densities differ only in the squared distance: $\log p_2 - \log p_1 = -\tfrac{1}{2}(9 - 1) = -4$, so $p_2/p_1 = e^{-4} \approx 0.0183$. Then $r_2 = 0.0183/1.0183 \approx 0.018$ and $r_1 \approx 0.982$. Note how fast a modest distance kills a responsibility, which is exactly why the mixed label sits near the midpoint.

M-step: weighted means and regularized scatter

Maximizing the bound with respect to the parameters is a weighted maximum-likelihood problem where point $n$ counts $r_{nk}$ times for component $k$. With $N_k = \sum_n r_{nk}$:

$$ \pi_k = \frac{N_k}{N}, \qquad \mu_k = \frac{1}{N_k}\sum_n r_{nk} x_n, \qquad \Sigma_k = \frac{1}{N_k}\sum_n r_{nk} (x_n - \mu_k)(x_n - \mu_k)^\top + \varepsilon I $$

for j in range(K):
    diff = X - mu[j]
    cov[j] = (resp[:, j, None] * diff).T @ diff / Nk[j] + reg * np.eye(d)

The $\varepsilon I$ term, usually 1e-6 relative to the feature scale, is not cosmetic. A component that captures a single point or a set of collinear duplicates drives $|\Sigma_k| \to 0$, the log-likelihood runs off to $+\infty$, and every downstream comparison in your capstone becomes meaningless.

Initialization follows directly from the previous lesson: run k-means, take the centroids as $\mu_k$, the cluster fractions as $\pi_k$, and hard assignments (or the k-means distances softened) as the first resp. Run several restarts and keep the model with the best final total_ll.

Exercise. Given resp = [[0.9, 0.1], [0.2, 0.8]] and X = [[0], [3]], compute $\pi$ and $\mu$.

Answer

$N_1 = 1.1$, $N_2 = 0.9$, so $\pi = (0.55, 0.45)$. Then $\mu_1 = (0.9 \cdot 0 + 0.2 \cdot 3)/1.1 \approx 0.545$ and $\mu_2 = (0.1 \cdot 0 + 0.8 \cdot 3)/0.9 \approx 2.667$. Both means move toward the points that claim them, not to the raw cluster means.

Reading the log-likelihood curve correctly

Three checks cover almost every bug:

  • Monotonicity. total_ll must be non-decreasing. A drop means a broken update: wrong axis in lse, missing log(pi) in the joint, a covariance that is no longer symmetric, or resp not summing to one per row.
  • Convergence. Stop on a relative change, for example (ll - prev) < 1e-6 * abs(ll), not on a fixed iteration count. EM converges linearly, and most datasets plateau in 30 to 100 iterations.
  • Collapse. If total_ll spikes by hundreds in one iteration, check min(np.linalg.eigvalsh(cov[j])). Near-zero eigenvalues mean a component died on a point or a line. Add regularization, drop the duplicate rows, or reduce $k$.

A word of caution for the capstone: training log-likelihood is comparable across runs on the same data and preprocessing, but it always rises with $k$. Never pick $k$ from training likelihood alone.

Exercise. At iteration 7 your run gains +412 nats and cov[3] has det = 3e-13. What happened, and what do you change first?

Answer

Component 3 collapsed onto essentially one direction with zero variance, so the Gaussian likelihood for those points diverges. Add reg * I to every covariance (or raise it from 1e-8 to 1e-6), then rerun with the same seed and confirm total_ll rises smoothly. If it recurs, look for exact duplicate rows or a $k$ that is too large for the data.

Field rules for an EM you can trust

  • Do your arithmetic in log space from the first line. Retrofitting stability after you see nan costs an afternoon.
  • Standardize your features first. A single unnormalized column makes cov badly conditioned and the Cholesky solve inaccurate.
  • Regularize every covariance, even when the run looks healthy. Set reg before the first iteration, not after the first crash.
  • Initialize with k-means and reuse the seed for your baseline comparison. Random responsibilities plus random restarts means you are comparing luck.
  • Keep log_resp, not just resp. The E-step already computed the objective, so monitoring should be free.
  • Compare against a library baseline with the same initialization, the same reg, and the same tolerance. Otherwise the gap you report is hyperparameters, not your implementation.

What you carry into the checkpoint

The GMM is the same alternating pattern as k-means, upgraded: the E-step is soft responsibilities computed from all components at once, and the M-step is a weighted moment fit. Nothing else changes. If k-means clicked, this is the same loop with a probability in place of an argmax.

What you carry into the checkpoint:

  • The ELBO decomposition $\log p(x) = \mathcal{L} + \mathrm{KL}$, and the fact that the E-step makes the bound tight while the M-step pushes it up. That is the whole proof of monotonicity, and it is what you should be able to say in one sentence when you write up the capstone.
  • The three M-step updates: $\pi_k = N_k / N$, $\mu_k = \frac{1}{N_k}\sum_n r_{nk} x_n$, $\Sigma_k = \frac{1}{N_k}\sum_n r_{nk}(x_n - \mu_k)(x_n - \mu_k)^\top$. Note the denominator is $N_k$, not $N$, and note it is not $N_k - 1$. The MLE does not correct for degrees of freedom.
  • Log-space everything: the log-sum-exp responsibilities, the log-likelihood accumulator, and the Cholesky solve instead of an explicit inverse.

Recall drill

The E-step computes normalised responsibilities in log space using the _____ trick, so that no single component underflows to zero. The M-step sets the mixing weight to _____, where N_k is the sum of that component's responsibilities over all points. Adding _____ to every covariance matrix keeps the likelihood finite when a component collapses onto few points. EM guarantees that the observed-data log-likelihood is _____ from one iteration to the next, so any decrease is a bug. When the responsibilities become hard zeros and ones, the GMM reduces exactly to _____ with spherical, equal covariances.

Front: What does the E-step of EM compute for a GMM? Back: The posterior responsibilities r_nk, the probability that point n was generated by component k.

Front: Why compute responsibilities with log-sum-exp? Back: The Gaussian log density can be very negative, so exponentiating first underflows to 0 and every responsibility becomes nan.

Front: M-step update for the mixing weights pi_k? Back: N_k / N, where N_k is the total responsibility assigned to component k.

Front: M-step update for the mean mu_k? Back: Responsibility-weighted average of the data: (1/N_k) * sum of r_nk * x_n.

Front: M-step update for the covariance Sigma_k? Back: Responsibility-weighted scatter divided by N_k (not N, and no Bessel correction).

Front: Where does the guarantee of monotonic likelihood come from? Back: The ELBO decomposition: the E-step makes the bound tight, the M-step maximises it, so the marginal log-likelihood cannot decrease.

Front: What is the ELBO in one sentence? Back: Expected complete-data log-likelihood under the responsibilities minus the entropy of the responsibilities.

Front: Why add reg * I to every covariance? Back: A component with near-zero variance in some direction makes the density diverge and the likelihood go to infinity or nan.

Front: How do you detect that EM has converged? Back: Track the observed-data log-likelihood and stop when the change is below a tolerance, typically 1e-4 to 1e-6.

Front: Why does the GMM log-likelihood have no closed-form maximum? Back: The log of a sum over components couples the parameters, so you need the alternating EM scheme instead of a single solve.

Front: How does the GMM relate to k-means? Back: k-means is the limit of EM with hard 0/1 responsibilities and shared spherical covariance with variance going to 0.

Front: What does a component weight of pi_k close to 0 tell you? Back: The component is starving, usually a sign that k is too large or the initialization placed it badly.

Front: Which quantity should you monitor at every iteration and why? Back: The total observed-data log-likelihood, because EM guarantees it is non-decreasing and any dip means an indexing or normalisation bug.

Front: What does a well-conditioned covariance matrix require of your features? Back: Standardised, roughly comparable scales, otherwise the determinant and the Cholesky solve lose precision.

Next up is the checkpoint, where PCA, k-means and the GMM sit side by side in one notebook: you will pick the right tool for a dimensionality reduction task and a density estimation task, and defend the choice with numbers rather than intuition.