Clustering and the K-Means Objective

PCA describes a data set by the directions along which it spreads. Sometimes the most important structure is not a direction but a set of groups. Picture three tight clumps of points lying roughly along a line. PCA finds the line and projects every point onto it, but the projections still sit in three clumps, and PCA has nothing to say about them. Finding those clumps is clustering.

A partition and its score

We have $n$ points $x_1, \dots, x_n \in \mathbb{R}^d$ and want to split them into $K$ groups. Describe a split by giving every point a cluster indicator $z_i \in \{1, \dots, K\}$: $z_i = 3$ means point $i$ goes in box 3. Any choice of $z_1, \dots, z_n$ is a partition.

There are a great many partitions: each point has $K$ choices, so up to $K^n$ of them. To pick one we need to score them. A natural requirement is that each cluster should be tight: its points should sit close to one another. Measure that by each point's squared distance from the mean of its own cluster. The mean of cluster $k$ is

the average of the points assigned to it (the indicator $\mathbf{1}[z_i = k]$ is 1 when point $i$ is in cluster $k$ and 0 otherwise). The score of the partition is the total squared distance from every point to its own cluster's mean:

Smaller is better. This is the k-means objective (also called the within-cluster sum of squares, or SSE).

A small example

Take five points on a number line, $x = 1, 2, 3, 10, 11$, and $K = 2$.

• Partition A: $\{1, 2, 3\}$ and $\{10, 11\}$. The means are 2 and 10.5. The score is $(1 + 0 + 1) + (0.25 + 0.25) = 2.5$.
• Partition B: $\{1, 2\}$ and $\{3, 10, 11\}$. The means are 1.5 and 8. The score is $(0.25 + 0.25) + (25 + 4 + 9) = 38.5$.

Partition A is far better, matching what the eye sees.

Why not just try every partition?

The goal is to find the partition with the smallest $F$. With only finitely many partitions, we could in principle score each one. But there are about $K^n$ of them: with $K = 2$ and $n = 1{,}000$ points, around $2^{1000}$, a number with more than 300 digits. Exact k-means is NP-hard, so no algorithm is expected to solve every instance in time polynomial in $n$ and $K$. We settle for a fast heuristic that finds a good partition, usually not provably the best.

Lloyd's algorithm

The standard heuristic is Lloyd's algorithm, usually just called k-means (strictly, k-means is the problem and Lloyd's algorithm is one way of attacking it).

1. Initialise: put every point in some cluster, giving $z^{(0)}_1, \dots, z^{(0)}_n$. (How to do this well is the subject of a later chapter.)
2. Repeat until nothing changes:
   • Update the means: for each cluster $k$, set $\mu^{(t)}_k$ to the mean of the points currently assigned to it.
   • Reassign: move every point to the cluster whose mean is nearest, $z^{(t+1)}_i = \arg\min_{k} \big\lVert x_i - \mu^{(t)}_k \big\rVert^2 ,$ keeping it where it is if its current mean is already among the nearest (so that ties never cause pointless moves).

When a full pass moves no point, every point is already closest to its own cluster's mean, and the algorithm stops: it has converged.

Four questions

Stated this simply, the algorithm raises four questions, which the next two chapters answer.

1. Does it always stop? Could points keep jumping between clusters forever?
2. What shape are the clusters it finds?
3. How should it be started? Different starting partitions can lead to different answers.
4. How do we choose K? Sometimes it is given (a teacher sorting students into five grades knows K = 5); usually it is not.