Researchers at Meta FAIR have released e-means, an algorithm presented as a direct replacement for Lloyd’s k-means algorithm. According to the authors, it achieves in 1.4 to 2.6 epochs the loss that Lloyd’s k-means obtains in 25 epochs. It is available in FAISS 1.16.0, under the module faiss.contrib.e_means.
- e-means atteint en 1,4 à 2,6 époques la perte que k-means de Lloyd obtient en 25 époques.
- Il remplace le recalcul complet par une moyenne mobile qui atténue progressivement les anciennes statistiques.
- Ses paramètres sont fixés par des lois d’échelle, et il est disponible dans FAISS 1.16.0 sous licence MIT.
- Les auteurs annoncent des gains de 10 à 45 fois sur GPU face à Flash-KMeans, et de 8 à 15 fois sur CPU face à Faiss k-means.
k-means is a clustering method that partitions data into groups, each represented by a center called a centroid. For those building vector indexes, quantizers, or data preparation pipelines, this computation is often the slowest offline step, and it is precisely the one targeted by e-means.
A moving average rather than a full recomputation
All variants of k-means maintain, for each group, a sum S and a count N, the centroid being the quotient of S by N. They differ mainly in how S and N evolve. Lloyd’s algorithm recalculates these values from the entire dataset at each iteration. MacQueen, in 1967, continuously adds batch statistics, so that old assignments never disappear.
e-means adopts a middle path: it attenuates old statistics at a rate α, according to the formula S ← (1−α)S + αs, and similarly for N. Old assignments thus gradually fade away, and centroids follow recent data. The Lloyd case corresponds to a full batch with α = 1. On the DEEP dataset, this case is 1.8% above the best loss achieved.
The authors also tested classic gradient descent. In its best configuration, it remains 1.6% above EMA on DEEP and 0.4% on SIFT. They hypothesize that the EMA method wins because centroids always remain inside the convex hull of their assigned points.
Parameters set by scaling laws
Two scaling laws, empirically fitted, determine the batch size B and the initial step α0 from the dataset size, number of clusters, and number of epochs. On held-out datasets, the median regret reaches 0.07% for the step law. The step then follows a cosine schedule, from α0 down to 0.1α0. Two other mechanisms complete the setup: the revival of ‘dead’ clusters, i.e., those that fall below a size threshold and are revived by splitting a large cluster, and a finalization on the exact mean of the last epoch.
According to the authors, the reported performance is 10 to 45 times faster on GPU than Flash-KMeans, and 8 to 15 times faster on CPU than Faiss k-means. On the VIBE benchmark, the 11 embedding datasets required between 1.6 and 2.4 epochs. For scaling, the DINO10B dataset, composed of 10 billion 1,024-dimensional embeddings, was partitioned into 16.7 million clusters in 2.5 epochs. GPU results were obtained on H100s. For configurations with 16 million centroids, float32 accumulators reach 64 GiB, requiring centroids to be distributed across multiple GPUs.
A single-threaded CPU demo, published in the announcement thread, gives 4.39 s for Lloyd over 25 epochs, versus 0.43 s for e-means over 2.2 epochs. e-means is distributed under the MIT license, via FAISS.
