k-means

The photograph of Lake Issyk has a hundred and fifty thousand pixels and tens of thousands of different colours. Can we keep six so that the picture stays the same—the same dark water, the same pale sky? We need to find the six colours that best stand for all the others, and for that, to split the pixels into six groups of similar ones. That is what k-means does. On the right is the same photograph painted only with the colours found; below is the cloud of its pixels in the RGB cube, with the group centres crawling through it.

Press “New start” a few times. Sometimes the centres spread through the cloud neatly, and sometimes two of them share the same blue of the sky while all the shades of green go to one. The algorithm stops every time—but not always in the best place. Below we will see why that happens and what to do about it.

The idea

If the group centres are known, it makes sense to assign each point to the nearest centre. If the groups are known, it makes sense to put each centre at the mean of its points. We know neither—so we start with some centres and alternate the two steps while anything changes.

1. Assign. Each point gets the number of its nearest centre.
2. Update. Each centre moves to the arithmetic mean of its points.

This is Lloyd's algorithm—what is usually meant by k-means. The quality of a split is measured by its inertia: the sum of squared distances from points to the centres of their groups. The smaller it is, the tighter the groups sit around their centres.

Why it stops

Neither of the two steps increases the inertia. The first step only changes which centre a point belongs to, and picks the nearest—the squared distance does not grow. The second step moves a centre to the mean, and the mean is exactly the point with the smallest sum of squared distances to the points of the group (a property of the arithmetic mean, proved by expanding the square). So the inertia after each step is no larger than before.

There are finitely many ways to split n points into k groups, and while the split changes, the inertia strictly decreases—the algorithm never returns to a split. So sooner or later the split stops changing, and the algorithm stops. In practice that takes tens of steps; on the lake photograph, about twenty.

But stopping is not winning. The algorithm finds a local minimum: a split that neither of the two steps can improve. It does not look for a better split, and finding the best one is hard in general: the k-means problem is NP-hard.

The start decides everything

Here are twelve points in three clear groups in the corners. Let us start badly on purpose: all three centres in one corner.

step 1: inertia 171.05, centres (1.00, 1.00) (4.39, 7.24) (2.00, 1.20)
step 2: inertia 93.94, centres (1.00, 1.00) (4.99, 8.50) (2.00, 1.87)
labels [0, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1]

The algorithm honestly lowers the inertia and stops at 93.94: the two upper groups have merged into one, and the lower one is split into a point and three. No step changes anything any more. But start from three points in different corners and one step is enough, with an inertia of 9.85—almost ten times less.

So the start is chosen with care. The standard way is k-means++ (Arthur and Vassilvitskii, 2007): the first centre is a random point, and each next one is random too, but a point is the more likely the farther it is from the centres already chosen (the probability is proportional to the squared distance). The centres thus spread across different groups by themselves. The authors proved that the expected inertia after such a start is at most O(log k) times worse than the optimum. Turn on the k-means++ button in the scene and compare. In practice the algorithm is also run several times from different starts, keeping the best result.

Code

Both halves of a step exactly as defined. On equal distances a point goes to the centre with the lower number, and the centre of an empty group stays put—such conventions make the result reproducible.

And this is the palette of a real photograph—the one with the lake. The program reads the picture in PPM format (a simple format: a header, then R, G, B bytes in a row), splits its 9600 pixels into six colours and prints the palette with shares. Python and Go give the same colours to the last digit: the start is the same, and so is their floating-point arithmetic.

What it costs

The assign step compares each of the n points with each of the k centres, and a comparison costs O(d) for points of dimension d: O(n·k·d) in all. The update step is O(n·d). There are usually few steps, from a handful to several tens, although in the worst case exponentially many may be needed—but such examples have to be built on purpose.

For a photograph of 150,000 pixels and 6 colours a step is about a million operations, an instant. For millions of points in hundreds of dimensions it is no longer an instant, and then mini-batches are used: each step is computed on a random part of the points. The scene does exactly that—the centres are fitted on 12,000 random pixels, and the whole picture is repainted.

How many groups to take

The algorithm does not know how many groups there “really” are. The inertia always falls as k grows: at k = n it is zero, every point a group of its own. So people look for where the fall slows down—the so-called “elbow” of the curve. For a palette the question is simpler: k is how many colours you are willing to keep. Move the k slider: at two colours the lake becomes a poster, at six it is recognisable, at twelve it hardly differs from the original.

Where k-means goes wrong

• Groups that are not round. The algorithm cuts space into regions of the nearest centre—convex pieces. Elongated, curved or nested groups get cut wrongly. For them there is DBSCAN, which looks for dense regions of any shape.
• Feature scale. If one coordinate is height in centimetres and another age in years, height will dominate the distance almost entirely. Features are first brought to one scale.
• Colour is not RGB. Euclidean distance in RGB is not how different colours look to the eye. For palettes it is better to work in Lab, where distances are closer to perception.
• Outliers. One distant point drags the mean after it. With many outliers, medians are used instead of means (k-medians), or actual points instead of means (k-medoids).

History

The idea of splitting by nearest centres was proposed by Hugo Steinhaus in 1956. The algorithm with two alternating steps was invented by Stuart Lloyd at Bell Labs in 1957—to quantise a signal in pulse-code modulation, that is, exactly for the question “which levels to keep so that the signal is distorted least”. He published it only in 1982. The name “k-means” appeared in a 1967 paper by James MacQueen.

Problems

The first problem is half a step: assign pixels to the nearest palette colours. The second is the whole algorithm with exact rules that make the answer unique.