e-means: outperforming k-means in 1/10th the epochs

Lloyd (full batch)

e-means (batch and step from the scaling laws)

Training MSE

epochs: 0.0

Gray dots are the data points; the colored dots are the centroids. Toy distribution adapted from Huy Vo's work on hierarchical k-means for data curation (10× sample, 90,000 points, 4,500 drawn for the plot), K = 300. Lloyd moves once per epoch (the staircase); e-means moves roughly 176 times per epoch, at batch 512 and peak step 0.083 from the two scaling laws below. By epoch 2 it is already below Lloyd's epoch-25 loss; it finishes after 4 epochs about 10% lower. True over 40 random restarts (wins all 40). MSE is the training error over all 90,000 points; this toy has no held-out set.

It’s just faster. By a lot.

Across the datasets tested (real embedding collections, production search indices, synthetic toy data, spanning a few thousand points to ten billion), e-means reaches a loss below Lloyd’s full 25-epoch convergence using, on average, 1.2 to 3 epochs. This is not a trade of quality for speed: the loss is lower, not merely close, and it arrives sooner. Measured end to end, that is a 10–45× wall-clock reduction on GPU against a strong k-means implementation, and 8–15× on CPU against Faiss. Speed and quality usually trade against each other in k-means, so in the ordinary sense this shouldn’t be possible: close to a free lunch, and the scaling laws below are why it holds.

Training (top) and held-out test (bottom) MSE, divided by Lloyd (Faiss) after 25 iterations from random vectors, mean over 10 seeds, one H100. e-means and Lloyd start from random vectors; Staleness-Reduction mini-batch k-means starts from greedy k-means++ and is tuned on validation data.

Test error against wall-clock time on one H100, 10 million vectors, K = 65,536. 25 iterations of a plain PyTorch Lloyd (data on the GPU, bf16) take 142 to 155 seconds, against 14 to 22 seconds for Flash-KMeans. Numbers on the curves mark epochs for e-means, iterations for Lloyd. Flash-KMeans, cuML and Faiss-GPU are drawn at one shared Lloyd error curve, each at its own measured time.

What e-means actually is

e-means is a mini-batch update for k-means, but mini-batch k-means already existed. What it changes over earlier mini-batch methods is specific:

• A reliable dead-cluster revival strategy. We tried several and found the one used here, faiss-style splitting of a populous cluster into a dead one, the most effective.
• Two scaling laws that set the batch size and step size automatically, removing the per-dataset tuning earlier mini-batch methods needed.
• Most importantly, the update rule itself. Earlier mini-batch k-means methods (Sculley’s, Staleness-Reduction, and others) update each centroid to the plain mean of whatever points land in it this batch. We found an exponential-moving-average update, which blends that batch mean into the centroid rather than replacing it outright, measurably more efficient.
• We also tested updating centroids with standard neural-network optimizers, Adam and SGD, instead of an EMA, and found both consistently less efficient. An EMA update stays, by construction, a weighted mean of the points assigned to it, which is more stable than a gradient step: it cannot overshoot past the data the way a gradient update can.

The EMA update itself is not new on its own; it is also how VQ-VAE (van den Oord et al., 2017) updates its codebook, though there it runs inside a neural network on a representation that keeps changing, rather than on a fixed, arbitrary dataset.

Concretely: Lloyd’s algorithm assigns every point to its nearest centroid, then sets each centroid to the exact mean of the points assigned to it, one full pass over the data per move. e-means keeps that same mean, but as a running quotient of two exponential moving averages instead of a single exact division: one tracking how many points cluster j has been getting, one tracking their sum.

\[ N_j \leftarrow (1-\alpha)\,N_j + \alpha\, n_j, \qquad\quad S_j \leftarrow (1-\alpha)\,S_j + \alpha\, s_j \]

where n_j and s_j are the count and the sum of this mini-batch’s points landing in cluster j (zero for both if none do). The centroid itself is read off as the ratio:

\[ c_j = S_j / N_j \]

The idea behind one update step. A mini-batch nudges the centroid toward its own mean rather than replacing it outright: the exact amount depends on how much accumulated weight N_j already carries, not a fixed fraction.

Set the batch to the full dataset and α = 1 and N_j, S_j become exactly the count and sum of the whole cluster, so their ratio is exactly one Lloyd iteration. e-means is the same update at a small batch and a small, annealed α instead: Lloyd is one extreme setting of it, not a different algorithm.

Using k-means? Try e-means instead.

e-means is on its way into faiss/contrib (in progress). Same call signature as Lloyd: points n, clusters k, an epoch budget T. Everything below runs underneath, calibrated automatically.

from faiss.contrib import emeans centroids = emeans.fit(x, k=k, epochs=T)

Repo: [LINK, coming soon].

Why the noisier update wins, not just runs faster

Lloyd can only go downhill, so it has no way out of a bad arrangement once it’s in one. A mini-batch update is noisy by construction, and that noise is a small perturbation at every step, the same reason stochastic gradient descent reaches better optima than full-batch gradient descent, not merely a cheaper way to reach the same one.

Fast kernels alone don’t fix this

