K-Means
1 Introduction
K-means is a clustering algorithm that assumes the existence of \(K\) clusters, where \(K\) is a hyperparameter that must be chosen in advance by the programmer. Each cluster has a center point given by the vector \({\bf m}_k\). For a given set of \(N\) data points \({\bf x}^{(1)}, {\bf x}^{(2)}, ..., {\bf x}^{(n)}\), the goal of K-means is to minimize the total Euclidean distance of all data points to their assigned cluster centers. Formally, we want to minimize the cost:
\[\min_{\{{\bf m}_k\},\{{\bf r}^{(n)}\}}\ \sum_{n=1}^N\sum_{k=1}^K r_{k}^{(n)} || {\bf m}_k -{\bf x}^{(n)}||^2 \nonumber\] where \({\bf r}^{(n)}\) is a one-hot vector with \(r_k^{(n)} \in \{ 0, 1 \}\), and \(r_{ k}^{(n)}=1\) means that \({\bf x}^{(n)}\) is assigned to cluster \(k\).
That is, if \({\bf x}^{(n)}\) is assigned to cluster \(\hat k\), \[{\bf r}^{(n)} = \underbrace{[0,0,...,1,...,0]^\top}_{\text{Only $\hat k$-th entry is 1}}\]
At this point in the course, we are accustomed to seeing a cost function, taking its gradient, and using gradient descent to learn parameter values, but that is actually not the approach we will take here. It is awkward to take a gradient when the \({\bf r}^{(n)}\) are discrete valued, and it would be intractable to look at all possible assignments of points. (Even just for \(K=2\) there are \(\mathcal{O}(2^N)\) possible combinations.)
Fortunately, there is a method we can use that, like gradient descent, will gradually approach stable values for the cluster centers (\({\bf m}_k\)) and point assignments (\({\bf r}^{(n)}\)). (And in fact, the discovery of this method predates the formal cost function!) It takes advantage of a pair of key observations:
If the \({\bf m}_k\) were fixed, it would be straightforward to find values of \({\bf r}^{(n)}\) that minimize the cost.
If the \({\bf r}^{(n)}\) were fixed, it would be straightforward to find values of \({\bf m}_k\) that minimize the cost.
Let’s look at each of these observations in turn.
2 The Assignment Step
We claim that if the cluster centers (\({\bf m}_k\)) were fixed, it would be straightforward to find values of the assignments (\({\bf r}^{(n)}\)) that minimize the cost.
Since each data point \({\bf x}^{(n)}\) is fixed and the cluster centers \({\bf m}_k\) are also fixed, we can simply calculate the distance from our point to each center and assign it to the center that is closest:
\[\begin{align*} r_{k}^{(n)} = \left\{ \begin{array}{ll} 1 & \mbox{if } k = \arg\min_j\|{\bf x}^{(n)} - {\bf m}_j\|^2 \\ 0 & \mbox{otherwise} \end{array} \right\} \end{align*}\]
Performing this reassignment over every point in the dataset gives what we call the Assignment Step, the first of two alternating steps in the K-means algorithm.
Question: After performing the Assignment Step on the dataset \(\{(1,2), (2,3), (4,4), (4,5)\}\) with cluster centers \(\{(1,1), (3,3)\}\) (plotted in the figure below), what is the value of each \({\bf r}^{(n)}\)?
Assignment Step Problem
Answer:
For each data point, we calculate its Euclidean distance to each cluster center and assign it to the cluster that is closer.
\((1,2)\): \(\sqrt{(1-1)^2+(2-1)^2}=1 \lt \sqrt{(1-3)^2+(2-3)^2}\approx 2.24 \implies\) \({\bf r}^{(1)} = [1,0]\)
\((2,3)\): \(\sqrt{(2-1)^2+(3-1)^2}\approx 2.24 \gt \sqrt{(2-3)^2+(3-3)^2}= 1 \implies\) \({\bf r}^{(2)} = [0,1]\)
\((4,4)\): \(\sqrt{(4-1)^2+(4-1)^2}\approx 4.24 \gt \sqrt{(4-3)^2+(4-3)^2} \approx 1.41 \implies\) \({\bf r}^{(3)} = [0,1]\)
\((4,5)\): \(\sqrt{(4-1)^2+(5-1)^2}= 5 \gt \sqrt{(4-3)^2+(5-3)^2}\approx 2.24 \implies\) \({\bf r}^{(4)} = [0,1]\)
Redrawing our plot with the above point assignments gives:
Assignment Step Solution
3 The Refitting Step
We claim that if the assignments (\({\bf r}^{(n)}\)) were fixed, it would be straightforward to find values of the cluster centers (\({\bf m}_k\)) that minimize the cost.
This can be shown by taking the gradient of the cost with respect to the cluster centers while holding the point assignments fixed:
\[\begin{align*} {\bf 0} =& \nabla_{{\bf m}_\ell} \left( \sum_{n=1}^N\sum_{k=1}^K r_{k}^{(n)} || {\bf m}_k -{\bf x}^{(n)}||^2 \right) \\ =&2\sum_{n=1}^N r_{\ell}^{(n)} ({\bf m}_\ell -{\bf x}^{(n)})\quad \\ \implies {\bf m}_\ell =&\frac{\sum_n r_{\ell}^{(n)}{\bf x}^{(n)}}{\sum_n r_{\ell}^{(n)}} \end{align*}\]
This has the effect of placing each center at the average position of all points assigned to that cluster, or in other words, at the cluster’s centroid.
Moving the center of each cluster to its centroid gives us the Refitting Step, the second of the two alternating steps in K-means.
Question: Continue the example in the last question box by now performing the Refitting Step. What are the new positions of each cluster center after refitting?
Answer:
We move each cluster center to the centroid of its assigned points, taking the average along each coordinate.
Cluster 1 (square) has only one assigned point at \((1,2)\), so its center moves to \((1,2)\).
For Cluster 2 (triangle), we take the average among its three data points for each coordinate: \((\frac{2+4+4}{3}, \frac{3+4+5}{3}) \approx (3.33, 4)\)
Redrawing our plot with new cluster center positions gives:
Refitting Step Solution
Question: Given the data points and their assignments, can we complete the Refitting Step without knowing the previous locations of the cluster centers?
Answer:
Yes, we can! Unlike in gradient descent, we are not moving a small amount from the previous position. Instead, we are moving our cluster centers directly to their centroids.
4 K-Means – Putting It All Together
The K-Means algorithm is as follows:
Initialize the values of the cluster centers. This may be done randomly or may, for example, by spreading the initial clusters evenly throughout the data space. (Strictly speaking, we could instead choose an initial assignment for each data point; it would just have the effect of swapping the order of steps 2 and 3.)
Perform the Assignment Step described above, treating the cluster centers as fixed.
Perform the Refitting Step described above, treating the point assignments as fixed.
Repeat steps 2 and 3 until convergence (i.e. until the point assignments stop changing).
In other words, after initialization, K-means alternates back and forth between the Assignment and Refitting steps until it arrives at a final, unchanging clustering. This optimization process is an example of block coordinate descent, in which one “block” of variables is optimized while the other block(s) remain fixed. We refer to one pass through both of these steps as one iteration of K-means.
In the interactive figure below, you can play with different initializations of K-means for up to \(K=5\). Start by clicking on the plot to set the initial locations of the cluster centers. Then, you can repeatedly press the “Step” button to step through one step at a time of K-means, alternating between Assignment and Refitting. Observe how the point assignments change and the centers move, gradually reaching a state where nothing changes on additional “Steps”. This indicates that convergence has been reached. At any point, you may press the “Clear All” button to start over.
Can you find two initializations that start with the same number of cluster centers, but result in a different clustering at convergence?
K-Means, One Step at a Time
Question: In the preceding questions, we performed one Assignment and one Refitting step. For this example, has K-means now converged? If not, perform one more full iteration of K-means.
Answer:
No, K-means has not converged. After the first Refitting Step, the cluster centers are at \((1,2)\) and \((3.33, 4)\), but the point assignments from the preceding Assignment Step are not yet stable with respect to these new center positions. (I.e. They will change if another Assignment Step is performed.) We therefore perform one more full iteration.
Assignment Step (iteration 2)
We again assign each point to its nearest cluster center:
\((1,2)\): \(\sqrt{(1-1)^2+(2-2)^2}=0 \lt \sqrt{(1-3.33)^2+(2-4)^2}\approx 3.07 \implies\) \({\bf r}^{(1)} = [1,0]\)
\((2,3)\): \(\sqrt{(2-1)^2+(3-2)^2}=\sqrt{2}\approx 1.41 \lt \sqrt{(2-3.33)^2+(3-4)^2}\approx 1.67 \implies\) \({\bf r}^{(2)} = [1,0]\) (This one changed.)
\((4,4)\): \(\sqrt{(4-1)^2+(4-2)^2}=\sqrt{13}\approx 3.61 \gt \sqrt{(4-3.33)^2+(4-4)^2}\approx 0.67 \implies\) \({\bf r}^{(3)} = [0,1]\)
\((4,5)\): \(\sqrt{(4-1)^2+(5-2)^2}=\sqrt{18}\approx 4.24 \gt \sqrt{(4-3.33)^2+(5-4)^2}\approx 1.20 \implies\) \({\bf r}^{(4)} = [0,1]\)
The assignment for \((2,3)\) has changed from \([0,1]\) to \([1,0]\), so K-means has not yet converged.
Refitting Step (iteration 2)
We move each cluster center to the centroid of its assigned points:
Cluster 1 (square): \((\frac{1+2}{2}, \frac{2+3}{2}) = (1.5, 2.5)\)
Cluster 2 (triangle): \((\frac{4+4}{2}, \frac{4+5}{2}) = (4, 4.5)\)
Redrawing our plot after this iteration gives:
After One More Iteration of K-Means
Has K-means converged now?
4.1 Convergence of K-Means
We have stated that K-means carries on alternating its Assignment and Refitting steps until convergence is reached. But does K-means always converge, or is it possible that it sometimes runs forever? Take a moment to consider what is happening in each of the two repeating steps:
- Whenever point reassignment occurs, points are always assigned to the closest center. This can only reduce the total cost.
- Whenever refitting occurs, the cluster center moves to the centroid of its currently assigned points. We have shown that this minimizes cost for the assignment, and can therefore only reduce the total cost.
- Whenever nothing changes in either step, we can be sure that nothing will ever change again because K-means is deterministic, and there are no changes in the initial conditions.
Taken together, we can infer from these statements that the cost of K-means will decrease monotonically until convergence. Convergence will always occur after a finite number of iterations because there are a finite number of point assignments, and we will never revisit old assignments.
One important caveat is that the algorithm does not guarantee convergence to the clustering with the lowest possible cost. In general, the cost function may have multiple local minima, and K-means may converge to a clustering that sits at a local minimum in the cost, rather than at the global minimum. You may have observed this when experimenting with the interactive figure above, in which it is possible to reach different final clusterings for the same number of clusters. For example, for \(K=3\), we can arrive at both of the following clusterings after convergence:
3 Clusters - Clustering A
3 Clusters - Clustering B
Although Cluster A has a lower total cost (approximately 162 compared to 186), it is still possible to end up converged on Cluster B. Which of the two we end up with depends on the initial cluster centers provided to K-means. For this reason, it is common in practice to run K-means multiple times with different random initializations to see if a better (lower cost) clustering can be obtained.
4.2 Limitations of K-Means
An important assumption baked into K-means is that clusters are (hyper)spherical (or circular in two dimensions). We can observe this by noting that the cost is based on the Euclidean distance, which equally penalizes all points at a given radius from the cluster center. In other words, the contours of the K-means cost function are spherical shells.
As a consequence, K-means is not well-suited to discovering clusters of irregular shapes. For example, the data in the figure below was generated from two Gaussian functions, one circular and one with much more variance along the horizontal axis than the vertical. Yet, if you experiment with the figure for \(K=2\), you will see that K-means is not able to neatly split these clusters.
A Troublesome Dataset for K-Means
To get around this limitation, sometimes practitioners will deliberately set K to a higher value than necessary, and then join some of the clusters together after K-means has completed. We will not discuss such methods in detail, but you can observe how they might be useful by experimenting with the above figure with \(K=5\).
Instead, in the next part of this chapter, we will discuss Gaussian Mixture Models, which are naturally well-suited to finding such clusters, without the need for post-hoc adjustments.
5 Soft K-Means
Soft K-means is a less commonly used extension of K-means that relaxes the assumption that all data points must belong to exactly one cluster. Instead, each cluster has some amount of responsibility for each data point, in the range \([0,1]\), where a responsibility of 0 indicates that the point does not at all belong to the cluster, and a responsibility of 1 indicates that the point belongs entirely to the cluster. To implement this change, we reformulate each \({\bf r}^{(n)}\) so that instead of being one-hot, their values must sum to one, giving the responsibility of each cluster for point \(n\).
We would still like it to be the case that the closest cluster centers are the most responsible for a point, so we need a function for both scaling responsibility inversely with distance and ensuring that the sum of all responsibilities is one. The softmax function (introduced in Chapter 4) can accomplish this for us. We then adjust the Assignment and Refitting steps as follows:
Assignment: Each data point \(n\) is given a soft “degree of assignment” to each cluster mean \(k\), based on responsibilities: \[\begin{align*} r_k^{(n)} &= \frac{\exp[-\beta \|{\bf m}_k- {\bf x}^{(n)}\|^2]}{\sum_{j=1}^K \exp[-\beta \|{\bf m}_j- {\bf x}^{(n)}\|^2]} \\ \implies {\bf r}^{(n)} &= \text{softmax}(-\beta \{\|{\bf m}_k- {\bf x}^{(n)}\|^2\}_{k=1}^K) \end{align*}\]
Refitting: The coordinates of each cluster center are adjusted to match the sample means of the data points they are responsible for: \[{\bf m}_k =\frac{\sum_n r_k^{(n)}{\bf x}^{(n)}}{\sum_n r_k^{(n)}}\]
Here, \(\beta\) is a hyperparameter that determines the “hardness” of K-means. When \(\beta\) is small, we get more mixing of responsibilities, and as \(\beta \to \infty\), we get back our original one-hot \({\bf r}^{(n)}\).
We can now refer to our original formulation, in which each data point can only be assigned to a single cluster, as “Hard K-means”. However, when you see a reference to just “K-means” in the wild, it is generally a reference to Hard K-means (the more commonly used method), unless otherwise specified. We will adopt the same convention here, referring to Soft K-means when necessary, and just K-means otherwise.
6 Summary
We have introduced K-means, the first of two unsupervised clustering methods in this chapter. Although the Gaussian Mixture Models (GMMs) that follow are strictly more powerful in the kinds of clusters that can be captured, K-means remains a popular method for its speed and simplicity. You will also see that many parts of GMMs are directly analogous to the concepts introduced in this section.