Clustering: k-means and Gaussian mixtures

Clustering means splitting data points into groups (clusters) so that points in the same group are similar. Nobody gives the algorithm the right groups: there are no labels. This page shows the two classic methods on points in the plane:

k-means describes each cluster by one point, its centre $\mu_k$ . It repeats two steps: give every point to its nearest centre, then move every centre to the mean of its points.
A Gaussian mixture model (GMM) describes each cluster by a Gaussian distribution: a centre, a covariance (the size, shape and direction of an ellipse) and a weight (the share of the data). It is fitted with the EM algorithm (expectation–maximization), which also repeats two steps.

In the first simulation you run both methods one step at a time on data that you can change. In the second one you see the same data clustered by both methods side by side, and you find out when k-means gives the wrong answer. The math behind the pictures is in section 4.

1. k-means and EM, step by step #

Press Step again and again, and watch the two steps of k-means. Each data point takes the colour of its cluster (blue, orange, aqua, violet, magenta, yellow for clusters 1 to 6). The large ringed dots are the centres $\mu_k$ ; the lines behind them are their paths. With k-means the coloured regions are the Voronoi cells: each region holds the places that are closest to one centre. Press New start to see the seeding step by step: with k-means++ the size of a grey dot shows its chance to be picked as the next centre. Switch the method to Gaussian mixture (EM): the ellipses show 1 and 2 standard deviations of each Gaussian, a point's colour mixes the cluster colours with its responsibilities, and the background shows the mixture density. Drag points, click to add points or a whole cloud, and drag a centre to pick your own start. The first chart shows the objective of this run after every step; the second chart is the elbow plot (k-means) or BIC (mixture) for $K = 1, \dots, 8$ . Click it to change $K$ .

2. When k-means fails #

Move the slider and watch the two clusterings change. The first panel shows k-means, the second a Gaussian mixture fitted with EM. Both get the true number of clusters $K$ and keep the best of several runs, so a failure is not bad luck. Colours are matched to the true clusters, and a ring marks each point in the wrong cluster. The chart shows the adjusted Rand index (ARI) of both methods against the slider value: $1$ means a perfect clustering, about $0$ means no better than random groups. Press Show the true clusters to see the answer, and try the other kinds of data.

3. Try this #

Can the objective go up? Keep blobs and $K = 4$ and press Step slowly. Question: after an assign step and after an update step, can the inertia $J$ (the sum of squared distances from each point to its centre) become larger? What to see: the chart only goes down. After an update step, some points lie in a region of another colour: the next assign step moves exactly these points. Here k-means stops after 2 iterations with $J \approx 180$ .
A bad start. Set Start to random points and press To the end. Question: do the four centres find the four blobs? What to see: no. Two centres share one blob, and one centre sits between two blobs. k-means stops with $J \approx 578$ , three times the $J \approx 180$ of the good answer, and it never escapes: this is a local minimum. Press New start a few times with both kinds of start. In our runs, random points ended in a worse local minimum in 16 of 40 starts, and k-means++ in 4 of 40. Step through a k-means++ start: the big grey dots (likely picks) are always far from the centres chosen so far.
How many clusters? Look at the elbow plot for blobs. Question: where does the curve bend? What to see: the best inertia drops fast up to $K = 4$ (2539, 1066, 591, 180) and then slowly (158, 142, …). The bend, the elbow, is at the true $K = 4$ . Now pick rings or elongated: the curve has no clear elbow. Switch the method to Gaussian mixture and read the BIC chart: for blobs and elongated the minimum is at the true $K$ .
Long clusters. In the second figure, pick elongated and move length / width from 1 to 7. Question: at what point does k-means start to cut across the clusters? What to see: the ARI of k-means falls from $0.98$ to about $0.45$ near a ratio of $3.5$ , while the mixture stays near $0.97$ . At the default setting k-means puts 70 points in the wrong cluster, but its inertia ( $J \approx 362$ ) is lower than the inertia of the true clusters ( $J \approx 616$ ). k-means is not stuck here: its objective prefers the wrong answer. In the first figure, choose elongated, Gaussian mixture, and spherical covariance: with circles only, EM fails like k-means.
One big cluster, two small ones. Pick uneven sizes. Question: which points does k-means give to the wrong cluster? What to see: the outer points of the big cluster, on the side of the small clusters. The true sizes are 180, 30 and 30; k-means finds about 160, 36 and 44. The mixture learns the weights $\pi = (0.75, 0.12, 0.12)$ and a large covariance for the big cluster, and gets every point right. On the chart, the ARI of k-means falls when the spread ratio grows above about 3, while the ARI of the mixture stays near $1$ up to a ratio of about 5.5.
Rings. Pick rings. Question: can either method separate the inner ring from the outer one? What to see: no; both ARIs stay near $0$ for every radius. The ARI of the mixture is even a little below $0$ : its clusters agree with the rings a little less than random groups would. A ring is not convex: the straight line between two of its points leaves the ring. A k-means region is always convex, and a Gaussian puts most of its mass near its centre, which here is in the middle of the other ring. In the first figure, the BIC of the mixture on rings keeps falling up to $K = 7$ : the mixture covers each ring with several small Gaussians. That is a good model of the density, but not two clusters.

