KMeans Clustering

KMeans clustering is one of the most fundamental and widely-used unsupervised learning algorithms for partitioning data into distinct groups. This post explores the mathematical foundation, implementation details, and practical considerations of this classic clustering technique.

Problem

Given a set of $n$ points $X=(x_1,x_2,...,x_n)$ where $x_i\in R^d$, assign the points to $k\le n$ clusters  such that the following loss function should be minimized. (Wiki Kmeans, 2024). Formally,

$$L(\mu )=\sum _{i=1}^n\underset{j\in 1,...,k}{\min }\Vert x_i-{\mu }_j{\Vert }_2^2$$

Algorithm (Naive KMeans)

The problem is NP-hard wrt finding the centroids which minimize the WCSS, but is approximated in practice using the following iterative refinement technique in naive KMeans:

1. Initialize the $k$ centroids $({\mu }_1,{\mu }_2,...,{\mu }_k)$ where ${\mu }_j\in R^d$ randomly

2. Repeat until convergence or max iterations $I$

   • Assignment step: For every point $x_i$, assign cluster label $l_i$ as the cluster with minimum $L^2$ distance:
   $$l_i=\arg \underset{j}{\min }\Vert x_i-{\mu }_j{\Vert }_2^2$$
   • Update step: Update centroid $c_j$ as:
   $${\mu }_j=\frac{1}{|\{i:l_i=j\}|}\sum _{i:l_i=j}{x}_i$$
   • Empty Cluster Handling: Check for clusters with zero assigned points and re-initialize them randomly using data points Time complexity for the above algorithm is $O(Inkd)$

Proof of Convergence

Lemma 1: The loss function is monotonically non-increasing for both the assignment and update step. (Andreas Krause, 2016) Proof: Let $l=(l_1,l_2,...,l_n)$ denotes the cluster assignment for the $n$ points.

Assignment step: The loss function $L(\mu ,l)$ can be written as:

$$L(\mu ,l)=\sum _{i=1}^n\Vert x_i-{\mu }_{l_i}{\Vert }_2^2$$

Consider a data point $x_i$ , and let $l_i$ be the assignment from the previous iteration and $z_i^{\ast }$ be the new assignment obtained as:

$$l_i^{\ast }\in \arg \underset{j\in 1,...,k}{\min }\Vert x_i-{\mu }_j{\Vert }_2^2$$

Let $l^{\ast }$ denote the new cluster assignment, the change in loss function can be written as

$$L(\mu ,l^{\ast })-L(\mu ,l)=\sum _{i=1}^n(\Vert x_i-{\mu }_{l_i^{\ast }}{\Vert }_2^2-\Vert x_i-{\mu }_{l_i}{\Vert }_2^2)\le 0$$

The above inequality holds because $l_i^{\ast }$ assigns each $x_i$ to the nearest cluster.

Update step: The loss function $L(\mu ,l)$ can be alternatively written as:

$$L(\mu ,l)=\sum _{j=1}^k(\sum _{i:l_i=j}\Vert x_i-{\mu }_j{\Vert }_2^2)$$

For the $j^{th}$ cluster, let’s denote the previous centroid as ${\mu }_j$ and the updated one as ${\mu }_j^{\ast }$

$${\mu }_j^{\ast }=\frac{1}{|\{i:l_i=j\}|}\sum _{i:l_i=j}{x}_i$$

Let ${\mu }^{\ast }$ denote the new centroids for all the $k$ clusters, change in loss function is given as:

$$L({\mu }^{\ast },l)-L(\mu ,l)=\sum _{j=1}^k((\sum _{i:l_i=j}\Vert x_i-{\mu }_j^{\ast }{\Vert }_2^2)-(\sum _{i:l_i=j}\Vert x_i-{\mu }_j{\Vert }_2^2)\le 0$$

We can verify after taking gradients of $L(\mu ,l)$ wrt each ${\mu }_j$ that $L(\mu ,l)$ is minimized with the updated definition of ${\mu }_j^{\ast }$, that’s why the inequality holds.

Claim: The K-Means algorithm terminate in finite number of steps.

1. If the clustering assignment changes, the newer one will have a lower cost (from Lemma 1)
2. If the cluster assignment doesn’t change, the centroids don’t change, and the algorithm terminates. Since the number of clusterings is finite and equal to $k^n$, the algorithm will eventually hit condition 2, and then terminate.

Implementation (in python using pytorch)

Show me the code

import torch
import torch.nn.functional as F
from typing import Tuple
 
# Two-loop solution
def k_means_clustering_twoloops(
    data: torch.Tensor,
    centers: torch.Tensor,
    max_iterations: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
        data: Points to be clustered [N, D]
        centers: Centroid matrix [K, D]
        max_iterations: max iterations clustering should run
    """
    for i in range(max_iterations):
        dist = torch.zeros((data.shape[0], centers.shape[0]))  # (N, K)
        for j in range(data.shape[0]):
            dist_j = torch.sqrt(
                ((data[j : j + 1] - centers) ** 2).sum(dim=1)
            )  # (K, )
            dist_j_sq = ((data[j : j + 1] - centers) ** 2).sum(dim=1) # (K, )
            dist[j, :] = dist_j
 
        labels = torch.argmin(dist, dim=1)  # (N,)
        for j in range(centers.shape[0]):
            points_in_cluster = data[labels == j, :]
            if points_in_cluster.shape[0] > 0:
                centers[j] = torch.mean(points_in_cluster, dim=0)
            else:
                # No point assigned to this cluster, pick a data point arbitraily as center
                centers[j] = data[torch.randint(0, data.shape[0], (1,)), :]
 
    return centers, labels
 
# One-loop solution
def k_means_clustering_oneloop(
    data: torch.Tensor,
    centers: torch.Tensor,
    max_iterations: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
        data: Points to be clustered [N, D]
        centers: Centroid matrix [K, D]
        max_iterations: max iterations clustering should run
    """
    N, _ = data.shape
    for i in range(max_iterations):
        # Alternatively torch.cdist API can be used
        dist = torch.sqrt(
            torch.sum(data**2, dim=1).unsqueeze(1)
            + torch.sum(centers**2, dim=1).unsqueeze(0)
            - 2 * data @ centers.T
        )  # [N, K]
        labels = torch.argmin(dist, dim=1)  # [N]
 
        # create [K, N] matrix M which has M(i, j) = 1 if point j belongs to cluster i, zero otherwise
        centers_to_labels = torch.zeros_like(dist).T  # [K, N]
        centers_to_labels[labels, torch.arange(N)] = 1
        centers_to_labels = centers_to_labels / torch.max(
            torch.sum(centers_to_labels, dim=1), torch.tensor(1e-12)
        ).unsqueeze(1)
        centers = centers_to_labels @ data  # [K, D]
 
        # Reinitialize the centers which didn't get a point based on random rows from data
        # Identify rows in centers that are all zeros
        zero_rows_mask = (centers == 0.0).all(dim=1)  # Boolean mask for zero rows in centers
 
        # Count how many rows are zero
        num_zero_rows = zero_rows_mask.sum().item()
        random_indices = torch.randint(0, data.shape[0], (num_zero_rows,))
        random_rows_from_data = data[random_indices]  # Select random rows
 
        # Replace zero rows in centers with the selected rows from data
        centers[zero_rows_mask] = random_rows_from_data
    return centers, labels

Clustering Example

Fig 1. Data Points

Fig 2. KMeans Clustering - One Loop After 100 Iterations

Run-time comparison

Show me the code

import timeit
 
# Time the two-loop k-means function
time_two_loops = timeit.timeit(
    lambda: k_means_clustering_twoloops(data, centers, max_iter), number=10
)
print(f"Time taken for two-loop k-means: {time_two_loops:.6f} seconds")
 
# Time the one-loop k-means function
time_one_loop = timeit.timeit(
    lambda: k_means_clustering_oneloop(data, centers, max_iter), number=10
)
print(f"Time taken for one-loop k-means: {time_one_loop:.6f} seconds")

Time taken for two-loop k-means: 41.656580 seconds Time taken for one-loop k-means: 0.507115 seconds

Single time version is $\sim 85$ faster than the two-loops version for a small datasets of 100 2-dim points with 10 iterations

Improved Initialization of Centroids with KMeans++

Naive KMeans initializes the centroids randomly, which could lead to convergence to bad local optimum. Here is an example below:

Fig 3. Poor clustering example due to bad initialization of centroids

K-Means++ aims to solve the initialization problem, while providing $\Theta (logk)$-competitive accuracy guarantees. The intuition is to select the initial centers which are further apart from each other. The authors provide preliminary resulting demonstrating that KMeans++ leads to both improvements in speed and accuracy in practice. (Arthur & Vassilvitskii, 2007). The algorithm works as follows:

1. Choose an initial center ${\mu }_1$ uniformly at random from $X$.
2. Choose the next center ${\mu }_i$ , selecting ${\mu }_i=x^{\prime }\in X$ with probability $\frac{D(x^{\prime }{)}^2}{\sum _{x\in X}D(x{)}^2}$, where $D(x)$ denote the shortest distance from a data point $x^{\prime }$ to the closest center we have already chosen.
3. Repeat Step 2. until we have chosen a total of k centers.
4. Proceed as with the standard k-means algorithm.

Show me the code

import torch
 
def kmeans_pp_initialization(data: torch.Tensor, num_clusters: int) -> torch.Tensor:
    """
    Performs k-means++ initialization.
 
    Args:
        data: The data points to be clustered [N, D].
        num_clusters: The desired number of clusters (K).
 
    Returns:
        A tensor of initial cluster centers [K, D].
    """
    n_samples = data.shape[0]
    centers = torch.zeros(num_clusters, data.shape[1], dtype=data.dtype)
    # Choose the first center randomly from the data points
    centers[0] = data[torch.randint(0, n_samples, (1,))]
 
    for i in range(1, num_clusters):
      distances = torch.min(torch.cdist(data, centers[:i]), dim=1)[0] # [N]
      probabilities = distances / torch.sum(distances)
      cumulative_probabilities = torch.cumsum(probabilities, dim=0) # [N]
 
      # Generate random number
      rand_val = torch.rand(1)
 
      # Find the index of the next center based on cumulative probabilities
      next_center_index = torch.searchsorted(cumulative_probabilities, rand_val)
      centers[i] = data[next_center_index]
 
    return centers

Check out the colab notebook for all the code pointers and plots.

How to select the number of clusters

There are several heuristics to identify the number of clusters $K$ based on Elbow Method, Silhouette Score, cross-validation or Dunn index, either using the KMeans loss $L(\mu )$ or other good of fitness metrics. (Neptune Kmeans, 2024)

Further Readings

• Dasgupta Kmeans Handout
• GeeksForGeeks KMeans
• Elkan Kmeans - Speed up using triangle inequality
• Stanford CS221 KMeans Handout - KMeans Compared with EM Algorithm

References

Andreas Krause. (2016). Learning and Intelligent Systems - Series 4, Clustering and KMeans.

Arthur, D., & Vassilvitskii, S. (2007). k-means++: the advantages of careful seeding. Proceedings of the Eighteenth Annual ACM-SIAM Symposium on Discrete Algorithms, 1027–1035.

Neptune Kmeans. (2024). K-Means Clustering Explained.

Wiki Kmeans. (2024). K-means clustering — Wikipedia, The Free Encyclopedia.