Gaussian mixture models and EM

The big idea

k-means gives each point exactly one cluster, and assumes every cluster is a round ball. Real groups overlap and stretch. A customer who is 60 percent "bargain hunter" and 40 percent "weekend splurger" is not well described by a single label.

A Gaussian mixture modelA probability model in which data are generated by first choosing one of k Gaussian components with given weights, then drawing a point from that component's normal distribution. (GMM) imagines the data was made by a two-step recipe: pick one of k bell-shaped clouds at random (with some weights), then draw a point from that cloud. Each cloud has its own centre, its own spread and its own tilt. Clustering means working backwards from the points to the clouds, and the answer for each point is a soft clusteringClustering in which each point gets a probability of belonging to each cluster instead of a single label.: a probability for each cloud.

The fitting algorithm, expectation-maximizationExpectation-maximization: an iterative algorithm for models with hidden variables that alternates estimating the hidden variables given the parameters (E-step) and the parameters given those estimates (M-step). (EM), is a chicken-and-egg loop. If you knew which cloud each point came from, fitting the clouds would be easy. If you knew the clouds, computing each point's membership would be easy. So guess one, compute the other, and alternate.

How it works

Step 1: The model

Each component j has a mixing weight pi_j (the share of points it generates), a mean vector mu_j and a covariance matrix Sigma_j (spread and tilt). The probability density of a point is a weighted sum of bells:

Fitting means choosing the parameters that make the observed data most likely, called maximum likelihoodChoosing model parameters that maximise the probability (likelihood) of the observed data under the model.. There is no closed-form solution, because we do not know which component made which point. That unknown is a latent variableA variable that is part of the model but never observed, such as which mixture component generated each point..

Step 2: E-step, compute responsibilities

Given the current parameters, compute for every point i and component j the responsibilityIn a mixture model, the posterior probability that a given point was generated by a given component., the probability that component j produced point i, with Bayes' rule:

Each row of responsibilities sums to 1. In the code, the point (3.6, 2.0), which sits between the clouds, ends up with responsibilities 0.391 and 0.609.

Step 3: M-step, refit each component with weighted points

Treat responsibilities as fractional memberships. Each component's new weight is its share of total responsibility, its mean is the responsibility-weighted average of all points, and its covariance is the weighted covariance around that mean. A point with responsibility 0.6 counts as 0.6 of a point for that component.

Step 4: Repeat until the log-likelihood stops improving

In the code the log-likelihood rises from -1908.0 after the first iteration to -1134.5 at convergence after 26. k-means is the special case of EM where every covariance is the same tiny circle and responsibilities are rounded to 0 or 1, which is why it has the same assign-update rhythm.

Going deeper

Code: EM from scratch, soft membership and choosing k

import numpy as np
from scipy.stats import multivariate_normal as mvn
from sklearn.cluster import KMeans
from sklearn.metrics import adjusted_rand_score
from sklearn.mixture import GaussianMixture

rng = np.random.default_rng(1)
# Two elongated, tilted clusters that overlap a little
A = rng.multivariate_normal([2, 2], [[1.6, 1.2], [1.2, 1.2]], 200)
B = rng.multivariate_normal([5, 2], [[0.3, -0.25], [-0.25, 0.9]], 200)
X = np.vstack([A, B]); truth = np.repeat([0, 1], 200)

# EM from scratch for k = 2
k = 2
mu = np.array([[3.0, 4.0], [4.0, 0.0]]); cov = np.array([np.eye(2)] * k); w = np.full(k, 1 / k)
prev = -np.inf
for it in range(1, 200):
    # E-step: responsibilities r[i, j] = P(cluster j | x_i)
    dens = np.column_stack([w[j] * mvn(mu[j], cov[j]).pdf(X) for j in range(k)])
    r = dens / dens.sum(1, keepdims=True)
    ll = np.log(dens.sum(1)).sum()
    # M-step: weighted means, covariances and mixing weights
    nk = r.sum(0)
    w = nk / len(X)
    mu = (r.T @ X) / nk[:, None]
    cov = np.array([((r[:, j, None] * (X - mu[j])).T @ (X - mu[j])) / nk[j] for j in range(k)])
    if it in (1, 2, 5) or ll - prev < 1e-6:
        print(f"iter {it:>2}: log-likelihood {ll:8.1f}, means {np.round(mu, 2).tolist()}")
    if ll - prev < 1e-6:
        break
    prev = ll
