Clustering: the K-means objective, Lloyd's algorithm and why it stops

So far the unsupervised part of the course has been about representation learning: PCA in Part 3 and its kernel version in Part 5. The second paradigm of unsupervised learning is clustering, and it is the topic of this part and the next.

The motivating picture has a few bunches of points in the plane. PCA would find the eigen direction of largest variance, which runs through the bunches, and replace every point by its proxy on that line (proxies were introduced in Part 2). The line does hold most of the information. But the proxies are not spread evenly along it. They sit in clumps, so even inside the subspace PCA found there is more structure to uncover.

Question for this part: given data points, how do we uncover the groups they fall into?

The input is as usual \(\mathbf{x}_1, \dots, \mathbf{x}_n\), each in \(\mathbb{R}^d\). The goal is to partition them into \(K\) clusters. The lecture also calls them partitions, groups or simply boxes: every point must go into exactly one of \(K\) boxes.

There are many ways to do this. With five points and three boxes, one option is \(\{\mathbf{x}_1, \mathbf{x}_2, \mathbf{x}_5\}, \{\mathbf{x}_3\}, \{\mathbf{x}_4\}\). Another is \(\{\mathbf{x}_1, \mathbf{x}_4\}, \{\mathbf{x}_2\}, \{\mathbf{x}_3, \mathbf{x}_5\}\).

To count the options naively, let every point choose its box independently. Each of the five points has three choices, so there are

possibilities. This count includes splits that leave some box empty. Even if those are excluded, the number grows exponentially in the number of points.

Not needed for the course: of the 243, exactly 150 use all three boxes, and 25 remain if you also ignore which box carries which number.

With so many candidates we need a way to say how good each one is. That means a performance measure: a number attached to every partition of the given points.

Cluster indicators

Attach to every point a cluster indicator \(z_i \in \{1, \dots, K\}\) that says which box \(\mathbf{x}_i\) goes into. If the tenth point is put in the fourth box, then \(z_{10} = 4\). Once \(z_1, \dots, z_n\) are given, the partition is fully determined.

Which box gets which number is a choice; the table numbers boxes in the order they were listed.

How good is a partition? L13 · 8:04

We want a function \(F(z_1, \dots, z_n)\) that returns one number for each partition. It depends on the data too, but the data is fixed, so only the \(z\)'s are written. There are many possible choices. A natural one asks each cluster to be homogeneous: its points should look alike.

One way to measure how alike a cluster's points are is to see how spread out they are around the cluster's mean. Add up the squared distance of each point to that mean. If all points in a cluster are identical, each one is at distance zero; if they are far apart, the total is large. Doing this for every cluster gives the measure the lecture adopts:

Here \(\mathbb{I}(u)\) is the indicator: 1 if the statement \(u\) is true and 0 otherwise. The numerator of \(\boldsymbol{\mu}_k\) adds up exactly the points assigned to box \(k\), and the denominator counts them, so \(\boldsymbol{\mu}_k\) is the average of box \(k\). The distance is the squared Euclidean (\(\ell_2\)) distance, and \(\boldsymbol{\mu}_{z_i}\) is the mean of the box that \(\mathbf{x}_i\) is in. So \(F\) is the total squared distance from every point to the mean of its own box.

Standard names: \(F\) is often called the K-means objective or the within-cluster sum of squares.

The lecture's example

Take five points and \(K = 2\) with \(z = (1, 2, 1, 1, 2)\). Box 1 holds \(\mathbf{x}_1, \mathbf{x}_3, \mathbf{x}_4\) and box 2 holds \(\mathbf{x}_2, \mathbf{x}_5\). Then

The lecture keeps the points abstract. To see actual numbers, give them illustrative coordinates \(\mathbf{x}_1 = (1, 1)\), \(\mathbf{x}_2 = (5, 4)\), \(\mathbf{x}_3 = (2, 1)\), \(\mathbf{x}_4 = (1, 2)\), \(\mathbf{x}_5 = (6, 5)\). Then \(\boldsymbol{\mu}_1 = (4/3, 4/3)\) and \(\boldsymbol{\mu}_2 = (5.5, 4.5)\), and the five terms are \(2/9,\ 1/2,\ 5/9,\ 5/9,\ 1/2\), so \(F = 7/3 \approx 2.33\).

