Theory of K-Means Clustering

K-Means Clustering is a fundamental algorithm in unsupervised learning that aims to partition a dataset into k distinct clusters. The theory behind K-Means is rooted in optimization and linear algebra, making it both powerful and widely applicable. This article explores the algorithm’s theoretical aspects, focusing on the mathematical principles that guide its operation.

1. The K-Means Objective Function

At the heart of the K-Means algorithm lies an optimization problem. The algorithm seeks to minimize the within-cluster sum of squares (WCSS), which is also known as the inertia or distortion. This objective function is defined as follows:

1.1 Objective Function Definition

Given a dataset $X = \{x_1, x_2, \dots, x_n\}$, where each $x_i$ is a point in $\mathbb{R}^d$, let $z_i \in \{1, \dots, k\}$ be its cluster assignment and let $\mu_1, \dots, \mu_k$ be the centroids. K-Means seeks a partition and centroids that minimize:

$$J(z, \mu) = \sum_{i=1}^{n} \|x_i - \mu_{z_i}\|_2^2$$

This is the within-cluster sum of squares, called inertia_ by scikit-learn. It depends on feature units and weights; it is not directly comparable across differently scaled representations or different datasets.

1.2 Minimizing the Objective Function

The K-Means algorithm alternates between two steps to minimize the objective function:

1. Assignment Step:

   • Each data point $x_i$ is assigned to the cluster whose centroid is closest, according to the Euclidean distance:
   $$\text{Assign } x_i \text{ to cluster } C_j \text{ such that } j = \arg\min_{j'} \|x_i - c_{j'}\|$$

2. Update Step:

   • After all points are assigned to clusters, each centroid $c_j$ is updated to be the mean of the points in its cluster:
   $$c_j = \frac{1}{|C_j|} \sum_{x_i \in C_j} x_i$$

These steps are repeated iteratively until the centroids no longer change significantly or a maximum number of iterations is reached.

2. Distance Metric: Euclidean Distance

The K-Means algorithm uses the Euclidean distance as the primary metric to determine the similarity between data points and centroids. The Euclidean distance between two points $x_i$ and $x_j$ in $\mathbb{R}^d$ is given by:

$$\|x_i - x_j\| = \sqrt{\sum_{l=1}^{d} (x_{il} - x_{jl})^2}$$

Where:

• $x_{il}$ and $x_{jl}$ are the $l$-th components of the vectors $x_i$ and $x_j$, respectively.

Euclidean distance is appropriate for K-Means because it emphasizes the overall magnitude of differences between points, which aligns well with the algorithm's objective of minimizing the sum of squared distances.

3. Convergence of K-Means

3.1 Convergence Guarantee

With deterministic tie handling, Lloyd's alternating assignment and update steps never increase $J$ and reach a fixed partition in finitely many distinct assignments. Implementations also stop at a tolerance or iteration limit. Being bounded below alone would not prove finite convergence, and numerical ties can affect which equivalent solution is returned.

3.2 Local Minima

However, K-Means does not guarantee convergence to the global minimum of the objective function. Instead, it may converge to a local minimum, where no further improvements are possible, but the solution is not necessarily the optimal one. The final clusters depend heavily on the initial placement of centroids, which can lead to different local minima.

To mitigate this issue, K-Means is often run multiple times with different initializations, and the solution with the lowest objective function value is selected.

4. Complexity of K-Means

4.1 Computational Complexity

The computational complexity of K-Means primarily depends on the number of data points $n$, the number of dimensions $d$, the number of clusters $k$, and the number of iterations $t$. Each iteration of K-Means involves two main operations:

1. Assignment Step: The complexity of assigning each point to the nearest centroid is $O(nkd)$, as we must compute the distance between each point and each centroid.

2. Update Step: The complexity of updating the centroids is $O(nd)$, since we must compute the mean of the points assigned to each cluster.

Therefore, the overall complexity of K-Means per iteration is $O(nkd)$, and for $t$ iterations, the total complexity is $O(tnkd)$.

4.2 Scalability

K-Means scales well with large datasets, but the number of clusters $k$ and the number of dimensions $d$ can significantly impact performance. Optimizations like the k-means++ initialization and efficient distance calculations (e.g., using spatial data structures) can improve scalability.

5. Variants of K-Means

Several variants of K-Means have been developed to address its limitations and improve performance:

5.1 K-Means++

K-Means++ chooses the first centroid randomly and samples later centroids with probability proportional to each point's squared distance from its nearest selected centroid. It usually improves starting solutions, but repeated starts remain useful and it does not guarantee the global optimum.

5.2 Mini-Batch K-Means

Mini-Batch K-Means is an extension of the standard K-Means algorithm that improves scalability by using mini-batches of data to update the centroids at each iteration. This reduces the computational burden when dealing with large datasets, as only a subset of the data is used in each update step.

5.3 Fuzzy C-Means (Soft K-Means)

Fuzzy C-Means assigns graded membership weights rather than one hard label. Those weights are constrained optimization memberships, not automatically calibrated probabilities. Its objective and update rules differ from ordinary K-Means.

6. Mathematical Insights: Why K-Means Works

6.1 Geometric Interpretation

The geometric interpretation of K-Means is that it partitions the data space into Voronoi cells around each centroid. Each cell contains all points that are closer to its centroid than to any other centroid. This partitioning minimizes the variance within each cluster, which is the sum of squared distances from each point to its centroid.

6.2 Optimization Perspective

K-Means is best described as alternating minimization of the within-cluster sum of squares. It is related to a limiting hard-assignment case of equal, spherical Gaussian mixtures, but ordinary K-Means is not the same as fitting a Gaussian mixture with expectation-maximization: it provides neither soft posterior memberships nor estimated covariance and mixing parameters.

7. Conclusion

The theoretical foundations of K-Means Clustering reveal why this algorithm is so effective for a wide range of clustering tasks. By iteratively minimizing the sum of squared distances between data points and their nearest centroid, K-Means effectively partitions data into meaningful clusters. Despite its simplicity, K-Means is rooted in rich mathematical concepts, including optimization, geometry, and linear algebra.

Understanding the theory behind K-Means empowers data scientists to apply the algorithm more effectively, recognize its limitations, and explore its various extensions. As we move forward, we will explore the practical implementation of K-Means using popular machine learning frameworks like Scikit-learn, PyTorch, and TensorFlow.