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

AI summary

Overview

Research area: Machine learning — clustering algorithms, specifically mini-batch k-means and its use in large-scale vector quantization, retrieval and data curation.

Technical level: Advanced. The paper assumes familiarity with Lloyd's algorithm, mini-batch stochastic updates, exponential moving averages (as used in VQ-VAEs), and scaling-law fitting. The high-level claims are accessible, but the method and the hyperparameter laws are technical.

Scope: The paper proposes e-means, an EMA-based mini-batch k-means optimizer with automatically derived batch-size and decay-rate scaling laws, and benchmarks it against Lloyd k-means, Faiss, Flash-KMeans, SuperKMeans and cuML from 1M to 10B vectors.

What This Paper Is About

Lloyd's k-means is the long-standing default for clustering a fixed dataset, but it needs many full passes over the data to converge, which is expensive at web scale. Mini-batch variants are faster per step but normally land at a worse final error, and they force practitioners to guess a batch size and a step size for each new dataset.

The paper's goal is a mini-batch k-means that converges to a lower loss than Lloyd while needing far fewer epochs, with no hyperparameter sweep — the batch size and decay rate are computed from the problem size (number of points, clusters and epochs) using fitted scaling laws.

Key Contributions

1. A modern EMA optimization recipe for k-means. The authors specialize the VQ-VAE exponential-moving-average centroid rule to the canonical k-means objective by adding a cosine-annealed step size, a modified initialization, an improved dead-cluster revival strategy, and two cheap finalization stages (last-epoch mean for training loss, slow EMA for held-out loss).

2. Scaling laws for automatic hyperparameter selection. They fit power laws on the peak step α₀ and on the batch size B as functions of (k, n, T, B), giving a formulaic rule that removes the need for a hyperparameter sweep.

3. A scalable method and package. e-means is evaluated on a wide collection of embeddings datasets against Faiss on CPU and Flash-KMeans, Faiss-GPU and cuML on GPU, plus the distance-pruning method SuperKMeans on CPU, with an IVF-based approximate assignment mode for extreme scale.

4. DINO10B, a new 10-billion-vector dataset of DINO image-patch embeddings (1024 dimensions) for evaluating clustering at very large scale, which also serves as an approximate nearest-neighbor benchmark.

Main Findings