With a measure in place, the goal is precise: among all partitions, find the one with the smallest \(F\).

There are only finitely many partitions, so one could imagine a brute-force algorithm: compute \(F\) for every partition and keep the smallest. The trouble is how many there are. Each point has \(K\) choices, so there are \(K^n\) ways to put \(n\) points into \(K\) boxes. That is an upper bound, since it counts splits with empty boxes, but the true count also grows exponentially.

Computer scientists call this problem NP-hard. It is not expected that any algorithm solves it in time polynomial in the quantities of interest, here \(n\) and \(K\). A running time like \(n^2K^3\) would be acceptable; \(K^n\) is not. Even with \(K = 2\) and \(n = 1000\) points, brute force would examine \(2^{1000} \approx 10^{301}\) partitions.

So the plan changes. Instead of solving the problem exactly, the next lecture gives a very popular heuristic algorithm for it.

The algorithm is called Lloyd's algorithm, and it is often also called the K-means algorithm. The lecture points out that this second name is a slight misnomer. K-means is the name of the problem, finding \(K\) means that represent a partition, and Lloyd's algorithm is one heuristic way to attack it that need not solve it exactly.

It needs a starting partition \(z_1^0, \dots, z_n^0 \in \{1, \dots, K\}\), where the superscript is the iteration number. For now any starting partition will do; good ways to choose one are the topic of L17 in Part 7.

1. Lloyd's algorithm
2. Initialize. Pick any partition \(z_1^0, \dots, z_n^0\). Set \(t = 0\).
3. Compute the means. For every \(k\), \(\displaystyle \boldsymbol{\mu}_k^t = \frac{\sum_{i=1}^n \mathbf{x}_i\,\mathbb{I}(z_i^t = k)}{\sum_{i=1}^n \mathbb{I}(z_i^t = k)}\).
4. Reassign. For every \(i\), \(\displaystyle z_i^{t+1} = \arg\min_{k} \big\lVert \mathbf{x}_i - \boldsymbol{\mu}_k^t \big\rVert^2\). If the current cluster's mean is among the closest, keep \(z_i^{t+1} = z_i^t\).
5. Check. If no point changed cluster, stop: the algorithm has converged. Otherwise set \(t \leftarrow t + 1\) and go back to computing the means.

The crux is the reassignment step. Every point compares its squared distance to the mean of its own box with its squared distances to the means of all the other boxes. If some other mean is strictly closer, the point jumps to that box in the next round. If the closest other mean is exactly as close as its own, it does not jump. 5:28

Each round moves the algorithm from one partition to another. It has converged when it reaches a partition in which no point wants to jump: no point has another box whose mean is strictly closer than the mean of its own box. From then on nothing changes.

"Repeat until convergence" quietly assumes that this moment comes. It might not. The algorithm could go from one partition to a second, then a third, and then back to the first. It would then cycle forever, always changing partitions and never converging. The hope behind the algorithm is that each jump leads to a better partition, but a hope needs an argument.

The lecture states the facts now and proves the first one in Lecture 15:

• Lloyd's algorithm always converges.
• The converged partition need not be optimal. It may not solve the K-means problem. That is expected, since the problem is NP-hard, and it is something we have to live with.
• In practice it works well. It is a very popular algorithm and it produces reasonable clusters. 10:06

Having the algorithm raises questions about where it works well and why. The lecture lists the four key ones, to be answered one by one:

1. Convergence. Does the algorithm converge? Answered in this part, below.
2. Nature of the clusters. If it converges, what kind of clusters does it produce? (L16, Part 7.)
3. Initialization. The algorithm does not say how to choose \(z^0\). How should we? (L17, Part 7.)
4. Choice of K. So far \(K\) is given. In a genuinely unsupervised setting nobody tells us how many natural groups the data has. (L18, Part 7.)

On the last question, Lecture 15 gives an example where \(K\) is natural. A teacher who wants to turn exam marks into grades already knows the buckets, say S, A, B, C and D, so the number of grades fixes \(K\). In other problems we are only handed the points, and finding a good \(K\) is part of the job.

Lecture 15 answers the first question: yes, Lloyd's algorithm converges. The lecture adds that understanding the argument also hints at what the final partitions look like, which Part 7 takes up. The proof rests on one simple fact about points and distances.