4. The math behind the figures #

k-means minimizes the inertia #

We have $n$ points $x_1, \dots, x_n \in \R^2$ and want $K$ clusters. k-means chooses centres $\mu_1, \dots, \mu_K$ and an assignment $a_i \in \{1, \dots, K\}$ for each point that make the inertia (the within-cluster sum of squares) small:

J(\mu, a) = \sum_{i=1}^n \lVert x_i - \mu_{a_i} \rVert^2 .

Lloyd's algorithm (the usual k-means algorithm) changes one group of variables at a time:

Assign step. Keep the centres fixed. Each term of $J$ is smallest when $a_i$ is the nearest centre: $a_i = \arg\min_k \lVert x_i - \mu_k \rVert^2$ .
Update step. Keep the assignment fixed. The sum of squares $\sum_{i : a_i = k} \lVert x_i - \mu \rVert^2$ is a convex quadratic in $\mu$ ; its gradient $-2 \sum_{i : a_i = k} (x_i - \mu)$ is zero at the mean: $\mu_k = \frac{1}{n_k} \sum_{i : a_i = k} x_i$ , where $n_k$ is the number of points in cluster $k$ .

Each step minimizes $J$ over one group of variables, so $J$ never increases. There are only finitely many assignments, and none can repeat with a lower $J$ , so the algorithm stops after finitely many steps, when no point changes cluster. But it stops at a local minimum, which depends on the start. Finding the global minimum is NP-hard in general, so in practice we run k-means from several starts and keep the lowest $J$ .

Why the regions are polygons. A place $x$ is closer to $\mu_j$ than to $\mu_k$ when $\lVert x - \mu_j \rVert^2 < \lVert x - \mu_k \rVert^2$ . Expand both sides; the $\lVert x \rVert^2$ terms cancel and we get a linear inequality:

2 (\mu_k - \mu_j)^\top x < \lVert \mu_k \rVert^2 - \lVert \mu_j \rVert^2 .

So the border between two clusters is a straight line: the perpendicular bisector of the two centres. A Voronoi cell is an intersection of half-planes, so it is a convex polygon. This explains the failures in section 2: k-means cannot follow a ring (not convex), and it puts the border halfway between two centres even when one cluster is much wider than the other (uneven variance). For long clusters, a cut across them gives a lower $J$ than the true clusters, so the objective itself prefers the wrong answer.

Worked example (one dimension). Points $1, 2, 3, 10, 11, 20$ with $K = 2$ and start centres $1$ and $20$ .

Assign: $10$ is closer to $1$ (distance $9$ ) than to $20$ (distance $10$ ); $11$ is closer to $20$ (distance $9$ ) than to $1$ (distance $10$ ). Clusters $\{1, 2, 3, 10\}$ and $\{11, 20\}$ , with $J = (0 + 1 + 4 + 81) + (81 + 0) = 167$ .
Update: the means are $\mu_1 = 4$ and $\mu_2 = 15.5$ ; now $J = (9 + 4 + 1 + 36) + (20.25 + 20.25) = 90.5$ .
Assign: $10$ is now closer to $15.5$ (distance $5.5$ ) than to $4$ (distance $6$ ). Clusters $\{1, 2, 3\}$ and $\{10, 11, 20\}$ , $J = 14 + 70.75 = 84.75$ .
Update: $\mu_1 = 2$ , $\mu_2 = 41/3 \approx 13.67$ , $J = 2 + 60.67 = 62.67$ . The next assign step changes nothing, so k-means stops. The inertia went $167 \to 90.5 \to 84.75 \to 62.67$ : down at every step.

