K-Means Clustering, machine learning lesson

A service records two usage measures for each account and wants a handful of representative profiles that summarize all of them. K-means picks $k$ centers, gives each account to its nearest center, and places the centers so that the total squared distance from the accounts to their centers is as small as possible. This lesson writes down that objective, runs the two alternating steps of Lloyd’s algorithm by hand on six points, and shows why the result depends on the starting centers and on the chosen $k$ , and why a cluster is not automatically a natural category.

Key terms

K-means clustering

Definition Let $\mathbf{x}_1, \dots, \mathbf{x}_n$ be points in $\mathbb{R}^p$ and $k$ an integer with $1 \le k \le n$ . K-means clustering seeks an assignment $c_i \in \{1, \dots, k\}$ of each point to a cluster and centers $\boldsymbol{\mu}_1, \dots, \boldsymbol{\mu}_k$ in $\mathbb{R}^p$ that minimize the within-cluster sum of squares

J = \sum_{i=1}^{n} \lVert \mathbf{x}_i - \boldsymbol{\mu}_{c_i} \rVert_2^2 .

Finding the global minimum is computationally hard in general. The standard method, Lloyd’s algorithm, starts from initial centers and alternates the assignment step and the update step until the centers stop changing. With a fixed tie rule and no empty clusters, $J$ never increases and the algorithm stops after finitely many iterations, because there are finitely many partitions and every update that moves a center strictly lowers $J$ . It stops at a partition that neither step can improve, a local optimum in this sense, which need not be the global minimum. The method assumes that squared Euclidean distance in the chosen feature representation is a useful notion of closeness.

In plain words K-means summarizes the data by $k$ representative points and judges the summary by how far, in squared distance, each point lies from its representative. It favors compact, roughly round groups of similar spread in the chosen units. It returns a partition for any data and any $k$ , whether or not the data contain distinct groups.

Example Six constructed usage profiles in a common dimensionless scale are $(0, 0)$ , $(0, 2)$ , $(2, 0)$ , $(8, 8)$ , $(8, 10)$ and $(10, 8)$ . With $k = 2$ the obvious candidate puts the first three points in one cluster and the last three in the other. The next cards compute the centers and the objective of this partition, run the algorithm, and compare starting points.

Centroid and within-cluster sum of squares

Definition The centroid of a nonempty cluster $C$ with $n_C$ points is the coordinatewise mean

\boldsymbol{\mu}_C = \frac{1}{n_C} \sum_{i \in C} \mathbf{x}_i .

The within-cluster sum of squares of the cluster is $W(C) = \sum_{i \in C} \lVert \mathbf{x}_i - \boldsymbol{\mu}_C \rVert_2^2$ , in squared feature units, and the objective $J$ of a partition with every center at its centroid is the sum of $W(C)$ over the clusters. For any point $\mathbf{c}$ in $\mathbb{R}^p$ ,

\begin{aligned} &\sum_{i \in C} \lVert \mathbf{x}_i - \mathbf{c} \rVert_2^2 \\[0.4em] &\quad = W(C) + n_C \lVert \boldsymbol{\mu}_C - \mathbf{c} \rVert_2^2 , \end{aligned}

so the centroid is the unique center that minimizes the cluster’s sum of squared distances. An empty cluster has no centroid.

In plain words The centroid is the average member, the balance point of the cluster, and it need not be one of the observed points. Because the distances are squared, one far point pulls the centroid toward it.

Example The cluster $(0, 0)$ , $(0, 2)$ , $(2, 0)$ has centroid $(2/3, 2/3)$ . The squared distances to it are $8/9$ , $20/9$ and $20/9$ , so $W = 48/9 = 16/3$ . Measured from $(0, 0)$ instead, the sum is $0 + 4 + 4 = 8$ , and the identity agrees, since $16/3 + 3 \cdot 8/9 = 8$ .

Assignment step and update step

Definition Given current centers $\boldsymbol{\mu}_1, \dots, \boldsymbol{\mu}_k$ , the assignment step gives each point the index of a nearest center,

c_i = \operatorname{argmin}_{j} \lVert \mathbf{x}_i - \boldsymbol{\mu}_j \rVert_2^2 ,

with a stated tie rule, which on this page is the lower center index. Given the assignments, the update step replaces each center by the centroid of the points assigned to it. Each step minimizes $J$ over its own variables with the others held fixed. In the assignment step every term of $J$ depends only on its own point’s choice, so the nearest center minimizes each term separately. In the update step $J$ splits into one sum per cluster, and for cluster $j$ with $n_j > 0$ points and coordinate $r$ of its center

