e-means: outperforming k-means (in one-tenth the epochs)

Lloyd’s algorithm, a full-batch method, is still the standard for k-means clustering. The Meta FAIR team replaces it with mini-batch updates in which each cluster’s running sum and count are exponential moving averages, and the centroid is their ratio. Scaling laws relating decay rate and batch size to final loss fix both hyperparameters for each problem size, so no sweep is needed.

On a range of embedding datasets, e-means reaches a loss below Lloyd’s at 25 epochs within 1.4 to 2.6 epochs, which the authors report as 10 to 45 times faster on GPU than leading k-means implementations, often beating k-means++ as well. It partitioned 10 billion DINO embeddings into 16 million clusters in 2.5 epochs. The code is released in Faiss.