Quick check

Which statement about Lloyd's algorithm is true?

k-means++: a better start #

k-means++ (Arthur and Vassilvitskii, 2007) chooses the start centres among the data points. The first centre is a uniformly random point. Each next centre is point $x_i$ with probability

P(x_i) = \frac{D(x_i)^2}{\sum_{j=1}^n D(x_j)^2},

where $D(x)$ is the distance from $x$ to the nearest centre chosen so far. This is $D^2$ sampling: points far from all centres, probably in a cluster without a centre yet, are likely to be picked. Points that already have a centre nearby are rarely picked. The authors proved that the start alone already has expected inertia at most $8(\ln K + 2)$ times the optimal $J$ ; Lloyd's steps only lower it further. It is the default start in scikit-learn.

In the worked example with first centre $1$ , the squared distances are $0, 1, 4, 81, 100, 361$ with sum $547$ , so the second centre is $20$ with probability $361/547 \approx 0.66$ .

Quick check

Three points on a line: $0$ , $1$ and $5$ . k-means++ picked $0$ as the first centre. What is the probability that the second centre is $5$ ?

Gaussian mixtures #

A Gaussian mixture says that each point was made in two steps: first pick a cluster $z \in \{1, \dots, K\}$ with probabilities $\pi_1, \dots, \pi_K$ (the mixing weights, $\pi_k \ge 0$ , $\sum_k \pi_k = 1$ ), then draw $x$ from the Gaussian of that cluster, $\mathcal N(\mu_k, \Sigma_k)$ . The cluster $z$ is a hidden (latent) variable. The density of one point is

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

and we fit $\theta$ by maximizing the log-likelihood $\ell(\theta) = \sum_{i=1}^n \log \sum_{k=1}^K \pi_k \, \mathcal N(x_i \mid \mu_k, \Sigma_k)$ . The sum inside the logarithm couples all clusters, so there is no closed-form maximum. If we knew the cluster of every point, the problem would be easy: fit one Gaussian per cluster. EM uses this idea with soft guesses of the clusters.

E-step (expectation). Compute the responsibilities: the posterior probability that point $i$ came from cluster $k$ , by Bayes' rule with the current parameters,

r_{ik} = p(z_i = k \mid x_i, \theta) = \frac{\pi_k \, \mathcal N(x_i \mid \mu_k, \Sigma_k)}{\sum_{j=1}^K \pi_j \, \mathcal N(x_i \mid \mu_j, \Sigma_j)} .

Each row $r_{i1}, \dots, r_{iK}$ sums to $1$ . These are the colour mixes in the figure.

M-step (maximization). Fit each Gaussian to all points, each point weighted by its responsibility. With $N_k = \sum_i r_{ik}$ (the "soft size" of cluster $k$ ):

\pi_k = \frac{N_k}{n}, \qquad \mu_k = \frac{1}{N_k} \sum_{i=1}^n r_{ik}\, x_i, \qquad \Sigma_k = \frac{1}{N_k} \sum_{i=1}^n r_{ik}\, (x_i - \mu_k)(x_i - \mu_k)^\top .

With spherical covariance the widget uses $\Sigma_k = \sigma_k^2 I$ with $\sigma_k^2 = \frac{1}{2 N_k} \sum_i r_{ik} \lVert x_i - \mu_k \rVert^2$ .

EM never decreases the log-likelihood, just as Lloyd's algorithm never increases $J$ . The widget stops when $\ell / n$ changes by less than $10^{-5}$ . Like k-means, EM finds a local maximum that depends on the start; the widget starts the means with k-means++ or random points, each covariance as a circle, and $\pi_k = 1/K$ .