\begin{aligned} \frac{\partial J}{\partial \mu_{jr}} &= 2 n_j \mu_{jr} \\[0.4em] &\quad - 2 \sum_{i : c_i = j} x_{ir} , \end{aligned}

which is zero at the coordinate mean. Hence neither step can increase $J$ . If a cluster receives no points, its mean is undefined and the algorithm needs a declared policy, such as stopping with an error or moving that center to a poorly represented point, and the objective has to be rechecked after any such move.

In plain words Each point goes to its nearest center, and then each center moves to the middle of its points. The two steps repeat, and each round lowers the total squared distance or leaves it unchanged. A new point is assigned by comparing it with every center, not by its resemblance to one member of a cluster.

Example Start the six points with centers $(0, 0)$ and $(8, 8)$ . The first three points are nearer the first center and the last three nearer the second, with squared distances 0, 4, 4, 0, 4 and 4, so $J = 16$ . The update moves the centers to $(2/3, 2/3)$ and $(26/3, 26/3)$ , each cluster now contributes $16/3$ , and

J = 2 \cdot \frac{16}{3} = \frac{32}{3} \approx 10.667 .

The next assignment step keeps every point in its cluster, so the update reproduces the same centers and the algorithm stops at a fixed point.

K-means initialization

Definition K-means initialization is the choice of the starting centers, and the end point of Lloyd’s algorithm depends on it. Common choices are $k$ distinct data points drawn at random and k-means++ seeding. That seeding draws the first center uniformly from the data points and each further center from the data points with probability proportional to $D(\mathbf{x})^2$ , the squared distance from $\mathbf{x}$ to the nearest center already chosen. The expected objective of the k-means++ seeds alone is at most $8(\ln k + 2)$ times the optimal value. Because one run can end at a poor local optimum, implementations usually run several initializations and keep the run with the smallest $J$ . That improves the result in practice without guaranteeing the global minimum. A fixed random seed makes the runs reproducible.

In plain words The algorithm only walks downhill from where it starts. Two starting centers inside one group can end in a split that a better start would avoid, and spreading the starting centers out and trying several starts lowers that risk.

Example On the six points, the starting centers $(0, 0)$ and $(0, 2)$ both lie in the lower-left group. The first assignment gives $(0, 0)$ and $(2, 0)$ to the first center and the other four points to the second, with $J = 368$ . The objective is then 97 after the update, $139/4 = 34.75$ after the next assignment moves $(0, 2)$ to the first cluster, and $32/3$ after the next update, which reaches the same partition and centers as the first run. That agreement holds for these six points only. The four constructed points $(0, 0)$ , $(4, 0)$ , $(0, 1)$ and $(4, 1)$ show a failure. Starting at $(0, 0)$ and $(0, 1)$ gives the bottom pair and the top pair with centers $(2, 0)$ and $(2, 1)$ , a fixed point with $J = 16$ . Starting at $(0, 0)$ and $(4, 0)$ gives the left pair and the right pair with $J = 1$ . The code below records $J$ after each assignment step.

import numpy as np

def lloyd(X, centers, max_iter=100):
    C = np.array(centers, dtype=float)
    history = []
    for _ in range(max_iter):
        d2 = ((X[:, None, :] - C[None, :, :]) ** 2).sum(axis=2)
        labels = d2.argmin(axis=1)                 # ties go to the lower index
        history.append(d2[np.arange(len(X)), labels].sum())
        if len(np.unique(labels)) < len(C):
            raise ValueError("empty cluster, a declared policy is needed")
        new = np.array([X[labels == j].mean(axis=0) for j in range(len(C))])
        if np.array_equal(new, C):
            return C, labels, history
        C = new
    raise RuntimeError("no fixed point within max_iter")

X = np.array([[0, 0], [0, 2], [2, 0], [8, 8], [8, 10], [10, 8]], dtype=float)
for start in ([[0, 0], [8, 8]], [[0, 0], [0, 2]]):
    C, labels, history = lloyd(X, start)
    print(labels, np.round(history, 3))    # [16. 10.667] and [368. 34.75 10.667]

Number of clusters

Definition The number of clusters $k$ is a choice made by the analyst, not a quantity the algorithm discovers. Let $J^*(k)$ be the smallest achievable within-cluster sum of squares with $k$ clusters. For $k < n$ it satisfies $J^*(k + 1) \le J^*(k)$ , since adding a center at a data point cannot raise the total, and $J^*(k) = 0$ once $k$ reaches the number of distinct points. So the smallest training objective always favors the largest $k$ , and separate local runs need not even decrease with $k$ , because each can stop at a different local optimum. Common diagnostics are an elbow plot of $J$ against $k$ , which looks for the point after which extra centers bring small reductions, the silhouette, which compares each point’s average distance to the other members of its cluster with its average distance to the nearest other cluster, and the stability of the clusters across resamples. None of them establishes true categories in a population. When the clusters feed a prediction, $k$ is a hyperparameter chosen on validation data, never on the test set.

