← All notes · ML revision (DA5400)

L18: Choosing K in K-means (Objective + Penalty, AIC/BIC) & Clustering Wrap-up

Source: https://www.youtube.com/watch?v=4-ycWgdMXD4 (≈17 min) Previous: L13 · L14 · L15 · L16 · L17

The big picture

So far K was always given. In practice it isn't: you get a pile of points and have to decide how many groups they really contain. This lecture argues that choosing K is a trade-off between fit and simplicity. It turns that trade-off into one rule, minimise objective + penalty, and names the two standard penalties (AIC and BIC). It then closes the clustering unit by tying together everything learned about Lloyd's algorithm.

The ideas build on each other in four steps:

1. The objective alone can't choose K, because it always prefers more clusters.
2. So we must also charge for complexity: objective + penalty(K).
3. AIC and BIC are principled penalties, but they need a likelihood view of the objective.
4. Wrap-up: the four questions about Lloyd's algorithm are now all answered.

Concept 1: The objective alone always says "more clusters"

The K-means objective F = Σᵢ ‖xᵢ − μ_{zᵢ}‖² measures how tightly points sit around their cluster means. For any fixed K it is the right thing to minimise. But comparing F across different K is misleading:

• More clusters means smaller clusters, and smaller clusters mean points are closer to their means. So the best achievable F keeps falling as K grows.
• At the extreme K = n, every point is its own cluster and its own mean, so F = 0, the lowest possible value.

If we picked K by minimising F, we would always get K = n. That isn't clustering at all. Clustering is meant to compress the data into a few representatives. So F still matters (within a given K, we want the best partition), but on its own it can't choose K.

Concept 2: Charge a price for each cluster (objective + penalty)

The fix is to make complexity cost something. Pick an increasing function Penalty(K) and choose

$$ K^\star=\arg\min_{K}\;\Big[\,F^\star(K)+\text{Penalty}(K)\,\Big], $$

where F*(K) is the best objective found with K clusters.

The "buying a cluster" intuition. Going from K to K + 1 costs Penalty(K+1) − Penalty(K). The purchase is worth it only if F drops by more than that price.

• Early on, it pays. Going from one big cluster to two or three cuts F dramatically, because real groups are being separated.
• Later, it doesn't. Once the true groups are found, extra clusters just split real ones and F barely moves, while the penalty keeps rising.

So F*(K) falls and the penalty rises. Their sum is U-shaped, and its lowest point is K*.

How you use it in practice:

1. For each candidate K, run Lloyd's algorithm with a good initialization (K-means++ or several random restarts) and keep the best F for that K.
2. Add the penalty to each.
3. Pick the K where the total is smallest.

The open question is which penalty to use. That depends on how much you are willing to pay per cluster.

Concept 3: Principled penalties — AIC and BIC

Two standard criteria give the penalty a principled form. Both use a likelihood L, a measure of how probable the data is under the fitted model, in place of the raw objective:

Key ideas:

• The −2 log L term plays the role of the objective (how well the model fits). The other term is the price paid for complexity.
• BIC charges more per cluster than AIC once n is larger than about 7, so it tends to choose fewer clusters.
• Each criterion is "optimal" under a particular assumption about how the data was generated (a probabilistic model). The course hasn't introduced probabilistic models or estimation yet, so the lecture only states these criteria and promises to come back to them. Section 6 of these notes does that derivation now.
• Other heuristics exist, notably the elbow method (look for the K where the F*(K) curve bends). They all follow the same idea: try several values of K and balance fit against complexity.

One subtlety for later: in the standard AIC/BIC formulas, the multiplier is the number of free parameters, not the number of clusters. For K-means in ℝᵈ, each cluster brings about d + 1 parameters (§6.3).

Concept 4: Wrap-up — what we now know about Lloyd's algorithm

The clustering unit began with the K-means objective (L13) and Lloyd's algorithm (L14), and four questions about that algorithm. All four now have answers, and together they give a complete practical picture:

What comes next: the course moves on to estimation: fitting the parameters of a probabilistic model to data. That will make the likelihood in AIC/BIC concrete, and it bridges unsupervised and supervised learning.

Going further (§6): d-dimensional derivations

This section goes beyond the video.

6.1 Proof that F*(K) is non-increasing, with F*(n) = 0

Claim: F*(K+1) ≤ F*(K) for every K < n.

Proof. Take an optimal K-partition with means μ₁, …, μ_K. As long as K < n, some cluster C contains a point x' ≠ its mean μ_C. (If every point equalled its own cluster's mean, F would be 0 already.) Make x' a cluster of its own, with mean x'. Every other point keeps its old centre, so the cost changes only by removing x''s term ‖x' − μ_C‖² > 0. Then re-centring C∖{x'} can only lower the cost further (Fact 1, L15). So this valid (K+1)-partition costs less than F*(K), which means F*(K+1) < F*(K). ∎

At K = n, every point is its own mean, so F*(n) = 0.

This is why a penalty is unavoidable: with no penalty, the argmin is always K = n.

6.2 K-means is a likelihood: the spherical-Gaussian model

Suppose each cluster generates points from a Gaussian with a shared spherical covariance:

$$ x_i \mid z_i=k \;\sim\; \mathcal{N}(\mu_k,\;\sigma^2 I_d),\qquad P(z_i=k)=\pi_k . $$

The density in ℝᵈ is

$$ p(x\mid\mu_k)=(2\pi\sigma^2)^{-d/2}\exp\!\Big(-\frac{\lVert x-\mu_k\rVert^2}{2\sigma^2}\Big). $$

Treat the assignments z as parameters, as K-means does (a hard "classification" likelihood):

$$ \log L(z,\mu,\pi,\sigma^2)=\sum_{i=1}^n\Big[\log\pi_{z_i}-\frac d2\log(2\pi\sigma^2)-\frac{\lVert x_i-\mu_{z_i}\rVert^2}{2\sigma^2}\Big]. $$

$$ \boxed{\;-2\log L=\frac{F(z,\mu)}{\sigma^2}+nd\log(2\pi\sigma^2)-2\sum_{k}n_k\log\pi_k\;} $$

• With σ and π fixed, maximising the likelihood is the same as minimising F. K-means is maximum likelihood for this model. That's the bridge to the estimation lectures.
• The MLE for σ²: differentiate with respect to σ² and set to 0, giving σ̂² = F/(nd). This is the average squared deviation per coordinate.
• The MLE for π: π̂ₖ = nₖ/n. Plugging in both gives

$$ -2\log\hat L = nd\,\log\!\Big(\frac{2\pi F}{nd}\Big) + nd - 2\sum_k n_k\log\frac{n_k}{n}. $$

The last term is n times the entropy of the cluster sizes, and it grows as clusters are split. That turns out to matter (§6.8).

6.3 AIC and BIC with the correct parameter count

Count the free parameters p in this model:

$$ \text{AIC}(K)=-2\log\hat L+2p,\qquad \text{BIC}(K)=-2\log\hat L+p\,\log n . $$

• The d-dimensional point: each new cluster costs d + 1 parameters, not 1. In high dimensions, adding a cluster is expensive. The lecture's "2K" and "K log n" are the 1-parameter-per-cluster picture.
• σ known (e.g. σ = 1): −2 log L = F/σ² + const, so BIC ≈ F/σ² + Kd·log n. That's literally "objective + penalty", with penalty slope d·log n per cluster.
• AIC vs BIC:
  • AIC estimates out-of-sample prediction error (an approximation to KL divergence). It tends to over-estimate K.
  • BIC approximates the log of the Bayesian marginal likelihood (via a Laplace approximation). It is consistent: if the model is right, it picks the true K as n → ∞. It is more conservative because log n > 2 once n > 7.

6.4 How fast does F*(K) fall? The K^(−2/d) law

