Understanding K-Means Clustering

Introduction

In our previous module, we concluded our exploration of advanced supervised learning methods by dissecting the kernel trick in Support Vector Machines, a powerful technique for handling non-linear data. We are now pivoting from a world where data comes with explicit labels to one where it does not. This marks our entry into unsupervised learning, a paradigm focused on discovering inherent structure, patterns, and representations directly from the data itself.

This lesson introduces one of the most fundamental and widely used unsupervised learning algorithms: K-means clustering. Our objective is to explain the algorithm, formally define its mathematical objective function, and detail the iterative optimization process it uses to find a solution. You will learn how K-means partitions data into groups and what "optimality" means in this context, providing a foundational pillar for understanding pattern discovery in unlabeled datasets.

What is Clustering?

At its core, clustering is the task of partitioning a set of data points into a number of groups, or clusters, such that points within the same cluster are more "similar" to one another than to those in other clusters. This is a cornerstone of exploratory data analysis, with applications spanning numerous domains:

• Marketing: Segmenting customers into groups with similar purchasing habits for targeted campaigns.
• Biology: Grouping genes with similar expression patterns to identify functional pathways.
• Computer Vision: Segmenting an image into regions of similar color or texture.
• Natural Language Processing: Grouping documents by topic.

To perform clustering, we first need a mathematical definition of "similarity." For algorithms like K-means, which operate on data in a vector space ${\mathbb{R}}^d$, similarity is typically quantified by a distance metric, most commonly the squared Euclidean distance.

The K-means Objective Function

The K-means algorithm aims to partition $N$ data points $\{{\mathbf{x}}_1,\dots ,{\mathbf{x}}_N\}$ into a pre-specified number, $K$, of clusters. The core idea is to find cluster centroids $\{{\mathbf{\mu }}_1,\dots ,{\mathbf{\mu }}_K\}$ and assign each data point to the nearest centroid, such that the total "within-cluster sum of squares" (WCSS) is minimized.

This objective can be expressed formally. Let $r_{nk}$ be a binary indicator variable, where $r_{nk}=1$ if data point ${\mathbf{x}}_n$ is assigned to cluster $k$, and $r_{nk}=0$ otherwise. The objective function, often denoted by $J$, is the sum of squared Euclidean distances from each point to the centroid of its assigned cluster:

The goal of K-means is to find the set of assignments $\{r_{nk}\}$ and centroids $\{{\mathbf{\mu }}_k\}$ that minimize this function $J$.

A deeper look at this objective reveals that it is equivalent to minimizing the average pairwise squared distance between points within the same cluster, a natural measure of cluster cohesion. The following resource provides a formal definition of the K-means criterion and shows this equivalence. It also makes a critical point for someone with your background: finding the global minimum of this objective is an NP-hard problem.

The Iterative Optimization Process: Lloyd's Algorithm

Since finding the exact solution to the K-means objective is computationally intractable, we rely on a heuristic iterative method known as Lloyd's Algorithm. This algorithm is a classic example of coordinate descent, an optimization strategy where you iteratively minimize a function with respect to different subsets of variables, one at a time. In this case, the two subsets of variables are the cluster assignments $\{r_{nk}\}$ and the cluster centroids $\{{\mathbf{\mu }}_k\}$.

The algorithm proceeds as follows:

1. Initialization: Randomly initialize the $K$ cluster centroids ${\mathbf{\mu }}_1,\dots ,{\mathbf{\mu }}_K$. A common strategy is to pick $K$ random data points from the dataset to serve as the initial centroids.

2. Iterative Refinement: Repeat the following two steps until the cluster assignments no longer change:

   • Assignment Step (E-step): For each data point ${\mathbf{x}}_n$, assign it to the closest cluster centroid. This step minimizes the objective $J$ with respect to the assignments $\{r_{nk}\}$, holding the centroids $\{{\mathbf{\mu }}_k\}$ fixed. $$r_{nk}=\{\begin{array}{ll}1& \text{if }k=\arg {\min }_j\Vert {\mathbf{x}}_n-{\mathbf{\mu }}_j{\Vert }^2\\ 0& \text{otherwise}\end{array}$$
   • Update Step (M-step): For each cluster $k$, recalculate its centroid ${\mathbf{\mu }}_k$ as the mean of all data points assigned to it. This step minimizes the objective $J$ with respect to the centroids $\{{\mathbf{\mu }}_k\}$, holding the assignments $\{r_{nk}\}$ fixed. $${\mathbf{\mu }}_k=\frac{{\sum }_{n=1}^N{r}_{nk}{\mathbf{x}}_n}{{\sum }_{n=1}^N{r}_{nk}}$$

The names "E-step" and "M-step" are borrowed from the more general Expectation-Maximization (EM) algorithm, which we will encounter later. For K-means, the "Expectation" corresponds to assigning points, and "Maximization" (or in this case, minimization) corresponds to updating the centroids.

The following resource provides a clear pseudocode for Lloyd's algorithm and a rigorous derivation of these two steps as a coordinate descent procedure on the objective function.

For a different perspective, the following two short video lectures walk through the exact same derivation, which may help reinforce the concepts.

Visualizing the Algorithm

The iterative nature of K-means is best understood visually. The algorithm begins with randomly placed centroids, and with each iteration, the centroids move towards the center of their assigned points, and the cluster boundaries (known as a Voronoi tessellation) shift accordingly. This process continues until the assignments stabilize.

This iterative process is guaranteed to converge. At each step (both assignment and update), the value of the objective function $J$ either decreases or stays the same. Since $J$ is bounded below by zero, the algorithm must converge to a local minimum. However, because the objective function is non-convex, this local minimum is not guaranteed to be the global minimum. The final result is sensitive to the initial placement of the centroids. For this reason, in practice, K-means is often run multiple times with different random initializations, and the solution with the lowest objective function value is chosen.

Conclusion

In this lesson, we have formally introduced K-means, a foundational algorithm in unsupervised learning. You now have a theoretical grasp of its core components.

The key takeaways are:

• Clustering is an unsupervised learning task for finding group structures in unlabeled data.
• The K-means objective function seeks to minimize the within-cluster sum of squares (WCSS), which is the total squared Euclidean distance between each data point and the centroid of its assigned cluster.
• Solving the K-means objective is NP-hard, necessitating a heuristic approach.
• Lloyd's Algorithm is an iterative optimization process based on coordinate descent. It alternates between an assignment step (assigning points to the nearest centroid) and an update step (re-computing centroids as the mean of assigned points).
• The algorithm is guaranteed to converge to a local minimum of the objective function, but its sensitivity to initialization means multiple runs are often required in practice.

This lesson has established the principle of finding structure through geometric partitioning. Our next topic will address another fundamental task in unsupervised learning: dimensionality reduction. We will begin by exploring Principal Component Analysis (PCA), a technique that finds a lower-dimensional representation of data while preserving as much variance as possible.