← All notes · ML revision (DA5400)

L14: K-means Clustering (Lloyd's Algorithm)

Source: https://www.youtube.com/watch?v=QrXQ5cMzphA (≈12 min) Previous: L13 notes, which set up the objective and showed it is NP-hard.

Scope note: this video presents the algorithm, says that it converges (proof deferred), and then lists four open questions: convergence, nature of the clusters, initialization and choice of K. It does not answer them yet; later lectures do. Section 6 ("Going further") derives the answers in full d-dimensional form so you have them ahead of time.

Timestamped outline

1. Naming

• K-means problem: minimise Σᵢ ‖xᵢ − μ_{zᵢ}‖² over all partitions (NP-hard, see L13).
• Lloyd's algorithm: a heuristic for that problem. People often call it "the K-means algorithm", which mixes up the problem with one method for solving it.

2. The algorithm

Step 0: Initialization. Pick a starting assignment z⁰ = (z⁰₁, …, z⁰ₙ), with each z⁰ᵢ ∈ {1..K}. The superscript is the iteration number. How to initialise is covered later.

Repeat until convergence:

Step 1: Compute means. For each k = 1..K:

$$ \mu_k^{t} = \frac{\sum_{i=1}^{n} \mathbb{1}(z_i^{t}=k)\,x_i}{\sum_{i=1}^{n} \mathbb{1}(z_i^{t}=k)} $$

Step 2: Reassignment. For each point i:

$$ z_i^{t+1} = \arg\min_{k\in\{1..K\}} \lVert x_i - \mu_k^{t} \rVert_2^2 $$

Tie rule: if the point's current cluster is among the closest, it stays. A point moves only when another mean is strictly closer. This rule matters for the convergence proof, because it stops the algorithm flipping back and forth between equal options.

Intuition: each box has a centre. Every point checks every centre and moves to the nearest one. Then the centres are recomputed.

3. Convergence: what it means

• Converged: no point wants to switch boxes, so z^{t+1} = z^t. Every point is already closest to its own cluster's mean.
• Worry: could the partitions cycle, A → B → C → A → …, and never stop?
• The lecture's claim: Lloyd's algorithm does converge (proof later). See §6.2 for the proof.

4. Caveats

• What it converges to may not be the optimal partition. You can't expect an efficient exact solver for an NP-hard problem.
• In practice it is very popular and gives reasonable clusters.

5. Open questions (answered in later lectures)

1. Convergence: does it always converge? (Yes, see §6.2.)
2. Nature of clusters: what shapes can it find? (Convex Voronoi cells, see §6.3.)
3. Initialization: how should z⁰ or μ⁰ be chosen? (see §6.5)
4. Choice of K: how many clusters? (see §6.6)

6. Going further: general d-dimensional derivations

This section goes beyond the video.

6.1 Why the mean? (Step 1 is optimal, for any dimension d)

Fix the assignments. For a cluster Cₖ with nₖ points, find the μ ∈ ℝᵈ that minimises

$$ J_k(\mu)=\sum_{i\in C_k}\lVert x_i-\mu\rVert^2=\sum_{i\in C_k}(x_i-\mu)^\top(x_i-\mu). $$

Expand: J = Σ xᵢᵀxᵢ − 2μᵀΣxᵢ + nₖ μᵀμ. Use the vector-calculus identities ∇_μ(aᵀμ) = a and ∇_μ(μᵀμ) = 2μ:

$$ \nabla_\mu J_k = -2\sum_{i\in C_k}x_i + 2n_k\mu = 0 \;\Rightarrow\; \mu^\star=\frac{1}{n_k}\sum_{i\in C_k}x_i . $$

The Hessian is ∇²J = 2nₖ I_d, which is positive definite. So J is strictly convex and μ* is the unique global minimiser. The calculation is the same in 1-D or 1000-D; it just runs coordinate by coordinate.

A useful identity (bias–variance split): for any μ,

$$ \sum_{i\in C_k}\lVert x_i-\mu\rVert^2=\sum_{i\in C_k}\lVert x_i-\bar x_k\rVert^2 + n_k\lVert \bar x_k-\mu\rVert^2 . $$

The cross term vanishes because Σ(xᵢ − x̄ₖ) = 0. So moving μ away from the mean adds exactly nₖ‖x̄ₖ − μ‖² to the cost.

6.2 Proof that Lloyd's converges

Treat the objective as a function of both the assignments and the centres:

$$ F(z,\mu)=\sum_{i=1}^n \lVert x_i-\mu_{z_i}\rVert^2 . $$

Lloyd's algorithm is coordinate descent on F:

• Step 1 (z fixed, optimise μ): by §6.1, F(zᵗ, μᵗ) ≤ F(zᵗ, μ) for every μ. In particular F(zᵗ, μᵗ) ≤ F(zᵗ, μ^{t−1}).
• Step 2 (μ fixed, optimise z): the terms for different points are independent, so picking each zᵢ by argmin minimises the sum. This gives F(z^{t+1}, μᵗ) ≤ F(zᵗ, μᵗ), with strict inequality whenever at least one point moves, thanks to the tie rule.

Chaining the two steps:

$$ F(z^{t+1},\mu^{t+1}) \le F(z^{t+1},\mu^{t}) < F(z^{t},\mu^{t}) \quad\text{whenever some point moves.} $$

1. F strictly decreases every time the partition changes.
2. So no partition can repeat: a repeat would give the same F value, which contradicts strict decrease. That rules out cycles.
3. There are at most Kⁿ partitions, which is a finite number.
4. So the algorithm stops after finitely many iterations. ∎

It reaches a local optimum: no single step (re-centring, or moving one point) can lower F. That is still not the global optimum.

6.3 What shape are the clusters? (Voronoi / linear boundaries)

Point x goes to cluster k rather than l when ‖x − μₖ‖² < ‖x − μₗ‖². Expand both sides. The ‖x‖² terms cancel:

$$ -2\mu_k^\top x+\lVert\mu_k\rVert^2 < -2\mu_l^\top x+\lVert\mu_l\rVert^2 \;\Longleftrightarrow\; 2(\mu_l-\mu_k)^\top x < \lVert\mu_l\rVert^2-\lVert\mu_k\rVert^2 . $$

This is a linear inequality in x. The boundary between two clusters is a hyperplane in ℝᵈ. It is the perpendicular bisector of the segment μₖμₗ: normal vector (μₗ − μₖ), passing through the midpoint (μₖ + μₗ)/2.

Cluster k is the intersection of K − 1 half-spaces, so it is a convex polyhedron. The set of all these regions is the Voronoi partition of ℝᵈ generated by the means.

Consequence: K-means cannot find non-convex clusters such as two concentric rings or two interlocking half-moons. (Kernel K-means and spectral clustering fix this, in the same way kernel PCA extends PCA.)

6.4 Matrix form and cost per iteration

Stack the data as rows of X ∈ ℝ^{n×d}. Let Z ∈ {0,1}^{n×K} be the one-hot assignment matrix (Zᵢₖ = 1 iff zᵢ = k), and let M ∈ ℝ^{K×d} hold the means as rows. Then:

$$ F = \lVert X - ZM\rVert_F^2,\qquad \text{Step 1: } M = (Z^\top Z)^{-1}Z^\top X . $$

ZᵀZ = diag(n₁, …, n_K) holds the cluster sizes, and ZᵀX holds the per-cluster sums. So Step 1 is exactly the least-squares solution for M with Z fixed. This is a nice link to linear regression.

Distances for Step 2, all at once: ‖xᵢ − μₖ‖² = ‖xᵢ‖² − 2xᵢᵀμₖ + ‖μₖ‖², i.e. D = diag(XXᵀ)𝟏ᵀ − 2XMᵀ + 𝟏 diag(MMᵀ)ᵀ.

Cost per iteration: O(nKd) for the distances plus O(nd) for the means. That is linear in the data size, which is why K-means scales so well.

6.5 Initialization (preview)

• Random partition or random K data points as centres. Simple, but it can land in a bad local optimum.

• K-means++ (Arthur & Vassilvitskii, 2007):

  1. Pick the first centre uniformly at random from the data.
  2. Pick each later centre with probability ∝ D(x)², the squared distance from x to the nearest centre chosen so far.

  This gives an expected cost of at most O(log K) times the optimum.

• In practice: run several random restarts and keep the run with the lowest F.

6.6 Choosing K (preview)

• F*(K) always decreases as K grows (F* = 0 when K = n). So you cannot simply minimise F over K.
• Elbow method: plot F*(K) against K and look for the "knee" where the gains flatten out.
• Penalised objective: minimise F*(K) + λ·K, a trade-off between fit and complexity (AIC/BIC-style).
• Other options: silhouette score, gap statistic.

6.7 Minimal NumPy implementation (any d)

import numpy as np

def lloyd(X, K, iters=100, seed=0):
    rng = np.random.default_rng(seed)
    n, d = X.shape
    z = rng.integers(K, size=n)                    # Step 0: random partition
    for _ in range(iters):
        # Step 1: means (re-seed empty clusters with a random point)
        M = np.stack([X[z == k].mean(0) if np.any(z == k) else X[rng.integers(n)]
                      for k in range(K)])
        # Step 2: squared distances n×K, via ||x||² - 2xᵀμ + ||μ||²
        D = (X**2).sum(1)[:, None] - 2 * X @ M.T + (M**2).sum(1)[None, :]
        best = D.argmin(1)
        # tie rule: stay unless strictly closer
        stay = D[np.arange(n), z] <= D[np.arange(n), best]
        z_new = np.where(stay, z, best)
        if np.array_equal(z_new, z):               # converged
            break
        z = z_new
    return z, M, ((X - M[z])**2).sum()

Key takeaways

1. Lloyd's algorithm alternates two steps: compute the means, then reassign each point to its nearest mean.
2. It is coordinate descent on F(z, μ). F strictly decreases, so no partition repeats, and the algorithm converges in finitely many steps.
3. It converges to a local, not global, optimum. Initialization matters.
4. Clusters are convex Voronoi cells with hyperplane boundaries in ℝᵈ.
5. Each iteration costs O(nKd), so it scales to large data.

Self-check questions

1. Why does the tie rule ("move only if strictly closer") matter for the convergence proof?
2. Show that the boundary between two centres in ℝ³ is a plane, and find its normal vector.
3. With z fixed, why is the mean the unique minimiser? Use the Hessian to argue it.
4. Can K-means separate two concentric circles? Why or why not?
5. Why can't you choose K by minimising F over K?
6. What is the cost per iteration for n = 10⁶, K = 10, d = 100? (≈ 10⁹ operations for the distances.)