What is k-means clustering?

As previously mentioned, many clustering algorithms don't scale to the datasets used in machine learning, which often have millions of examples. For example, agglomerative or divisive hierarchical clustering algorithms look at all pairs of points and have complexities of \(O(n^2 log(n))\) and \(O(n^2)\), respectively.

This course focuses on k-means because it scales as \(O(nk)\), where \(k\) is the number of clusters chosen by the user. This algorithm groups points into \(k\) clusters by minimizing the distances between each point and its cluster's centroid (see Figure 1).

As a result, k-means effectively treats data as composed of a number of roughly circular distributions, and tries to find clusters corresponding to these distributions. But real-world data contains outliers and density-based clusters and might not match the assumptions underlying k-means.

k-means clustering algorithm

The algorithm follows these steps:

1. Provide an initial guess for \(k\), which can be revised later. For this example, we choose \(k = 3\).

2. Randomly choose \(k\) centroids.

   

   Figure 1: k-means at initialization.

3. Assign each point to the nearest centroid to get \(k\) initial clusters.

   

   Figure 2: Initial clusters.

4. For each cluster, calculate a new centroid by taking the mean position of all points in the cluster. The arrows in Figure 4 show the change in centroid positions.

   

   Figure 3: Recomputed centroids.

5. Reassign each point to the nearest new centroid.

   

   Figure 4: Clusters after reassignment.

6. Repeat steps 4 and 5, recalculating centroids and cluster membership, until points no longer change clusters. In the case of large datasets, you can stop the algorithm before convergence based on other criteria.

Because the centroid positions are initially chosen at random, k-means can return significantly different results on successive runs. To solve this problem, run k-means multiple times and choose the result with the best quality metrics. (We'll describe quality metrics later in this course.) You'll need an advanced version of k-means to choose better initial centroid positions.

Though a deep understanding of the math is not necessary, for those who are curious, k-means is a special case of the expectation-maximization algorithm. See lecture notes on the topic from UPenn.