Meta's e-means: 45x Faster K-Means with EMA & Scaling Laws

Background: The Persistent Cost of Lloyd's K-Means

K-Means has been a cornerstone of clustering for nearly 70 years, yet the classic Lloyd algorithm still requires repeated full-dataset scans — typically 25 epochs in Faiss — to converge centroids. Each epoch computes distances from every point to every centroid, making billion-scale clustering prohibitively expensive.

Mini-Batch History and Its Limitations

Mini-batch K-Means (dating to 1967) updates centroids after each batch, avoiding full scans. However, it accumulates historical statistics indefinitely; early inaccurate assignments persist and degrade final quality. Mini-batch methods often converge faster but yield higher error than well-trained Lloyd, and require careful tuning of batch size and learning rate.

e-means: Redesigning Centroid Updates with EMA

Meta FAIR's e-means replaces the cumulative average with exponential moving average (EMA) updates. For each mini-batch, it computes per-cluster sample counts and vector sums, applies EMA to both, and sets the centroid to the ratio of the two moving averages. This gives recent assignments higher weight while gradually forgetting early noisy assignments.

The authors note EMA was previously used in vector quantization codebook updates. They argue EMA keeps centroids inside the convex hull of assigned points, providing stability at large update steps where stochastic gradient descent (SGD) becomes unstable. Experiments confirm SGD diverges with large steps on big batches, while EMA remains stable.

Escaping Local Optima via Stochastic Mini-Batch Updates

Lloyd's deterministic updates can trap centroids in poor local minima. Mini-batch noise helps e-means escape. In a controlled experiment with two imbalanced clusters (target k=8), Lloyd assigned two centroids to the small cluster and too few to the large one in 26% of 200 random restarts; e-means reduced this to 3.5%. On five linearly arranged clusters, Lloyd failed 46% of the time vs. 0% for e-means. An animated figure (Figure 2) illustrates centroids breaking free from suboptimal configurations.

Dead Cluster Revival by Splitting Large Clusters

Centroids that receive too few points become "dead". e-means checks every five batches and revives dead centroids by splitting the largest active cluster rather than random reinitialization. On the DEEP dataset, without revival 26.3% of centroids had <10 points; with splitting this dropped to 0.3%. An ablation over 3 epochs showed test MSE increased ~2.24% without revival (Figure 3 illustrates the splitting mechanism).

Cosine Annealing, Initialization, and Final Polish

Centroids are initialized by random sampling from data with negligible initial EMA weight. A cosine annealing schedule gradually reduces the EMA step size to prevent late-stage oscillation. After the last epoch, a single full pass recomputes centroids as exact means of their assigned points, lowering training error at minimal cost. For streaming data, a slower EMA can improve test error.

Scaling Laws Eliminate Hyperparameter Search

Batch size and step size jointly determine speed and accuracy. The authors ran ~75,000 experiments varying dataset size, cluster count, and epochs, then fitted two Scaling Laws that predict optimal batch size and initial step size. Fitted on DEEP, the laws transfer to SIFT, SPACEV, and text embeddings; on held-out data the median error increase versus per-dataset tuning was only ~0.19%. The laws also offer a throughput-oriented vs. quality-oriented batch choice.

Large-Scale Benchmarks: Up to 45.5× GPU Speedup

Experiments compare e-means against Faiss K-Means, Flash-KMeans, and other mini-batch baselines on image (DINO, SIFT) and text embeddings. Using MSE as metric, e-means reaches Flash-KMeans' 25-epoch test error in 1.4–2.6 epochs, translating to 9.5–45.5× wall-clock speedup on GPU. The peak 45.5× occurs on 1 billion DINO vectors. CPU speedups range 8–15×. A fused GPU kernel merges distance computation and nearest-centroid selection, avoiding the massive distance matrix.

10 Billion Vectors, 16.77 Million Clusters

Scaling further, the team built DINO10B: 10 billion 1024-dim vectors (~10 TB) from 60M images using DINOv3 patch features. With exact assignment, Lloyd 25 epochs took 6,784 GPU hours; e-means 2.5 epochs took 779 GPU hours and achieved lower test error. Using inverted file index (IVF) approximate assignment, e-means hit the target error in 3 epochs (~23.2 GPU hours), finishing in ~3 hours on a single 8×H100 server (Table 2).

Downstream Impact: Faster Index Training, Better Recall

e-means is now in the Faiss repository under faiss.contrib.e_means (not yet in official releases). In 11 vector datasets, IVF indexes built from e-means centroids achieve higher Recall@10 at 1% scan rate for the same training time, or match fully-trained baselines in less time. Final recall is comparable or better (Figure 6).

Conclusion

By combining EMA updates, dead-cluster splitting, cosine annealing, and scaling-law hyperparameter selection, e-means delivers a drop-in replacement for Lloyd's K-Means that is dramatically faster at billion-scale while maintaining or improving clustering quality. The code and reproduction scripts are publicly available in the Faiss contrib directory.