Flash-KMeans, released by Berkeley in 2026, is a much faster k-means library: better-engineered GPU code, a large real speedup over a naive PyTorch loop and earlier GPU implementations. But it’s still full-batch Lloyd underneath, one centroid update per pass over the entire dataset, just computed quickly. That full-batch ceiling is a separate problem from how fast the code runs, and there’s a lot of speed left on the table once you fix it: the 10–45× reported above is measured against Flash-KMeans itself, already fast, not against a naive baseline it would be easy to beat.

The deep learning machinery, precisely

The step size α is annealed to zero on a cosine schedule rather than held constant, the same curve a learning-rate schedule follows, for the same reason: a step that never shrinks never lets the estimate settle. Dead-cluster revival, a faiss-style rule that replaces a centroid a mini-batch keeps missing by splitting a populous one, is necessary to keep the codebook from collapsing, but not sufficient on its own: without the annealed step and the scaling laws below, reviving the right clusters at the right rate is exactly the per-dataset tuning problem this recipe exists to remove. The last epoch closes the run on its own assignment, each centroid set to the exact mean of the points it received, free, since that assignment was already computed.

What makes this transfer across datasets without tuning is that the batch size B and the peak step α₀ are not chosen by hand. Given only (n, k, T), two power laws return both, fit once on one dataset (deep1b) and validated on datasets, scales and dimensions the fit never saw:

\[ B = c_B \, k^{b_k} \, n^{b_n} \, T^{b_T} \]

\[ \alpha_0 = \min\!\Big(c_\alpha \, k^{a_k} \, B^{a_B} \, n^{a_n} \, T^{a_T},\; 1\Big) \]

law	constant	exponent of				regret, other datasets	regret, held-out size	R²
		k	B	n	T
step α0	0.824	−0.46	0.90	−0.45	−0.11	0.07%	0.07%	0.95
batch B, knee at 0.2%	0.035	0.13	n/a	0.88	0.81	0.19%	0.20%	0.96
batch B, knee / 8	0.0044	0.13	n/a	0.88	0.81	0.07%	0.06%	n/a

Regret = (MSE with the law − best MSE over the grid) / best MSE, median over configurations. Other datasets: SIFT, SPACEV, arXiv and AG-News, none seen by the fit. Held-out size: one value of n left out of the fit and predicted. Above two million points, the batch knee divided by 8 is used: same throughput, lower regret.

The two scaling laws, fit on deep1b. Left (T = 4): training error against step size for several batch sizes, 3 seeds; stars mark each batch's best step, what the step law fits. Right: each point at its own best step, 200 seeds; diamonds mark 0.2% above the floor, the batch knee the batch law fits.

The same recipe, from a few thousand points to ten billion

The headline evidence is breadth, not a single favorable number. Across roughly thirty datasets (the 11-dataset VIBE embedding benchmark, production search indices, synthetic toy data, real-world corpora at every scale we could source), the same two scaling laws are applied with zero per-dataset tuning, every time:

• VIBE (11 datasets, 0.2M–1.5M vectors): 8.5–21.2× GPU and 9.0–15.0× CPU speedup to reach Lloyd@25’s loss, in roughly 2 epochs.
• DINO10B (10 billion image embeddings, 16 million clusters): below Lloyd’s loss in 2.5 epochs versus 25.
• A 13-billion-vector production Instagram index, half a million clusters: matches a custom, hand-built DINO-team codebook in 7 ordinary epochs.
• Verified recently at the small end too: 1,000 and 10,000 points, K = 2 and 16. Same recipe, same result, 2 and 4 epochs respectively. This is not a large-scale-only effect.

Epochs needed to reach Lloyd@25's test loss on DINO10B, n = 1,000 to 10,000,000,000, K the power of two nearest n/500. Epoch count only, no wall-clock speedup shown here. Same recipe at every point, no re-tuning: the epoch count needed stays in a 1.4–4 band across seven orders of magnitude.

Test error against wall-clock time on one H100, relative to Lloyd after 25 iterations, across the 11-dataset VIBE embedding benchmark (Jääsaari et al. 2025), d = 200–3,072.

In many cases, it beats k-means++ too

A smaller finding, worth stating plainly as secondary rather than the headline: across 40 random restarts per method on several datasets, e-means’s worst restart beats Lloyd’s best restart, including when Lloyd is initialized with greedy k-means++, the strongest initialization in common use. This happens because e-means keeps refining the codebook long after a converged Lloyd run has locked in; it is a consequence of the update rule, not something tuned for.

Final MSE across 40 restarts (seeds 10–49). Top row: held-out test MSE; bottom: training MSE; dashed lines mark medians. In every panel the distributions do not overlap.

Try it yourself

A CPU-only, dependency-light notebook runs e-means against Lloyd k-means on any distribution, cluster count, and dataset size you give it. [Colab notebook, coming soon].

---

This post accompanies our ICLR 2027 submission (S. Bessa, P.-E. Mazare, G. Szilvasy, H. Touvron, D. Picard, M. Douze, H. Jégou), currently under anonymous review. A public citation and the faiss/contrib repo link will be added here on release, expected early November 2026.