For data spread uniformly (no real clusters) over a region of ℝᵈ, quantisation theory (Zador's theorem) gives

$$ F^\star(K)\;\propto\;n\,K^{-2/d}\quad\text{for large }K . $$

Intuition: K cells share a volume V, so each cell has volume V/K and radius ∝ (V/K)^(1/d). The squared radius is ∝ K^(−2/d).

Consequences:

• d = 1: F ∝ K⁻², a steep drop even when the data has no clusters.
• Large d: the exponent −2/d is close to 0, so F*(K) falls very slowly without real structure. A sharp drop in high dimensions is therefore strong evidence of real clusters.
• This is why "how much did F drop?" must be compared with what you'd expect from structureless data. That's exactly what the gap statistic (§6.7) does.

6.5 The elbow method, made precise

Plot F*(K) against K. The elbow is where the marginal gain ΔF(K) = F*(K−1) − F*(K) suddenly collapses. Two simple ways to formalise it:

• Largest second difference: K* = argmax_K [ΔF(K) − ΔF(K+1)].
• Ratio test: pick the first K where ΔF(K+1)/ΔF(K) drops below some threshold (e.g. 0.2).

Weakness: with real data the bend is often gradual, so the elbow is subjective. It's best used as a quick visual check.

6.6 Silhouette score (any d, any distance)

For point i in cluster A:

• a(i) is the mean distance from i to the other points in A. This measures cohesion.
• b(i) = min over B ≠ A of the mean distance from i to the points in B. This measures separation from the nearest other cluster.

$$ s(i)=\frac{b(i)-a(i)}{\max\{a(i),\,b(i)\}}\in[-1,1],\qquad \bar s(K)=\frac1n\sum_i s(i). $$

• s ≈ 1: the point is well inside its cluster.
• s ≈ 0: the point sits on a boundary.
• s < 0: the point is probably in the wrong cluster.

Choose K = argmax s̄(K). It needs no penalty, but it's undefined for K = 1 and costs O(n²d).

6.7 Gap statistic (Tibshirani, Walther & Hastie, 2001)

Compare log F*(K) on your data with its expected value on reference data with no clusters, e.g. points drawn uniformly over the bounding box, B times:

$$ \text{Gap}(K)=\frac1B\sum_{b=1}^{B}\log F_b^{\star}(K)\;-\;\log F^\star(K). $$

Choose the smallest K with Gap(K) ≥ Gap(K+1) − s_{K+1}, where s is the standard error over the B reference sets. This automatically accounts for the K^(−2/d) baseline drop from §6.4, and it can also conclude K = 1, which silhouette can't.

6.8 Tested comparison: true K = 4, d = 10, n = 400

import numpy as np

def kmeanspp(X, K, rng):
    n = len(X); M = [X[rng.integers(n)]]; D2 = ((X - M[0])**2).sum(1)
    for _ in range(1, K):
        M.append(X[rng.choice(n, p=D2 / D2.sum())]); D2 = np.minimum(D2, ((X - M[-1])**2).sum(1))
    return np.array(M)

def best_kmeans(X, K, restarts=10):
    best = (np.inf, None)
    for s in range(restarts):
        M = kmeanspp(X, K, np.random.default_rng(s))
        for _ in range(200):
            z = ((X[:, None] - M[None])**2).sum(-1).argmin(1)
            M_new = np.stack([X[z == k].mean(0) if np.any(z == k) else M[k] for k in range(K)])
            if np.allclose(M_new, M): break
            M = M_new
        F = ((X - M[z])**2).sum()
        if F < best[0]: best = (F, z)
    return best

def silhouette(X, z):
    D = np.sqrt(((X[:, None] - X[None])**2).sum(-1)); s = np.zeros(len(X))
    for i in range(len(X)):
        same = (z == z[i]); same[i] = False
        if not same.any(): continue
        a = D[i, same].mean()
        b = min(D[i, z == k].mean() for k in set(z) if k != z[i])
        s[i] = (b - a) / max(a, b)
    return s.mean()

rng = np.random.default_rng(0)
true_K, d, n_per = 4, 10, 100
X = np.vstack([rng.normal(c, 1.0, (n_per, d)) for c in rng.uniform(-6, 6, (true_K, d))])
n = len(X)

print(" K        F*(K)   drop      AIC      BIC     sil")
prev = None
for K in range(1, 9):
    F, z = best_kmeans(X, K)
    nk = np.bincount(z, minlength=K); nk = nk[nk > 0]
    neg2LL = n*d*np.log(2*np.pi*F/(n*d)) + n*d - 2*(nk*np.log(nk/n)).sum()  # incl. mixing weights
    p = K*d + 1 + (K - 1)                            # means + shared σ² + weights
    bic = neg2LL + p*np.log(n)
    aic = neg2LL + 2*p
    sil = silhouette(X, z) if K > 1 else float('nan')
    drop = "" if prev is None else f"{prev-F:7.0f}"
    print(f"{K:2d} {F:12.0f} {drop:>7s} {aic:8.0f} {bic:8.0f} {sil:7.3f}")
    prev = F

Output when I ran it:

All four methods agree on K = 4:

• Elbow: the drop collapses from 8,855 to 118 after K = 4.
• AIC and BIC both reach their minimum at K = 4.
• Silhouette peaks at 0.701 at K = 4.
• Note F*(4)/(nd) = 3966/4000 ≈ 0.99, which is close to the true σ² = 1. The MLE σ̂² = F/(nd) recovers the noise level once K is right.

What went wrong first (a real pitfall): my first version of BIC used only nd·log(F/(nd)) + Kd·log n. It dropped the mixing-weight term −2Σnₖ log(nₖ/n) and counted only the means. That version kept decreasing past K = 4 (9209 → 7538 → 4839 → 205 → 144 → 95 → 43 → −23) and would have picked K = 8. Two reasons:

• For hard assignments, splitting a true Gaussian cluster always lowers F somewhat. Here the drop was about 118, which is nd·log(3966/3848) ≈ 121 in −2 log L.
• That gain beats a penalty of only d·log n ≈ 70 per cluster.

The missing cluster-size entropy term (the −2Σnₖ log(nₖ/n) part of −2 log L̂) is what charges for splitting. Lesson: "objective + penalty" only works if the objective is the full likelihood of a well-specified model and p counts all the parameters.

6.9 Where this leads: GMMs and EM

Replace the hard assignment zᵢ with soft probabilities, and allow a separate covariance Σₖ for each cluster. That gives the Gaussian mixture model, fitted with EM (the soft version of Lloyd's algorithm). BIC on GMMs is the standard principled way to choose K, and it's exactly where the course's "estimation" part is heading.

Key takeaways

1. F*(K) is non-increasing and reaches 0 at K = n, so minimising F alone always picks K = n.
2. Choose K* = argmin [F*(K) + Penalty(K)]. Add a cluster only if F drops by more than the extra penalty.
3. K-means is maximum likelihood for spherical Gaussians with shared variance: −2 log L = F/σ² + nd·log(2πσ²) − 2Σnₖ log πₖ.
4. In ℝᵈ, each cluster adds d + 1 parameters. So AIC = −2 log L̂ + 2(Kd + K) and BIC = −2 log L̂ + (Kd + K)·log n. BIC is the more conservative of the two.
5. Without real structure, F*(K) falls like K^(−2/d), very slowly in high dimensions. The gap statistic uses this as its baseline.
6. On test data with 4 true clusters, the elbow, AIC, BIC and silhouette all picked K = 4. The naive BIC, missing the cluster-size term, did not.

Self-check questions

1. Prove F*(K+1) ≤ F*(K) without assuming the optimal (K+1)-partition is related to the K one.
2. Derive σ̂² = F/(nd) by differentiating −2 log L with respect to σ².
3. For d = 50, n = 10,000, how much must F/σ² drop for BIC to accept one more cluster? (More than (d+1)·ln n ≈ 51 × 9.21 ≈ 470, ignoring the entropy term.)
4. Why does uniform data in d = 2 still show F*(K) falling quickly, and why does that fool the elbow method?
5. Why can silhouette never choose K = 1, while the gap statistic can?
6. Explain in one sentence why the cluster-size entropy term stops BIC from over-splitting.