e-means: beating k-means in one-tenth the epochs

k-means is one of the few problems in machine learning where full-batch optimisation is still the default. Lloyd's algorithm, which goes back to a 1957 Bell Labs technical note, sweeps the entire dataset every iteration, assigns each point to its nearest centroid, and recomputes the centroids. At a billion points, that "recount everything" step defines the algorithm's economics.

Meta FAIR's e-means paper changes the premise. It updates each cluster's sum and count with an exponential moving average on every mini-batch, then sets batch size and step size from the problem's dimensions. The paper reports reaching Lloyd's 25-epoch loss in 1.4–2.6 epochs. The project site gives the range as 1.2–3. This article is a review of the paper and public code. CodeBridge did not reproduce the clustering experiments.

The problem is fewer passes, not better centroids

The k-means objective is to choose k centroids C that minimise the mean squared distance from a fixed dataset X. Lloyd assigns each point to its nearest centroid and replaces the centroid with the exact mean of its assigned points. The loss decreases monotonically, but each update reads the whole dataset once. Typical runs use 20 to 50 iterations.

Mini-batch k-means has existed for decades, but most variants follow MacQueen's 1967 update. MacQueen keeps accumulating the batch sums and counts. The centroid becomes the uniform mean of every point assigned to that cluster since training began, including assignments made under earlier, possibly bad, centroids. The more data you see, the heavier that old history becomes.

e-means starts by admitting the difference. Even on a static dataset, trusting recent assignments more can converge faster. The observation is that a moving average, long used where the distribution moves, can be retuned for static k-means.

The update rule: forget the old sums

For each cluster j, keep a sum S_j and a count N_j, with the centroid at c_j = S_j / N_j. The three algorithms differ only in how they update S_j and N_j. Let s_j and n_j be the sum and count of the current batch's points assigned to cluster j.

Lloyd      S_j ← s_j                    N_j ← n_j
MacQueen   S_j ← S_j + s_j              N_j ← N_j + n_j
e-means    S_j ← (1−α)S_j + α·s_j       N_j ← (1−α)N_j + α·n_j

In e-means, α is a forgetting rate. Older assignments fade by α, and the centroid tracks recent assignments rather than the whole history. If α = 1 and the batch is the full dataset, e-means is exactly Lloyd. Lloyd is a special case of e-means. That framing matters: the new algorithm does not replace the old one so much as place it at one extreme of a hyperparameter.

Making α small does not automatically make things faster. Small batches add gradient noise and can oscillate; large batches reduce the number of updates per epoch. e-means works because of three other parts.

Part 1: scaling laws set the batch and step

e-means claims that users do not need to choose batch size B or initial step α₀. Instead, two scaling laws compute them from the problem's characteristics (n, k, T), where n is the number of points, k the number of clusters and T the number of epochs.

B    = c_B · k^{b_k} · n^{b_n} · T^{b_T}
α₀   = min(c_α · k^{a_k} · B^{a_B} · n^{a_n} · T^{a_T}, 1)

The paper offers two batch laws. The default law prioritises throughput; the quality law trades throughput for lower loss at the same budget. The project site reports R² of 0.95 for the step law and 0.96 for the default batch law, with regret of roughly 0.07–0.2% on held-out datasets and held-out sizes. "No hyperparameter sweep" rests on that fit.

Cosine scheduling is layered on top. α anneals from α₀ to 0.1α₀ by the end. A finalization stage then ends the run on each cluster's exact last-epoch mean, reducing loss for a negligible runtime cost.

Part 2: revive dead clusters by splitting a large one

The chronic failure of mini-batch k-means is the dead cluster. When a centroid receives no points, the effective codebook shrinks. At 16,384 clusters the loss of capacity matters.

The usual fix is to replace a dead centroid with a random point from the batch, as EnCodec does. e-means instead samples an active cluster with probability proportional to its size and splits it in two, a strategy adapted from Faiss k-means.

The ablation shows the effect. On one million vectors, k = 16,384, five epochs and five seeds, the DEEP dataset produces:

On SIFT the proportional split leaves 10 dead clusters (0.1%), resampling 750 (4.6%) and no revival 2,080 (12.7%). Fewer dead clusters and lower loss move in the same direction.

