← Atelier
Lesson 06 · Unsupervised learning · ~30 min + code

k-means, against Lloyd's algorithm

You have users described by two features and want \(k\) groups, so that every user is close to their group's center. What does "best" mean, how do you find it, and when does the standard algorithm fail?

1 · The objective

What makes a clustering good?

Pick \(k\) centers \(\mu_1,\dots,\mu_k\) and give each point \(x_i\) to one of them. k-means scores the result by the total squared distance from each point to its own center, the inertia:

\[ J = \sum_{i=1}^{n} \big\lVert x_i - \mu_{c(i)} \big\rVert^2 \]

Two things can change: the assignments \(c(i)\) and the centers \(\mu_j\). Optimizing both at once is hard (NP-hard in general). Each one alone is easy.

2 · Two moves

Fix one, solve the other

The centers are fixed. Which assignment of each point minimizes \(J\)?

\(J\) is a sum of one term per point, and each point's term depends only on its own assignment. So each point independently picks the center that makes its term smallest: the nearest one. That carves the plane into regions, one per center (a Voronoi diagram).

Now the assignments are fixed. Where should a center go to minimize the squared distances to its points?
\[ \begin{aligned} \frac{\partial}{\partial \mu}\sum_{i \in S}\lVert x_i - \mu\rVert^2 &= -2\sum_{i\in S}(x_i - \mu) = 0 \\ \Rightarrow\quad \mu &= \frac{1}{|S|}\sum_{i\in S} x_i \end{aligned} \]

The mean. (The median minimizes the sum of absolute distances: that's k-medians, more robust to outliers.)

Alternate the two moves and you have Lloyd's algorithm, the thing everyone calls k-means. Each move can only lower \(J\) or leave it unchanged, and there are finitely many assignments, so it always stops. But it stops at a point where neither move helps. That's a local minimum, not necessarily the best clustering. The game shows the difference.

3 · Play it

Place the centers

Drag the ✕ centers. Points take the color of their nearest center, and the shaded regions show who owns what. Your score is the inertia \(J\). Get within 3% of the best clustering, then compare yourself with Lloyd.

4 · Initialization

Where you start decides where you stop

Level 2 showed Lloyd getting stuck. The usual fix is choosing the starting centers well and running several restarts.

You've picked the first starting center. How should you pick the next one?

That's k-means++ (Arthur & Vassilvitskii, 2007). Far points are likely picks, so new centers tend to land in uncovered clusters. It's still random, so a single outlier doesn't reliably grab a center, which is the weakness of always taking the farthest point. Uniform sampling often puts two centers in one blob. k-means++ also comes with a guarantee: the expected result is within \(O(\log k)\) of optimal before Lloyd even runs.

300 runs of Lloyd on level 2's data, by starting rule

"Found the best" means final inertia within 0.1% of the best any run reached. Runs are recomputed in your browser.

5 · Scale and shape

k-means sees only distances

Real features come in different units. Here users differ in sessions per week (three real groups) and average session length in seconds (noise, 0 to 3,600).

You run k-means with k = 3 on the raw features. What do you get?

Distances in seconds are about 200× larger than distances in sessions, so the seconds axis decides every assignment. Standardize the features (or choose a metric deliberately) and the real groups appear:

Raw features
Standardized features

Same data and same k. Colors are the clusters k-means found. Axes are drawn in each feature's own units.

The same objective also explains k-means's other blind spots. It prefers round clusters of similar size: squared distance to a single center can't describe a ring, a long thin cloud, or a small dense cluster next to a big diffuse one. It also has no notion of noise: every outlier is assigned somewhere and drags that center.

6 · New cases, no hints

Apply it somewhere else

You plot inertia \(J\) against \(k\) and pick the \(k\) with the lowest \(J\). What goes wrong?

More centers can always fit at least as well, and every point its own center gives \(J = 0\). Look for an elbow, use silhouette or the gap statistic, or better, choose \(k\) by how the clusters are used downstream (segment sizes, retrieval recall, business constraints).

During Lloyd's algorithm one center ends up with no points assigned. What should a robust implementation do?

The mean of an empty set is undefined. Leaving the center in place wastes one of your \(k\) clusters. Re-seeding at a poorly served point (or splitting the largest cluster) keeps all \(k\) useful and still lowers \(J\).

Your data is two concentric rings. What does k-means with k = 2 produce?

Clusters from nearest-center assignment are convex regions (Voronoi cells), and a ring isn't one. Use a method built on connectivity or density (spectral clustering, DBSCAN), or first transform the features, for example to radius.

7 · Implement, unaided

Write k-means from scratch

This is the kind of exercise a coding screen can ask for, so treat it like one. State the complexity per iteration (\(O(nkd)\)), the stopping rule, and the empty-cluster policy before you type. ⌘/Ctrl+Enter runs.

Exercise 1: k-means++ initialization

Exercise 2: Lloyd's algorithm

8 · Staff-level follow-ups

Answer out loud, then check

Why does Lloyd's algorithm converge, and to what?
  • It's coordinate descent on \(J\): the assignment step is optimal given the centers, and the update step is optimal given the assignments. So \(J\) never increases.
  • There are finitely many partitions, so it terminates, but at a local minimum (a fixed point). Global k-means is NP-hard.
  • In practice: k-means++ initialization, several restarts, and keep the lowest \(J\).
How does k-means relate to Gaussian mixture models?
  • A GMM fit by EM uses soft assignments (responsibilities) and learns means, covariances and weights.
  • k-means is the limit of EM for a GMM with equal, spherical covariances \(\sigma^2 I\) as \(\sigma \to 0\): responsibilities become hard, nearest-center assignments.
  • Prefer a GMM when clusters are elliptical, overlap, or you need probabilities of membership.
How is k-means used in large-scale retrieval?
  • IVF indexes cluster the corpus embeddings with k-means (the coarse quantizer). A query scans only the nprobe nearest clusters: nprobe trades recall for latency.
  • Product quantization runs k-means separately on sub-vectors and stores codebook indices, compressing vectors for memory and fast approximate distances.
  • Training uses a sample and mini-batch or GPU k-means (e.g. Faiss). As embeddings drift, re-train the centroids, or cluster sizes skew and recall drops.
How would you run k-means on a billion points?
  • Mini-batch k-means: update centers from small random batches with per-center learning rates. Much cheaper, with slightly worse \(J\).
  • Parallel initialization (k-means||) instead of the sequential k-means++. Fit on a sample, then assign everything.
  • The assignment step is a nearest-neighbor search: use GPUs or ANN for large \(k\). Distribute by sharding the points and aggregating per-center sums and counts.
How do you evaluate a clustering when there are no labels?
  • Internal: inertia (only for comparing at fixed \(k\)), silhouette, Davies–Bouldin, and stability across seeds and subsamples.
  • External, when some labels exist: adjusted Rand index, NMI.
  • Best: the downstream task. Does segmenting improve a model, a targeting decision, or retrieval recall at fixed latency?

Sources: S. Lloyd, Least squares quantization in PCM (1982); Arthur & Vassilvitskii, k-means++: The Advantages of Careful Seeding (SODA 2007); D. Sculley, Web-Scale K-Means Clustering (WWW 2010); Jégou, Douze & Schmid, Product Quantization for Nearest Neighbor Search (TPAMI 2011); Johnson, Douze & Jégou, Billion-scale similarity search with GPUs (2017).