k-means, and what it assumes about your data

The algorithm in four lines

1. Pick k random points as initial centres.
2. Assign every row to its nearest centre.
3. Move each centre to the mean of the rows assigned to it.
4. Repeat 2 and 3 until nothing moves.

That is the whole of k-means. It converges quickly, it is easy to implement, and it is the most used clustering algorithm by a wide margin.

from sklearn.cluster import KMeans
km = KMeans(n_clusters=5, n_init=10, random_state=0).fit(X_scaled)
km.labels_, km.cluster_centers_, km.inertia_

inertia_ is the sum of squared distances from each point to its assigned centre. It is what the algorithm minimises, and it will feature in the next lesson.

What it is quietly assuming

Every step above encodes an assumption, and each one fails on some real dataset in a way that produces confident nonsense.

Clusters are spherical. Step 2 assigns by Euclidean distance to a centre, so the region belonging to each cluster is bounded by straight lines equidistant between centres. Two elongated, parallel groups get cut across rather than separated along their length. If your data has stretched or curved groups, k-means will not find them, and it will still return an answer.

Clusters are similar in size and density. The mean is pulled towards wherever the points are, so a large dense cluster next to a small sparse one tends to have its centre annexe part of the small one.

Every point belongs somewhere. There is no “none of these”. An outlier 40 standard deviations away is assigned to whichever cluster is nearest and drags that centre towards it.

k is known. The algorithm cannot tell you there are three groups when you asked for five. It will produce five.

Initialisation matters more than people expect

k-means finds a local minimum, and different starting centres reach different ones. The classic failure is two centres landing inside one real cluster while a genuine second cluster shares a third.

k-means++ initialisation (the default in scikit-learn) chooses starting centres that are spread out, by picking each new centre with probability proportional to its squared distance from the nearest existing one. It reduces bad outcomes substantially, and n_init=10 runs the whole thing ten times and keeps the best inertia, which reduces them further.

Do not lower n_init to save time unless you have measured that it does not matter for your data.

Scaling, again

k-means is a distance method, so a feature in rupees dominates one in years completely. Standardise before clustering, always.

There is a subtlety worth noticing. Standardising asserts that every feature should count equally in the distance, and that is itself a choice about what similarity means. If age matters twice as much as income for your purpose, you can say so by scaling age by two after standardising: clustering has no target to tell it what matters, so what matters is entirely encoded in how you scaled.

This is the deepest thing about unsupervised learning: the answer depends on your definition of similar, and you supply that definition whether or not you notice.

Variants worth knowing

• MiniBatchKMeans: updates centres from small random batches. Much faster on large data, slightly worse clusters. Above a few hundred thousand rows, use it.
• Gaussian mixture models: fit ellipse-shaped clusters with orientation, and give each point a probability of belonging to each cluster rather than a hard assignment. Strictly more flexible; k-means is the special case with spherical, equal-sized components and hard assignment.
• k-medoids: uses actual data points as centres rather than means, so the centre is a real example you can look at, and it tolerates outliers better.

What clusters are actually for

Honesty about the use case, because clustering is often run without one.

• Segmentation for marketing or product decisions: the most common use, and the one most likely to produce groups that get named and then treated as real.
• Compression: replace a row with its cluster id. Genuinely useful as a feature for a supervised model.
• Exploration: look at the centres to understand what kinds of rows exist. Frequently the best value.
• Labelling triage: cluster unlabelled data, label a few examples per cluster, and you have a rough labelled set for a fraction of the effort.

What clusters are not: a discovery of real categories in the world. k-means will partition uniformly random noise into five tidy groups and report a respectable inertia. The groups will be reproducible, describable and entirely artefactual. Before believing any clustering result, run it on shuffled data and see whether the output looks meaningfully different.