Part 3: mini-batches escape local minima

The more interesting observation is that e-means does not just run faster. It reaches a lower loss. The authors trained toy datasets over 200 seeds and counted how often each method stayed stuck in a pathological configuration.

• Two blobs of very different sizes, k = 8: Lloyd stays stuck on 26% of starts; e-means on 3.5%.
• Nine blobs in a 3×3 grid, k = 10: Lloyd is more than 1% worse on 43% of starts; e-means on 4%.
• Five blobs on a line, k = 5: Lloyd stays stuck on 46% of starts; e-means on 0%.

The explanation is that Lloyd's loss always decreases, so once it enters a local minimum it has no reason to leave. Mini-batches add noise, explore more widely and create chances to escape a bad arrangement. That is the basis for claiming both speed and quality in a problem where the two usually trade off.

The evidence: 24 datasets and DINO10B

The paper says it validated e-means on 24 embedding and synthetic datasets and configurations. The project site publishes a results map covering 16 datasets, 100K to 100M points and 234 configurations. The headline numbers:

• On the 11 VIBE embedding datasets, e-means reaches Lloyd's 25-epoch test error in 1.6–2.4 epochs. The datasets cover text, code, image and multimodal embeddings, dimensions 200–3,072, with k = 8,192 or 16,384.
• On 10 million vectors, k = 65,536, one H100, 25 epochs of plain PyTorch Lloyd take 142–155 seconds and Flash-KMeans takes 14–22 seconds. e-means is faster still.
• e-means partitions DINO10B, 10 billion DINO embeddings, into 16 million clusters in 2.5 epochs, with a loss below Lloyd's at 25 epochs.

These are the authors' numbers. CodeBridge did not install the Faiss implementation and rerun the datasets. "10–45×" is GPU wall-clock against Flash-KMeans; "8–15×" is CPU against Faiss Lloyd. Read the baseline and hardware with the multiplier.

Where to use it first

e-means fits pipelines that cluster embeddings at scale: a coarse quantiser for a search index, deduplication in data curation, candidate clustering in recommendation, topic partitioning of a RAG document set. The workload is static, the cluster count is large, and the cost of a full run is high, so fewer epochs turns directly into lower cost.

The interface is small.

from faiss.contrib.e_means import Emeans

em = Emeans(d, k, nepochs=T)          # T defaults to 5
centroids, labels = em.train(x)       # x: PyTorch tensor (n, d)

Use NumpyEmeans for NumPy. You do not have to throw away an existing Lloyd pipeline. Run e-means on the same data, compare loss against Lloyd, and check the cluster-size distribution and downstream metric before switching.

Limits: static data and the average trap

First, e-means assumes a static dataset. The paper explicitly distinguishes it from codebook learning in VQ-VAEs, where the distribution keeps moving. Moving averages came from that setting, but e-means' scaling laws and finalization are fitted to a fixed X. Online use needs separate validation.

Second, lower loss does not always mean better downstream quality. In search or classification, cluster balance, interpretability or alignment with a known label may matter more. The paper's metrics are test MSE and cluster-size imbalance. Re-measure with your own task metric.

Third, "no hyperparameters" is true inside the range where the scaling laws were fitted. Move far outside the paper's n, k and T, and the laws lose predictive power. GPU kernels and CPU implementations also gain differently, so measure wall-clock in the deployment environment.

Insight: reread an old algorithm with modern optimisation

The real message of e-means is not that k-means is slow. It is that k-means was the one problem that kept full-batch optimisation out of habit rather than necessity. Moving averages, scaling laws, cosine scheduling, dead-cluster revival and fused kernels all existed in deep learning. Reassembling them for a 1957 problem cut the epochs by an order of magnitude.

If you are considering adoption, the question is not "is e-means faster than k-means?" It is "is my data static, is my cluster count large, and which metric besides loss must I protect?" Answer those, then run both algorithms side by side on the same data. The paper's 1.4–2.6 epochs is a hypothesis for that experiment, not a guarantee for your pipeline.

To go deeper on embedding retrieval architecture, the RAG systems course continues with retrieval and generation design. For model-routing cost, read Muse Ember. For the memory boundary of a local model, read Nota AI's 1.58-bit 30B.