• Lower loss than Lloyd, on every dataset tested. e-means ends below Lloyd on both train and test MSE across the reported datasets and scales, while the mini-batch baselines (Sculley's and Staleness-Reduction) beat Lloyd at one epoch but converge at a higher MSE.

• Epoch efficiency. On the embeddings datasets, e-means beats the loss of Lloyd k-means at 25 epochs in 1.4 to 2.6 epochs, which the abstract translates into 10–45× GPU wall-clock speedups versus state-of-the-art k-means implementations. Table 4 reports the range as 9.5× to 45.5× against Flash-KMeans.

• Speedups grow with scale. In Table 4, reaching Flash-KMeans' 25-epoch MSE takes 2.0 epochs at 1M points with k = 16,384 on DEEP* (12.8× speedup) and 1.6 epochs at 1B points with k = 1,048,576 on DINO10B (45.5× speedup). The largest measured speedup is 45.5× and the smallest is 9.5×.

• DINO10B result. e-means partitions the 10-billion-vector DINO10B dataset into 16 million clusters in only 2.5 epochs, reaching a loss below Lloyd's at 25 epochs. In the Figure 1 description the same comparison is stated as beating the 25-epoch loss in 2 epochs.

• IVF helps only at very large n. GPU hours to reach Lloyd's test loss at 25 epochs on DINO10B subsets: at n = 100M, k = 2¹⁸, exact assignment costs 0.028 GPU-hours versus 0.030 with IVF (no gain); at n = 1B, k = 2²¹, 4.56 versus 1.05 (4.3× cheaper); at n = 10B, k = 2²⁴, 779 versus 22.8 (34× cheaper), training 16.7M centroids in less than one GPU-day.

• Beats k-means++-initialized Lloyd at convergence. Across 40 seeds on three datasets, e-means with random initialization reaches the lowest train and test losses, outperforming Lloyd with greedy k-means++ initialization by a large margin.

• Three epochs beat 25-iteration Flash-KMeans on all eight configurations in Table 5. With Flash-KMeans using greedy k-means++ initialization, e-means still has the lower mean error after 5 epochs on all eight configurations, with the margin narrowing to between 0.06% and 0.33%.

• Ablation (DEEP, n = 10⁶, k = 16,384, T = 3, medians over 30 seeds).* Starting from train MSE 0.3613 and test MSE 0.3880: removing dead-cluster revival costs +2.97% train and +2.24% test (the largest single effect); replacing revival with data resampling costs +2.06% train and +1.15% test; removing the cosine schedule costs +0.26% train and −0.06% test; removing finalization costs +0.14% train and +0.02% test; k-means++ initialization gains −0.18% train and −0.11% test.

• The scaling laws transfer. Fitted on the DEEP dataset only using 75k optimization runs, the laws show out-of-distribution regret of 0.07% (step) and 0.19% (batch, knee at 0.2%), with R² of 0.95 and 0.96 respectively. The batch law incurs a median 0.19% excess MSE on held-out datasets, close to its 0.2% target.

• Batch size at scale. For n > 2M the authors divide the batch law by 8 because throughput was equivalent with smaller batches.

• Stated limitation. The scaling laws are tuned for relatively high-dimensional datasets and are not optimal on very low-dimensional, tiny datasets; the authors note that local intrinsic dimensionality, which matters here, is generally unknown.

Methodology in Plain English

Instead of recomputing each centroid as the exact mean of everything assigned to it (as Lloyd does after a full pass), e-means keeps two running numbers per cluster — a decayed count and a decayed sum of assigned vectors — updated from small batches. The centroid is simply the ratio of the two.

Two mechanisms make this work better than a naive running mean. First, the step size is annealed over training on a cosine schedule from a peak α₀ down to a floor of 0.1·α₀, so centroids move a lot early and stabilize late. Second, every five batches the algorithm rescues "dead" clusters whose decayed count falls below 0.1 of the mean count, rebuilding them by splitting a donor cluster in the same way Faiss does.

Because the user still has to pick a peak step and a batch size, the authors sweep both over many (n, k, T) configurations, record where the validation error is lowest (for the step) or where it rises 0.2% above its floor (the "knee" for the batch), and fit simple power laws in log space. At run time the formulas output these hyperparameters directly, so no per-dataset tuning is required.

At the end of training, two optional finalization steps produce final centroids: a last-epoch mean, which reuses assignments already computed in the final epoch and is guaranteed to lower training MSE, and a slower-moving EMA that generalizes better to unseen vectors. For very large cluster counts, the assignment step can be replaced by an IVF approximate nearest-neighbor search over the centroids.

Why This Matters

Impact on research. The paper challenges the assumption that full-batch Lloyd is the quality ceiling for k-means on static datasets, and it argues that explicitly fitted hyperparameter laws — not tuned constants — are what make mini-batch clustering practical. It also contributes a 10B-scale benchmark dataset, DINO10B, for both clustering and approximate nearest-neighbor search.

Real-world applications (all named in the paper):

• Retrieval and indexing, where k-means centroids define inverted-file buckets and quantizers.
• Compression and quantization of embeddings, including codebook learning inside VQ-VAEs.
• Data curation of very large unlabeled corpora, e.g. partitioning billions of DINO features.
• LLM stacks, where k-means has been used to compress the attention cache and to sparsify attention.
• Image segmentation and general data analysis.

Industry relevance. The claimed 10–45× wall-clock savings and the elimination of hyperparameter sweeps translate directly into lower GPU cost for web-scale clustering pipelines. The method is released as a library link in the Faiss repository plus a project website, and it degrades gracefully — random initialization, no k-means++ setup, no tuning — which is what production pipelines need.

Future Directions

• Low-dimensional and tiny datasets. The scaling laws are explicitly not optimal there (Appendix D.2), and a dimensionality term was tried but dropped as unstable, so extending the laws to these regimes remains open.

• Estimating intrinsic dimensionality. The authors identify unknown local intrinsic dimensionality as the key obstacle to refining the laws, since it is the quantity that actually governs behavior rather than the ambient dimension d.

• Better initialization at low epoch budgets. k-means++ slightly improves results at T = 3 but is sequential and slow; the paper notes it is better to train longer than to pay for it, leaving room for a cheap, parallelizable initialization.

• Approximate assignment as a first-class mode. IVF assignment only pays off above roughly 100M points in these experiments, so tuning the IVF configuration and understanding its effect on convergence at intermediate scales is an open engineering question.

Target Audience

Practitioners and researchers working on large-scale clustering, vector search and quantization — particularly people maintaining or choosing between k-means implementations (Faiss, Flash-KMeans, cuML, SuperKMeans) who care about wall-clock time to a target loss rather than per-iteration cost. It is also relevant to anyone fitting empirical scaling laws for optimizer hyperparameters, and to engineers building embedding pipelines over hundreds of millions to billions of vectors, who will find the formulaic batch/step rules and the DINO10B benchmark most directly useful.

Authors’ abstract

Lloyd's k-means algorithm is a long-running standard for clustering, with applications in data analysis, retrieval, quantization and data curation. In this paper, we introduce an exponential moving average optimization procedure for k-means with mini-batch updates, powered by modern deep learning machinery and new insights. We also establish scaling laws on how decay rate and batch size impact final performance. This gives us a formulaic rule to fix them for multiple problem…