Take any points \(\mathbf{x}_1, \dots, \mathbf{x}_l \in \mathbb{R}^d\) and look for the single vector \(\mathbf{v}\) that is closest to all of them in the sense of total squared distance.

Using the average \(\frac1l\sum_i\) instead of the sum would not change the answer, since the average is the sum scaled by a constant. The lecture works with the sum because \(F\) was defined as a sum.

The lecture leaves the derivation as an exercise, recommended especially if you are not used to differentiating with respect to a vector. Treat the objective as a function of \(\mathbf{v}\), take its gradient (one partial derivative per component) and set it to zero:

A quick check that this point is a minimum, which the lecture does not do: write \(\bar{\mathbf{x}}\) for the mean and split \(\mathbf{x}_i - \mathbf{v} = (\mathbf{x}_i - \bar{\mathbf{x}}) + (\bar{\mathbf{x}} - \mathbf{v})\). Expanding the square, the cross terms add up to \(2(\bar{\mathbf{x}} - \mathbf{v})^\top \sum_i (\mathbf{x}_i - \bar{\mathbf{x}})\), which is zero because the deviations from the mean sum to zero. So

Any \(\mathbf{v}\) other than the mean pays exactly \(l\,\lVert\bar{\mathbf{x}} - \mathbf{v}\rVert^2\) extra.

Look at some iteration \(t\). The current assignment is \(z_1^t, \dots, z_n^t\) (superscript: iteration, subscript: point), each a value in \(\{1, \dots, K\}\), and \(\boldsymbol{\mu}_k^t\) is the mean of cluster \(k\) under this assignment.

If no point is reassigned at step \(t\), the algorithm has converged and there is nothing to prove. So assume it has not converged. Then at least one point is unhappy with its own mean: it has found another box whose mean is strictly closer. The reassignment step produces \(z_1^{t+1}, \dots, z_n^{t+1}\), and some points have moved. The question is whether this is a better partition, measured by \(F\).

The argument compares three sums of squared distances. They differ only in which box and which mean each point is measured against.

The proof shows \(F(z^{t+1}) \le \text{intermediate} < F(z^t)\), one inequality at a time.

The claim is

The right side is the current objective: since \(z_i^t\) is by definition the box of \(\mathbf{x}_i\) at round \(t\), it adds up each point's squared distance to its own mean. The left side measures each point against the mean of the box it wants to go to, still using the round-\(t\) means \(\boldsymbol{\mu}_1^t, \dots, \boldsymbol{\mu}_K^t\). Compare the two sums term by term.

1. A point that stays has \(z_i^{t+1} = z_i^t\). Both sides measure it against the same mean, so its two terms are equal.
2. A point that moves jumped only because the mean of its new box is strictly closer than its own mean. Its term on the left is strictly smaller than its term on the right.
3. No term on the left is larger than its partner on the right, since \(z_i^{t+1}\) picks the closest round-\(t\) mean.
4. Because the algorithm has not converged, at least one point moves. So the left sum is strictly smaller than the right sum.

This inequality holds simply by the algorithm's choice: it jumps only on a strict improvement.

The left side is not the objective of the next partition. The points have only announced where they want to go. Once they all move, each box has new members and therefore a new mean, and \(F(z^{t+1})\) measures every point against that new mean. The left side is an intermediate quantity that sits between the two objectives.

The second claim compares the next objective with the intermediate quantity:

Focus on a single cluster \(k\). At round \(t\) it had some members and the mean \(\boldsymbol{\mu}_k^t\). During reassignment, some members are happy and stay. Some find a closer mean elsewhere and leave. Some points from other clusters find \(\boldsymbol{\mu}_k^t\) closer than their own mean and arrive. Call the resulting set of members \(C_k = \{\, i : z_i^{t+1} = k \,\}\).

• In the intermediate quantity, every \(i \in C_k\) contributes \(\lVert\mathbf{x}_i - \boldsymbol{\mu}_k^t\rVert^2\), because \(z_i^{t+1} = k\) means \(\boldsymbol{\mu}^t_{z_i^{t+1}} = \boldsymbol{\mu}^t_k\). Each such point either stayed after comparing itself with \(\boldsymbol{\mu}_k^t\), or compared itself with \(\boldsymbol{\mu}_k^t\) and jumped in.
• In \(F(z^{t+1})\), every \(i \in C_k\) contributes \(\lVert\mathbf{x}_i - \boldsymbol{\mu}_k^{t+1}\rVert^2\), where \(\boldsymbol{\mu}_k^{t+1}\) is the mean of exactly these points.