print("final covariances:", np.round(cov, 2).tolist())

# Soft membership of a point between the clusters
x = np.array([3.6, 2.0])
d = np.array([w[j] * mvn(mu[j], cov[j]).pdf(x) for j in range(k)])
print(f"point {x.tolist()}: P(cluster) = {np.round(d / d.sum(), 3).tolist()}")

hard = r.argmax(1)
print(f"\nARI  GMM (EM above) {adjusted_rand_score(truth, hard):.3f}   "
      f"k-means {adjusted_rand_score(truth, KMeans(2, n_init=10, random_state=0).fit_predict(X)):.3f}")

# Choosing the number of components with BIC (lower is better)
print("\nk   BIC")
for kk in range(1, 6):
    print(f"{kk}   {GaussianMixture(kk, n_init=3, random_state=0).fit(X).bic(X):.1f}")

iter  1: log-likelihood  -1908.0, means [[3.26, 2.67], [3.98, 1.31]]
iter  2: log-likelihood  -1323.3, means [[3.2, 2.64], [4.04, 1.36]]
iter  5: log-likelihood  -1304.5, means [[2.82, 2.48], [4.44, 1.53]]
iter 26: log-likelihood  -1134.5, means [[2.16, 2.11], [5.04, 1.93]]
final covariances: [[[1.17, 0.96], [0.96, 1.02]], [[0.31, -0.23], [-0.23, 1.04]]]
point [3.6, 2.0]: P(cluster) = [0.391, 0.609]

ARI  GMM (EM above) 0.874   k-means 0.739

k   BIC
1   2721.9
2   2335.0
3   2371.8
4   2409.5
5   2441.5

On these two tilted, overlapping clouds the GMM recovers the true groups with an adjusted Rand index of 0.874, while k-means, limited to round clusters, gets 0.739.

Choosing the number of components

Likelihood always improves with more components, so compare models with a penalty for parameters. The BICBayesian information criterion: minus twice the log-likelihood plus the number of parameters times the log of the number of points. Lower is better. is the usual choice. In the code it is lowest (2335.0) at k = 2, the true number, and rises for k = 3, 4 and 5.

Practical issues

• Covariance type: full covariances (any tilt) need d(d+1)/2 parameters per component, which explodes in high dimensions. covariance_type="diag", "tied" or "spherical" constrain them.
• Collapse: a component can shrink onto a single point, sending its likelihood to infinity. Libraries add a small value to the diagonal (reg_covar).
• Initialization: EM finds local optima; scikit-learn initializes from k-means and supports several restarts (n_init).
• Density estimation: a fitted GMM gives p(x) for any point, so low-density points can be flagged as anomalies ( Anomaly detection ), and you can sample new points from it.
• EM is general: the same E and M steps fit hidden Markov models, handle missing data and train many latent-variable models.

Trade-offs

Common mistakes

• Reading responsibilities as certainty: 0.61 versus 0.39 means the point is genuinely ambiguous.
• Choosing k by the highest likelihood, which always prefers more components; use BIC or held-out likelihood.
• Fitting full covariances with 100 features and a few hundred points.
• Forgetting that EM can converge to a poor local optimum; run several initializations.
• Expecting a GMM to model crescent or ring shapes; it can only approximate them with many ellipses.

In the wild

Gaussian mixtures were for decades the acoustic model in speech recognition, modelling how each speech sound is spread in feature space, before neural networks took over. Today they appear in background subtraction for video (each pixel's colour as a small mixture), in density-based anomaly detection, in customer segmentation where soft membership is useful, and as a teaching ground for EM, which itself powers hidden Markov models and many missing-data methods.

Check yourself

Recap