Meta FAIR Open-Sources the e-means Clustering Algorithm: Mini-Batch Exponential Moving Averages Replace Lloyd Iterations, and the Paper Reports Beating 25-Epoch k-means Loss in 1.4 to 2.6 Epochs

Mini-batch moving averages instead of full passes

Lloyd's algorithm, the usual way to run k-means, scans the whole dataset every epoch before recomputing each cluster's centre. A paper submitted to arXiv on November 6 by Swann Bessa and six co-authors at Meta FAIR proposes e-means: it takes one mini-batch at a time, keeps exponential moving averages of each cluster's vector sum and point count, and divides the two to get the centroid. The project page notes that setting the batch to the full dataset and the step to 1 makes one e-means iteration exactly one Lloyd iteration; the authors' hyperparameter search found that this extreme setting performs markedly worse.

Batch size and step size from a formula, with no tuning

Earlier mini-batch k-means variants needed batch size and learning rate tuned by hand. The paper fits a set of scaling laws that compute both directly from the dataset size, the number of clusters and the planned number of epochs, so no hyperparameter sweep is needed. Initialization is also changed; "dead" clusters that end up with no points are not replaced by random samples, as is commonly done, but revived by picking an existing cluster in proportion to its size and splitting it in two; the decay rate follows a cosine schedule, and the run ends on each cluster's exact last-epoch mean.

The authors' numbers, and the validation still missing

According to the abstract, across a range of embedding datasets e-means beats the loss of Lloyd k-means at 25 epochs in 1.4 to 2.6 epochs, which translates to 10-45x GPU wall-clock speedups, and it frequently beats k-means++; it partitions DINO10B, a dataset of 10 billion DINO embeddings, into 16 million clusters in only 2.5 epochs. The project page puts it as 1.2 to 3 epochs, 10-45x faster on GPU against Flash-KMeans and 8-15x faster on CPU against Faiss.

The code is in the contrib/e_means directory of facebookresearch/faiss. The paper is currently a v1 preprint that has not been peer reviewed, the figures above are the authors' own measurements, and the tests cover embedding-type data; how it performs on other data distributions awaits third-party replication.