In plain words More centers always fit the training points at least as closely, so the objective alone cannot choose $k$ . The useful question is which number of groups gives a summary that is stable and serves the purpose at hand. Calling a cluster “high value” or “high risk” needs evidence beyond the algorithm, which never saw what such a name means.

Example For the six points the smallest objectives are $608/3 \approx 202.67$ for $k = 1$ , $32/3 \approx 10.67$ for $k = 2$ , $22/3 \approx 7.33$ for $k = 3$ , then 4, 2 and 0 for $k = 4$ , 5 and 6. Nearly all of the reduction comes from the second center. A local run can still misbehave. With $k = 4$ , the starting centers $(0, 0)$ , $(1, 1)$ , $(8, 8)$ and $(9, 9)$ are already a fixed point with $J = 8$ , which is above the best three-cluster value $22/3$ .

Check yourself

A proposed cluster contains $(1, 1)$ , $(3, 1)$ and $(2, 4)$ . Find its centroid and its within-cluster sum of squares. Another cluster has centroid $(6, 1)$ . The new point $(4, 1)$ is only 1 unit from the member $(3, 1)$ of the first cluster. Which cluster does the assignment step give it?

Show the answer

The coordinate means are $(1 + 3 + 2)/3 = 2$ and $(1 + 1 + 4)/3 = 2$ , so the centroid is $(2, 2)$ . The squared distances to it are $1 + 1 = 2$ , $1 + 1 = 2$ and $0 + 4 = 4$ , so the within-cluster sum of squares is 8.

The new point has squared distance $4 + 1 = 5$ to the center $(2, 2)$ and $4 + 0 = 4$ to the center $(6, 1)$ , so the assignment step gives it to the second cluster. K-means compares a point with the centers, not with individual members. If the point joins the data, the next update moves the second center toward it, and the two steps continue until nothing changes.

Run k-means with $k = 3$ on the six profiles of this lesson, starting from the centers $(0, 0)$ , $(8, 8)$ and $(100, 100)$ . Find the first assignment and its objective. What happens in the update step?

Show the answer

The center $(100, 100)$ is far from every point, so the first assignment is the same as with two centers. The points $(0, 0)$ , $(0, 2)$ and $(2, 0)$ go to the first center, the other three go to the second, and $J = 0 + 4 + 4 + 0 + 4 + 4 = 16$ .

The third cluster is empty, and the mean of no points is undefined, so the update step cannot be carried out as written. An implementation needs a declared policy, for example stopping with an error, restarting from new seeds, or moving the empty center to a point far from its current center and then recomputing the objective. A center placed outside the data is a poor start. Seeds chosen as distinct data points, as in k-means++, cannot give an empty cluster in the first assignment, because each seed is at distance zero from its own point and at a positive distance from the other seeds.

The points 0, 1, 2, 3 and 20 lie on a line. Find the partition into two clusters with the smallest within-cluster sum of squares, and explain what it says about k-means and outliers. In one dimension an optimal partition can be found among the splits into a left group and a right group.

Show the answer

The four splits give the following sums, each the left group’s sum of squares plus the right group’s.

Left group	Right group	Sum of squares
$\{0\}$	$\{1, 2, 3, 20\}$	$0 + 245 = 245$
$\{0, 1\}$	$\{2, 3, 20\}$	$1/2 + 614/3 \approx 205.17$
$\{0, 1, 2\}$	$\{3, 20\}$	$2 + 144.5 = 146.5$
$\{0, 1, 2, 3\}$	$\{20\}$	$5 + 0 = 5$

The best partition puts 20 alone and summarizes the other four points by the single center 1.5. One of the only two centers is spent on one extreme point, because squared distance makes a far point very expensive. With all five points in one cluster the centroid would be 5.2, pulled far from where most points lie. With $k = 3$ the best partition is $\{0, 1\}$ , $\{2, 3\}$ , $\{20\}$ with $J = 1$ . An outlier should be investigated and handled by a documented rule before the clusters are read, and naming the singleton cluster as a group of customers would claim far more than one point supports.

K-Means Clustering

Key terms

K-means clustering

Centroid and within-cluster sum of squares

Assignment step and update step

K-means initialization

Number of clusters

Check yourself

How this lesson connects

Builds on

Used later in