Why EM never decreases the log-likelihood

For any distribution $q_i$ over the cluster of point $i$ , the log-likelihood of that point splits into two parts (see the ELBO lesson):

\log p(x_i \mid \theta) = \underbrace{\sum_k q_i(k) \log \frac{p(x_i, z_i = k \mid \theta)}{q_i(k)}}_{\text{ELBO}(q_i,\, \theta)} + \underbrace{\mathrm{KL}\big(q_i \,\big\|\, p(z_i \mid x_i, \theta)\big)}_{\ge 0} .

The E-step sets $q_i = p(z_i \mid x_i, \theta_{\text{old}})$ , which is $r_{i\cdot}$ . Then the KL term is $0$ , so $\sum_i \text{ELBO}(q_i, \theta_{\text{old}}) = \ell(\theta_{\text{old}})$ . The M-step maximizes $\sum_i \text{ELBO}(q_i, \theta)$ over $\theta$ ; for Gaussians the maximizer is exactly the formulas above (set the gradients to zero; for $\pi$ use a Lagrange multiplier for $\sum_k \pi_k = 1$ ). So

\ell(\theta_{\text{new}}) \ge \sum_i \text{ELBO}(q_i, \theta_{\text{new}}) \ge \sum_i \text{ELBO}(q_i, \theta_{\text{old}}) = \ell(\theta_{\text{old}}) .

