# Finding an Infinite Loop in Scikit-Learn’s K-Means

2024-04-29 · https://ibraverse.ca/tech/scikit-learn-kmeans-infinite-loop/

> An input that makes scikit-learn’s K-means loop forever: the bug in its empty-cluster handling, the report I filed, and the fix.


I went looking in
[scikit-learn's K-means](https://scikit-learn.org/stable/modules/generated/sklearn.cluster.KMeans.html)
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](https://en.wikipedia.org/wiki/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](https://ibraverse.ca/tech/scikit-learn-kmeans-infinite-loop/01-kmeans-bug-repro_hu_9aeef2d22951f5a4.webp)
*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](https://github.com/scikit-learn/scikit-learn/issues/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](https://scikit-learn.org/stable/developers/contributing.html).
I followed the steps, and I opened an issue that was resolved quickly, in my
opinion. The fix, [scikit-learn#28165](https://github.com/scikit-learn/scikit-learn/pull/28165)
by Jérémie du Boisberranger, shipped in scikit-learn 1.4.1 in February 2024.

## Sources

- scikit-learn — [`KMeans` documentation](https://scikit-learn.org/stable/modules/generated/sklearn.cluster.KMeans.html)
- scikit-learn — [issue #28055](https://github.com/scikit-learn/scikit-learn/issues/28055), the bug and its discussion
- scikit-learn — [pull request #28165](https://github.com/scikit-learn/scikit-learn/pull/28165), the fix, in the [1.4.1 release notes](https://scikit-learn.org/stable/whats_new/v1.4.html)
- scikit-learn — [contributing guide](https://scikit-learn.org/stable/developers/contributing.html)
- Wikipedia — [Bell Labs](https://en.wikipedia.org/wiki/Bell_Labs), where Lloyd proposed the algorithm in 1957