← All notes · ML revision (DA5400)

L15: Convergence of the K-means (Lloyd's) Algorithm

Source: https://www.youtube.com/watch?v=5geXmawaImk (≈33 min) Previous: L13 · L14

L14's "Going further" section previewed this proof in compact form. This lecture is the official version: a three-quantity sandwich argument. Section 7 extends it in d dimensions: deriving Fact 1, the exact drop per iteration, counting partitions, how many iterations it takes, and a concrete local-minimum example.

Timestamped outline

1. When K is known vs. unknown

• Known K: a teacher grouping students into grades S, A, B, C, D has K = 5 fixed in advance.
• Unknown K: you're just given points and need to find natural groupings. Choosing K is a separate question, covered later.

2. Fact 1: the mean minimises the sum of squared distances

For points x₁, …, x_ℓ ∈ ℝᵈ:

$$ v^\star=\arg\min_{v\in\mathbb{R}^d}\sum_{i=1}^{\ell}\lVert x_i-v\rVert^2=\frac{1}{\ell}\sum_{i=1}^{\ell}x_i . $$

• Minimising the sum and minimising the average give the same v, since one is a positive multiple of the other.
• The lecture leaves the derivation as an exercise (take the gradient with respect to v and set it to 0). It is done in full in §7.1.

3. Proof setup

• At iteration t: assignments zᵗ = (z₁ᵗ, …, zₙᵗ), and means μₖᵗ of each cluster under zᵗ.
• Objective: F(z₁, …, zₙ) = Σᵢ ‖xᵢ − μ_{zᵢ}‖², where the μ's are the means of that partition.
• If the algorithm has already converged at t, there is nothing to prove. So assume at least one point moves, giving zᵗ⁺¹ ≠ zᵗ.

4. The three quantities

I is not an objective value. It mixes the new assignment with the old means. It is only a stepping stone.

5. The sandwich: L ≤ I < R

Step A: I < R (the algorithm's choice).

• A point that stays (zᵢᵗ⁺¹ = zᵢᵗ) contributes exactly the same term to I and to R.
• A point that moves does so only because the new mean is strictly closer: ‖xᵢ − μᵗ_{zᵢᵗ⁺¹}‖² < ‖xᵢ − μᵗ_{zᵢᵗ}‖².
• At least one point moves, so I < R strictly.

Step B: L ≤ I (Fact 1, one cluster at a time).

• Group both sums by the new clusters Cₖᵗ⁺¹ = {i : zᵢᵗ⁺¹ = k}:

$$ L=\sum_{k=1}^{K}\sum_{i\in C_k^{t+1}}\lVert x_i-\mu_k^{t+1}\rVert^2, \qquad I=\sum_{k=1}^{K}\sum_{i\in C_k^{t+1}}\lVert x_i-\mu_k^{t}\rVert^2 . $$

• Every point in Cₖᵗ⁺¹ is in I measured against the same vector μₖᵗ. Some points stayed and compared themselves to μₖᵗ; others jumped in because μₖᵗ was closer.
• μₖᵗ⁺¹ is the mean of exactly the set Cₖᵗ⁺¹. By Fact 1 (with v = μₖᵗ), the inner sum for L is ≤ the inner sum for I, for every k. Adding over k gives L ≤ I.

Putting it together:

$$ F(z^{t+1}) \;\le\; I \;<\; F(z^{t}) \qquad\text{whenever a reassignment happens.} $$

Intuition: points that move pay less right away (Step A), and then the means re-centre, which makes things better still (Step B).

6. Why strict decrease means the algorithm converges

• A strictly decreasing F means no partition can ever repeat, since a repeat would give the same F value.
• There are only finitely many partitions (at most Kⁿ).
• Each reassignment crosses one partition off the list. Eventually reassignments must stop, so the algorithm converges. ∎

Two caveats:

1. The proof gives no useful bound on the number of rounds. Kⁿ is only a worst-case ceiling. In practice Lloyd's converges quickly.
2. Convergence does not mean optimality. It reaches a partition where every point is happy (a local minimum), which may be far from the global minimum.

7. Going further: d-dimensional derivations and extensions

This section goes beyond the video.

7.1 Deriving Fact 1 in ℝᵈ (the exercise the lecture skipped)

Let J(v) = Σᵢ ‖xᵢ − v‖² = Σᵢ (xᵢ − v)ᵀ(xᵢ − v).

Method 1: one coordinate at a time. Write v = (v₁, …, v_d). Then

$$ J(v)=\sum_{i=1}^{\ell}\sum_{j=1}^{d}(x_{ij}-v_j)^2 \;\Rightarrow\; \frac{\partial J}{\partial v_j}=-2\sum_{i=1}^{\ell}(x_{ij}-v_j)=0 \;\Rightarrow\; v_j=\frac1\ell\sum_i x_{ij}. $$

J separates into d independent 1-D problems. Each is solved by the 1-D mean, so the d-dimensional answer is just the vector of coordinate means.

Method 2: vector calculus. Expand: J(v) = Σᵢ‖xᵢ‖² − 2vᵀ(Σᵢxᵢ) + ℓ‖v‖². Use

• ∇_v(aᵀv) = a
• ∇_v(vᵀv) = 2v

to get

$$ \nabla_v J = -2\sum_i x_i + 2\ell v \overset{!}{=} 0 \;\Rightarrow\; v^\star=\bar x . $$

Is it a minimum? The Hessian is ∇²J = 2ℓ·I_d ≻ 0, so J is strictly convex and v* is the unique global minimiser.

Method 3: no calculus (completing the square). Write xᵢ − v = (xᵢ − x̄) + (x̄ − v) and expand:

$$ \sum_i\lVert x_i-v\rVert^2=\sum_i\lVert x_i-\bar x\rVert^2+2(\bar x-v)^\top\underbrace{\sum_i(x_i-\bar x)}_{=0}+\ell\lVert\bar x-v\rVert^2 . $$

$$ \boxed{\;\sum_i\lVert x_i-v\rVert^2=\sum_i\lVert x_i-\bar x\rVert^2+\ell\,\lVert\bar x-v\rVert^2\;} $$

The second term is ≥ 0 and equals 0 only when v = x̄. This proves Fact 1 and says exactly how much worse any other v is. That is the key to §7.2.

7.2 The exact drop per iteration

Apply the boxed identity to each new cluster, with v = μₖᵗ:

$$ I - L = \sum_{k=1}^{K} n_k^{t+1}\,\lVert \mu_k^{t+1}-\mu_k^{t}\rVert^2 \quad(\ge 0). $$

Add the gain from Step A:

$$ F(z^t)-F(z^{t+1})=\underbrace{\sum_{i:\,\text{moved}}\Big(\lVert x_i-\mu^t_{z_i^t}\rVert^2-\lVert x_i-\mu^t_{z_i^{t+1}}\rVert^2\Big)}_{\text{reassignment gain} \;>\;0} +\underbrace{\sum_k n_k^{t+1}\lVert\mu_k^{t+1}-\mu_k^t\rVert^2}_{\text{re-centring gain}\;\ge 0}. $$

• The reassignment gain comes from points getting closer to a centre.
• The re-centring gain comes from centres moving. It is weighted by cluster size.
• The algorithm stalls exactly when both are zero: no point wants to move, and the means are fixed.

7.3 How many partitions are there, really?

Kⁿ counts labelled assignments, including ones that leave boxes empty. The number of ways to split n points into exactly K non-empty, unlabelled clusters is the Stirling number of the second kind:

$$ S(n,K)=\frac{1}{K!}\sum_{j=0}^{K}(-1)^j\binom{K}{j}(K-j)^n \;\approx\; \frac{K^n}{K!}\ \text{for large }n. $$

Example: S(5,3) = 25, compared with 3⁵ = 243. That is still exponential in n, so the lecture's conclusion holds.

7.4 How many iterations? (what's known)

• Worst case: there are inputs in the plane (d = 2) where Lloyd's takes 2^Ω(n) iterations (Vattani, 2011). So the finiteness argument is essentially as good as you can do in the worst case.
• Smoothed analysis: if the input is slightly perturbed by random noise, the expected number of iterations is polynomial in n (Arthur, Manthey & Röglin, 2009). This explains why it converges fast in practice.
• In practice: implementations stop when the drop F(zᵗ) − F(zᵗ⁺¹) falls below a tolerance, or after a maximum number of iterations.

7.5 A concrete bad local minimum

Take four points at the corners of a wide rectangle, with K = 2:

$$ x_1=(0,0),\;x_2=(0,1),\;x_3=(4,0),\;x_4=(4,1). $$

Why bottom vs. top is stuck: check x₁ = (0,0).

• Distance² to its own mean (2,0) is 4.
• Distance² to the other mean (2,1) is 4 + 1 = 5.

So x₁ stays, and by symmetry so does every other point. Lloyd's algorithm has converged at 16 times the optimal cost. This is why initialization matters, which is the next big topic.

7.6 Edge case: empty clusters

If every point leaves cluster k, then μₖᵗ⁺¹ is 0/0, which is undefined. The proof still works if you either drop the cluster (F can only stay the same) or re-seed it with the point farthest from its centre (F strictly drops). Libraries do one of these silently.

7.7 Why the proof works: a broader view

The proof only used two facts:

1. Assigning each point to its nearest centre minimises F when the centres are fixed.
2. The mean minimises the cost when the clusters are fixed.

Any distance for which some "centre" minimises the within-cluster cost gives an algorithm that converges in the same way. Examples:

• K-medians: ℓ₁ distance, and the centre is the coordinate-wise median.
• K-medoids: the centre must be one of the data points.
• Bregman divergences (e.g. KL divergence): the mean is still the optimal centre (Banerjee et al., 2005).

The general pattern is alternating minimisation (block coordinate descent). The EM algorithm for Gaussian mixtures, coming later in the course, follows the same template. That makes K-means a "hard-assignment" version of EM.

7.8 Seeing it numerically (any d)

import numpy as np
rng = np.random.default_rng(0)
X = np.vstack([rng.normal(c, 1.0, (100, 10)) for c in (0, 4, 8)])  # 300 points in R^10
K, n = 3, len(X)
z = rng.integers(K, size=n)
F = lambda z, M: ((X - M[z])**2).sum()
while True:
    M = np.stack([X[z == k].mean(0) for k in range(K)])
    D = ((X[:, None, :] - M[None])**2).sum(-1)
    best = D.argmin(1)
    z_new = np.where(D[np.arange(n), z] <= D[np.arange(n), best], z, best)
    print(f"F = {F(z, M):.2f}")
    if np.array_equal(z_new, z): break
    z = z_new

The printed values strictly decrease, as the proof says.

Key takeaways

1. Fact 1: the mean is the unique minimiser of the sum of squared distances in any dimension. Moving the centre by Δ costs exactly ℓ‖Δ‖² extra.
2. The sandwich: F(zᵗ⁺¹) ≤ I < F(zᵗ). Step A uses the algorithm's strict-move rule; Step B uses Fact 1 on each cluster.
3. Strict decrease plus finitely many partitions means guaranteed convergence.
4. The proof gives no useful iteration bound (the worst case is exponential), but convergence is fast in practice.
5. Convergence is to a local optimum, which can be arbitrarily worse than the global one (e.g. the rectangle example, 16× worse).

Self-check questions

1. In the proof, why is the intermediate quantity I not an objective value?
2. Which step uses the "strictly closer" rule, and which uses Fact 1?
3. Using the completing-the-square identity, compute Σ‖xᵢ − v‖² for the points (0,0), (2,0), (1,3) and v = (0,0), without summing directly. (Mean (1,1). Spread around the mean = 2 + 2 + 4 = 8. Shift penalty = 3·‖(1,1)‖² = 6. Total = 14.)
4. Why doesn't "F strictly decreases" alone guarantee stopping for a continuous problem, and what extra fact rescues the K-means case?
5. In the rectangle example, how wide must the rectangle be for the bottom-vs-top split to be stuck? (For width w and height 1: distance² to own mean is (w/2)², to the other is (w/2)² + 1. So it is stuck for every w > 0, and it is worse than left-vs-right whenever w > 1.)
6. Is S(n,K) or Kⁿ the right count of partitions for the convergence argument? Does it matter?