The first inequality holds because KL is never negative (Jensen's inequality).

Quick check

One dimension, two Gaussians with $\pi_1 = \pi_2 = 0.5$ , $\sigma_1 = \sigma_2 = 1$ , $\mu_1 = 0$ and $\mu_2 = 2$ . What is the responsibility $r_{1}$ of the first Gaussian for the point $x = 0.5$ ?

What EM adds to k-means #

Shape and orientation. Each cluster has its own covariance $\Sigma_k$ , an ellipse. The border between two Gaussians is where $\log \pi_j \mathcal N(x \mid \mu_j, \Sigma_j) = \log \pi_k \mathcal N(x \mid \mu_k, \Sigma_k)$ : a quadratic curve (a conic), not only a straight line. This fixes the long clusters.
Sizes. The weights $\pi_k$ and the covariances let a wide, large cluster win the points far from its centre. This fixes the uneven clusters.
Soft assignments. A point between two clusters gets two responsibilities near $0.5$ , not a forced choice, so we can see which assignments are uncertain.
A probability model. $p(x \mid \theta)$ is a density: it gives a likelihood for choosing $K$ (below), and a low $p(x \mid \theta)$ flags an unusual point (anomaly detection).

k-means is a limit of EM. Fix all covariances to $\sigma^2 I$ and all weights to $1/K$ , and let $\sigma \to 0$ . Then $r_{ik} \propto \exp(-\lVert x_i - \mu_k \rVert^2 / 2\sigma^2)$ goes to $1$ for the nearest centre and $0$ for the others: the E-step becomes the assign step, and the M-step for $\mu_k$ becomes the update step. So k-means is "hard EM" for round clusters of the same size and weight, and this is exactly the case where it works well.

What EM does not fix: a Gaussian is still one blob around one centre, so a ring stays out of reach. Methods built on neighbours, such as spectral clustering, DBSCAN, or a kernel version of k-means, can follow such shapes.

A trap: the likelihood has no maximum. Put one mean exactly on a data point, $\mu_k = x_i$ , with $\Sigma_k = \sigma^2 I$ . The term $\mathcal N(x_i \mid \mu_k, \sigma^2 I) = 1/(2\pi\sigma^2)$ grows without limit as $\sigma \to 0$ , and so does $\ell$ . A Gaussian that sits on one point (or on points along a line) can collapse. Practical code adds a small ridge to each covariance (scikit-learn adds $10^{-6} I$ , called reg_covar), uses a prior, or restarts. The widget adds $10^{-6} I$ too, and it stops with a message when a covariance becomes almost singular. To see it: press Clear all, set Click to adds a cloud and click once, then set it to adds a point and click once far away. With $K = 2$ and Gaussian mixture, press To the end. k-means++ gives the lonely point its own Gaussian, and that Gaussian shrinks onto it.

Choosing the number of clusters $K$ #

The inertia always goes down when $K$ grows (with $K = n$ it is $0$ ), so we cannot simply minimize it.

Elbow plot. Plot the best $J$ for each $K$ and look for the bend: after it, more clusters reduce $J$ only a little. This works for well-separated blobs and often gives no clear answer otherwise.
BIC. For a probability model, the Bayesian information criterion $\text{BIC} = -2\, \ell(\hat\theta) + p \ln n$ trades fit against size. Here $\hat\theta$ are the fitted parameters and $p$ is the number of free parameters. A 2D mixture with full covariances has $p = (K - 1) + 2K + 3K = 6K - 1$ ( $K - 1$ weights, $2$ numbers per mean, $3$ per covariance). Smaller is better. On blobs the BIC is smallest at $K = 4$ .
Other tools: the silhouette score, the log-likelihood on held-out data, and, most important, whether the clusters are useful for the task.

Measuring the result. When the true groups are known (as in the second figure), the adjusted Rand index compares two clusterings: it counts the pairs of points that are put together or apart in both, and corrects for the agreement expected by chance. It is $1$ for identical clusterings (up to renaming the clusters) and about $0$ for random ones.

Key takeaways

k-means minimizes the inertia $J = \sum_i \lVert x_i - \mu_{a_i} \rVert^2$ by repeating two steps: assign each point to its nearest centre, then move each centre to the mean of its points.
Both steps minimize $J$ over one group of variables, so $J$ never increases and the algorithm stops; but it stops at a local minimum that depends on the start, so we use several starts.
k-means++ picks each new start centre with probability proportional to $D(x)^2$ , the squared distance to the nearest centre so far; its expected $J$ is within a factor $O(\log K)$ of the optimum.
The border between two k-means clusters is the perpendicular bisector of their centres, so each cluster region is a convex polygon. k-means fails on long, non-convex or unequal-size clusters, and for long clusters a wrong answer can have a lower $J$ than the true clusters.
A Gaussian mixture models $p(x) = \sum_k \pi_k \mathcal N(x \mid \mu_k, \Sigma_k)$ . EM alternates the E-step $r_{ik} = \pi_k \mathcal N(x_i \mid \mu_k, \Sigma_k) / \sum_j \pi_j \mathcal N(x_i \mid \mu_j, \Sigma_j)$ and the M-step (weighted $\pi_k = N_k/n$ , means and covariances); the log-likelihood never decreases.
EM adds a covariance (shape and direction) and a weight for each cluster, soft assignments, and a density. k-means is the limit of EM with equal weights and covariances $\sigma^2 I$ , $\sigma \to 0$ .
The GMM likelihood is unbounded: a Gaussian that shrinks onto one point sends it to infinity, so practical code adds a small ridge to the covariances.
Choose $K$ with the elbow of the inertia curve or with $\text{BIC} = -2\ell + p \ln n$ ; neither works for clusters that are not blob-shaped, such as rings.

ML Fundamentals: the track has a full lesson on clustering, Gaussian mixtures and EM, with anomaly detection and hierarchical clustering.
The ELBO: variational inference, VAEs and importance-weighted bounds: EM as coordinate ascent on the ELBO, and how VAEs replace the exact E-step by a network.
Jensen's inequality and the KL divergence: the inequality behind the EM guarantee.
Positive definite matrices and Cholesky: covariance matrices and their ellipses.
PCA as SVD of centered data: the axes of a covariance ellipse are its eigenvectors.
Gaussian integrals and Laplace's method: working with the multivariate Gaussian.
Nonconvex landscapes: local minima, the same problem k-means and EM have.
Probability and Statistics: maximum likelihood, Bayes' rule and latent variables.
Generative models: from mixtures of Gaussians to VAEs and diffusion models.

1. k-means and EM, step by step #

2. When k-means fails #

3. Try this #

4. The math behind the figures #

k-means minimizes the inertia #

k-means++: a better start #

Gaussian mixtures #

What EM adds to k-means #

Choosing the number of clusters KKK #

5. Related lessons #

Choosing the number of clusters $K$ #