I went looking in scikit-learn’s K-means and found a state it can’t get out of: the code that handles empty clusters creates the empty cluster it’s handling.

K-means

The k-means algorithm is a popular unsupervised machine learning technique used for clustering data. Its goal is to partition a set of data points into k clusters where each data point belongs to the cluster with the nearest mean. The standard algorithm was first proposed by Stuart Lloyd of Bell Labs in 1957.

K-means has two inputs: the number of clusters k, and the data points. Some implementations also support optional initial centroids. These are the main steps:

  1. Initialization — choose k initial cluster centroids from the data points (this could be random or based on some statistics).
  2. Assignment step — assign each data point to the nearest cluster centroid based on a distance metric (typically Euclidean distance). This step creates clusters.
  3. Update step — recalculate the centroids of the clusters based on the current assignment of data points. The new centroid is the mean of all data points assigned to that cluster.
  4. Repeat — alternate between the assignment step and the update step until the centroids no longer change significantly, or a maximum number of iterations is reached, or some other stop condition.

The bug

I noticed that a weird behaviour happens if we encounter an empty cluster in a very particular condition. An empty cluster can occur due to several reasons, but mainly poor initialization and data distribution.

In scikit-learn, to solve the problem of empty clusters, they decided to associate at least one data point to each empty cluster, and they were relocating data points that are far away from their centroids to do that job.

The idea is good, but there was no check that the algorithm is not generating an empty cluster when trying to fill an empty cluster. And I noticed that if the following three conditions are satisfied, the algorithm gets stuck in an infinite loop:

  1. There are duplicate input data points.
  2. The number of clusters is equal to the total number of data points.
  3. Very specific initial positions.

Here is a minimal example that reproduced the bug:

Scatter plot of the failing case: green input data points, several of them duplicated at the same coordinates, and blue cluster centres
Green marks the input data, blue the cluster centres. Note the duplicated points sharing coordinates.

The fix

There are several strategies to address this issue:

  • Ensure that the relocation algorithm does not create an empty cluster when attempting to fill one.
  • In certain cases, it may be beneficial to remove duplicate points.
  • Reduce the number of clusters used if the specified number of clusters exceeds the number of unique input data points.

You can learn more about the strategy that was finally used by reading the discussion of the resolved issue: scikit-learn#28055.

To propose a fix for an open-source algorithm in general, one way is to open an issue in their GitHub repo. Scikit-learn did a good job documenting how you should do it: scikit-learn contributing. I followed the steps, and I opened an issue that was resolved quickly, in my opinion. The fix, scikit-learn#28165 by Jérémie du Boisberranger, shipped in scikit-learn 1.4.1 in February 2024.

Sources