Meta FAIR Introduces e-means: A Faster Alternative to k-means Clustering
Meta FAIR researchers have launched e-means, a drop-in replacement for Lloyd’s k-means algorithm. This new clustering algorithm achieves the loss that Lloyd’s k-means reaches in 25 epochs within just 1.4 to 2.6 epochs. It is included in FAISS 1.16.0 under faiss.contrib.e_means. For developers working on vector indexes, quantizers, or data-curation pipelines, clustering often represents the slowest offline step, which e-means aims to optimize.
TL;DR
- Size: It is an algorithm tested from 100K to 10 billion points, supporting up to 16.7M clusters.
- Runs on: CPU or GPU, compatible with PyTorch tensors or NumPy. Paper GPU runs utilize H100s.
- Performance: Lower loss than Lloyd’s k-means in approximately one-tenth the epochs across 24 dataset configurations (self-reported).
- Best: 10-45x GPU speedup over Flash-KMeans at equal loss.
- Worst: 16M-centroid runs require 64 GiB of float32 accumulators, necessitating centroid sharding across GPUs.
What is e-means?
e-means is a mini-batch k-means algorithm developed by Meta FAIR that updates centroids using exponential moving averages.
How does e-means work?
Each k-means variant maintains a sum S and a count N for each cluster, with the centroid calculated as S divided by N. The methods differ in how S and N are updated.
- Lloyd: Recomputes S and N from the entire dataset during each iteration.
- MacQueen (1967): Continuously adds batch statistics, leading to stale assignments persisting.
- e-means: Decays old statistics at a rate α: S ← (1−α)S + αs, with a similar approach for N.
This decay allows older assignments to fade, enabling centroids to better reflect recent data. The research team notes that Lloyd’s method is a corner case, operating with a full batch and α = 1, which results in a loss 1.8% higher than the best loss.
The team also tested plain gradient descent as an alternative, which performed 1.6% (DEEP) and 0.4% (SIFT) worse than e-means at its best settings. The authors suggest that e-means may have an advantage because centroids consistently remain within the hull of their assigned points.
What replaces hyperparameter tuning?
Two fitted scaling laws determine the batch size B and initial step α0 based on dataset size n, cluster count k, and epochs T. On held-out datasets, the median regret was 0.07% for the step law, which then follows a cosine schedule from α0 to 0.1α0.
What else changed?
- Dead-cluster revival: Clusters below a size threshold are revived by splitting a larger cluster.
- Finalization: Training concludes with each cluster’s exact last-epoch mean.
- Fused GPU kernels: Nearest-centroid assignment is fused, leading to reduced epoch times.
How fast is e-means in benchmarks?
All performance metrics below are self-reported by the authors.
- Epochs: e-means achieves Lloyd’s 25-epoch loss in 1.4 to 2.6 epochs.
- GPU: 10-45x wall-clock speedup compared to Flash-KMeans.
- CPU: 8-15x speedup over Faiss k-means.
- VIBE benchmark: All 11 embedding datasets required 1.6 to 2.4 epochs.
- Billion scale: DINO10B, which includes 10 billion 1,024-dimensional embeddings, was split into 16.7M clusters in 2.5 epochs.
A release thread indicates that a single CPU thread demo showed Lloyd taking 4.39 seconds for 25 epochs, while e-means completed 2.2 epochs in just 0.43 seconds. One step of Lloyd processed 100,000 points in 0.176 seconds, while e-means processed 512 points in 0.001 seconds.
Mini-batch noise also aids in avoiding poor local minima. In a 5-blob toy example, Lloyd was stuck in 46% of 200 attempts, while e-means experienced no stagnation.
Revival is significant at larger k. In the DEEP dataset with 16,384 clusters, proportional splitting resulted in 47 dead clusters (0.3%). Without revival, 4,313 clusters (26.3%) became inactive.
How does e-means compare with other k-means options?
Key Takeaways
- e-means replaces Lloyd’s full-batch update with EMA mini-batches.
- It matches the performance of 25 Lloyd epochs in just 1.4 to 2.6 epochs.
- It offers 10-45x speed improvements on GPU and 8-15x on CPU (self-reported).
- No tuning is required, as scaling laws determine batch and step sizes.
- e-means is available now in FAISS 1.16.0.
Related link: e-means.github.io
Source: MarkTechPost