So both sides measure the same set of points, once against an arbitrary vector \(\boldsymbol{\mu}_k^t\) and once against their own mean. Fact 1 says their own mean always wins. Written out, step by step:

A box that ends up with no members contributes nothing to either side. The lecture does not discuss that case.

The lecture's picture of this step is that after each round of moves, the points of every cluster are, on the whole, drawn closer together around their centre.

Put the two steps side by side. The first sum below is \(F(z^t)\) and the last is \(F(z^{t+1})\); the intermediate quantity sits between them:

Strict decrease alone does not prove convergence, since in general a quantity can keep decreasing forever. Here each value belongs to a partition, and three observations finish the proof.

1. No partition can repeat. Every round with a reassignment gives a partition with strictly smaller \(F\) than all the partitions visited before it. A repeated partition would need the same value of \(F\), which is impossible.
2. There are finitely many partitions. At most \(K^n\). The number is huge but finite.
3. So the reassignments must stop. Think of crossing a partition off the list each time it is visited. A finite list runs out, so eventually a round comes with no reassignment, which is exactly convergence.

The finiteness of the set of partitions is what turns "the objective keeps going down" into "the algorithm stops".

To summarize: every reassignment strictly lowers the sum of squared distances from points to the means of their clusters, so at some point every point is happy with its own cluster and nothing changes any more. The lecture is careful about two things this does not claim.

• It does not say how many rounds. The finiteness argument would allow on the order of \(K^n\) rounds in the worst case, but it does not say the algorithm needs that many. In practice Lloyd's algorithm converges very fast on most data sets. The argument only guarantees that it stops.
• It does not say the result is good. Converging does not mean reaching the smallest objective over all partitions. The algorithm can stop at a partition where everybody is happy and nothing changes, yet some other partition has a smaller \(F\). In this sense it can converge to a local minimum.

What the argument does give is important: the algorithm can never get stuck in a loop of partitions, so "repeat until convergence" is a valid stopping rule. Without it we would not know whether the algorithm ever stops.

The next lecture studies the nature of the clusters that Lloyd's algorithm produces. That, initialization and the choice of \(K\) are in Part 7.

Clustering

Unsupervised task of partitioning data points into \(K\) groups of similar points.

Partition (cluster, box)

A split of the points in which every point belongs to exactly one of \(K\) groups.

Cluster indicator \(z_i\)

The number in \(\{1, \dots, K\}\) of the box that \(\mathbf{x}_i\) is assigned to; \(z_1, \dots, z_n\) determine the partition.

Cluster mean \(\boldsymbol{\mu}_k\)

Average of the points with \(z_i = k\). Also called the centroid of cluster \(k\).

K-means objective \(F\)

\(\sum_i \lVert\mathbf{x}_i - \boldsymbol{\mu}_{z_i}\rVert^2\): total squared distance from each point to its own cluster mean.

K-means problem

Minimizing \(F\) over all \(K^n\) assignments; NP-hard.

NP-hard

No algorithm is expected to solve the problem in time polynomial in \(n\) and \(K\).

Lloyd's algorithm

Heuristic for the K-means problem: alternate between computing cluster means and reassigning points to the nearest mean.

Reassignment step

\(z_i^{t+1} = \arg\min_k \lVert\mathbf{x}_i - \boldsymbol{\mu}_k^t\rVert^2\); a point moves only if another mean is strictly closer.

Convergence

Reaching a partition where no point wants to change cluster, after which nothing changes.

Fact 1

The mean of a set of points minimizes the sum of squared distances to them.

Intermediate quantity

\(\sum_i \lVert\mathbf{x}_i - \boldsymbol{\mu}^t_{z_i^{t+1}}\rVert^2\): new boxes, old means. It lies between \(F(z^{t+1})\) and \(F(z^t)\).

Local minimum

A converged partition that no single round can improve, but whose \(F\) is not the smallest possible.