e-means K-Means: Faster Clustering at Lower Loss
See how e-means speeds up k-means clustering, why its mini-batch updates can improve loss, and what its benchmarks mean for production workloads.
e-means: A Faster Way to Run K-Means Clustering
K-means is a familiar tool for grouping vectors, but the standard version comes with a painful cost: each Lloyd iteration does a full scan of the dataset to reassign points and recompute centroids. On large embedding collections, those repeated passes can dominate training time. Mini-batch variants help with the compute per update, but they can also make results less steady - or leave some clusters effectively unused.
e-means, developed at Meta FAIR, tries to hit a better balance. It uses mini-batch updates plus exponential moving averages, automatically selects batch size and step size, revives dead clusters, and then finishes with a finalization stage. In the reported evaluations, it reaches loss levels below Lloyd's while using fewer epochs and less wall-clock time. Still, you'll want to read those numbers with the usual context: the dataset, the hardware, and how the method is implemented all matter.
This article walks through how e-means works, what's different about its update rule, and what the reported speedups might mean for real production workloads like embedding quantization and large codebooks. It also covers when the method is likely a good fit - and what you should check before swapping it into an existing clustering pipeline.
What is e-means and what problem does it solve?
K-means takes a fixed dataset of vectors and an integer k, then looks for k centroids that minimize the average squared distance from each point to its nearest centroid. With points X = {x₁, ..., xₙ} and centroids C = {c₁, ..., cₖ}, the objective is to reduce the mean squared distance between every point to its assigned centroid.
The usual baseline is Lloyd's algorithm. Each iteration assigns every point to its nearest centroid, computes the exact mean of the points assigned to each cluster, and replaces the centroid with that mean. It's simple and often reliable. The cost is simple too: every iteration needs a full pass over the dataset, including nearest-centroid assignments.
That cost becomes a real issue when you're working with millions of vectors or when the codebook has a lot of clusters. An embedding pipeline might also cluster data repeatedly as a model or product changes. A faster clustering step can reduce GPU or CPU time, speed up experimentation, and lower the cost of building a codebook. But runtime improvements only matter if the cluster quality is still good enough.
What e-means changes
e-means is a mini-batch optimization method for the static k-means setting. Instead of using the entire dataset at every update, it processes batches and updates cluster statistics with an exponential moving average (EMA). The method also includes automatic rules for selecting batch size and step size, initialization and cluster-revival strategies, a scheduled decay in its update rate, and a finalization stage.
The project reports that e-means can get a test loss below Lloyd's in 25 epochs, using 1.2 to 3 epochs for the method across a range of datasets and scales. In the reported comparisons, that corresponds to 10-45× less wall-clock time on GPU versus Flash-KMeans and 8-15× less on CPU versus Faiss. These are results from the evaluated configurations, not a blanket guarantee.
The intended workload matters
e-means is aimed at clustering a fixed dataset. That's not the same as a codebook inside a model like a VQ-VAE, where the vector distribution changes during training. EMA-style updates show up in moving-distribution settings already, but e-means is tuned specifically for a static clustering objective.
In practice, this distinction matters: if your distribution keeps shifting, a method designed for one-time clustering may not be the best match. If you're doing one-time clustering, then the reverse is also true. Before comparing methods, it helps to pin down whether your data is static, how often you recluster, how many centroids you need, and whether your main constraint is runtime, loss quality, or both.
How does e-means update centroids differently?
All the update rules here can be framed with two per-cluster statistics: a sum Sj and a count Nj. The centroid is their ratio, cj = Sj / Nj. What changes from method to method is how those statistics are updated as new data comes in.
Lloyd, MacQueen, and EMA updates
Lloyd's algorithm uses the full dataset as its batch. Each iteration recomputes each cluster's count and sum from the current assignments. That gives an exact mean for the points currently assigned to each cluster, but you pay for it with a full data pass every time.
MacQueen's method accumulates counts and sums over time rather than discarding older assignments. The centroid can therefore reflect points assigned under earlier centroid positions. If assignments change as centroids move, older observations can stick around as history.
Staleness-Reduction Mini-Batch K-Means later introduced a hybrid approach, with reported speedups when the number of clusters is relatively small.
For each e-means mini-batch, let sj and nj be the sum and count of points assigned to cluster j in that batch. The method updates the stored statistics using an EMA, mixing the new batch statistics into the existing values. In simplified form:
- Sj ← (1 − α)Sj + αsj
- Nj ← (1 − α)Nj + αnj
- cj ← Sj / Nj
Here, α controls how strongly the current batch affects the estimate. As training goes on, older assignments lose influence, so the centroid tracks recent batches instead of keeping the entire history.
Why use an exponential moving average?
The e-means team reports that EMA updates were more stable than gradient-descent-based centroid updates at large learning rates. One geometric explanation they give is that an EMA update keeps the centroid within the convex hull of the points contributing to the update, while a gradient step can overshoot that region. That's a helpful intuition, but it's not the same thing as a universal proof that EMA always wins.
The update rate also changes over the run. e-means uses a cosine-scheduled decay from an initial α₀ down to 0.1α₀. Early on, updates respond more strongly to incoming batches; later, updates become more conservative. Then a finalization stage computes each cluster's exact mean from the last epoch's assignments, with a reported negligible runtime cost.
None of this is "mini-batch only." Mini-batching is just one ingredient. The EMA choice, the schedule, how it deals with dead clusters, and the final exact-mean step all feed into the behavior reported in the paper. If you treat e-means as just "Lloyd with smaller batches," you miss the parts that make the comparisons meaningful.
Why can e-means outperform Lloyd's algorithm?
A tempting assumption is that more exact updates should produce better clustering. Lloyd's updates do compute exact means for the current assignments, and under the standard k-means objective, its loss decreases monotonically. But monotonic improvement doesn't guarantee you land in the best local solution you could reach.
K-means is non-convex, so the end result depends partly on initialization and the path the centroids take through the optimization landscape. You can make steady progress toward some local minimum and still finish in a bad configuration. In experiments described by the e-means team, repeated runs over 200 seeds on toy datasets showed that e-means was less likely than Lloyd's to end up in certain pathological cases.
Mini-batches add useful noise
A mini-batch estimates assignment statistics from only part of the dataset. That adds noise to the update. In some optimization setups, noise can actually help you move away from a narrow or poor local minimum rather than reinforcing the same bad path every time.
That's a reasonable explanation for the reported results: Lloyd's full pass can lock onto the current configuration, while e-means' partial updates can explore different paths. Of course, the same noise can also hurt. If it's too large, centroids may move erratically and loss may vary more between runs. e-means tries to manage this trade-off with its scaling laws and update schedule rather than forcing users to tune everything by hand.
Lloyd is a special case, not the best comparison
If the batch size B equals the dataset size n and α equals 1, e-means' statistics line up with the full cluster counts and sums. In that setting, a single update is essentially equivalent to a Lloyd iteration, just without the cosine scheduling. So yes - Lloyd can be viewed as an extreme parameter choice inside the broader e-means update framework.
The e-means hyperparameter study also found that this full-batch setting wasn't in the lowest-loss region of its tested batch-size and step-size grid. That helps explain why the method's advantage isn't only about implementation speed. Its update strategy can change the optimization path within a fixed training budget.
Still, lower loss isn't the only definition of success. Your application might care more about downstream retrieval quality, codebook usage, assignment latency, or reproducibility. It's worth validating the metric your product actually cares about, rather than assuming that a lower clustering MSE automatically improves everything downstream.
How does e-means choose batches and handle dead clusters?
Mini-batch methods have their own balancing act: batch size and update rate affect both convergence and runtime. Too small a batch can make updates noisy. Too large a batch can reduce the number of updates you get per unit of compute. Many implementations deal with this by tuning these values through experiments - which costs time and can reduce reproducibility across different dataset sizes.
Scaling laws replace manual tuning
e-means fits scaling laws to choose the initial step size α₀ and batch size B based on problem characteristics like number of points n, number of clusters k, and training budget T. The reported laws were tested on held-out datasets and held-out problem sizes. The study reports low median regret relative to the best settings in its search grid, with R² values of 0.95 for the step-size fit and 0.96 for the default batch-size fit.
The method provides a default batch-size law and a separate quality-oriented law. The quality-oriented choice trades throughput for better clustering quality in the reported experiments. Practically, this gives you a couple of sensible presets instead of a long list of knobs: use the default if throughput matters most, or pick the quality-oriented setting if the loss improvement is worth extra runtime.
Automatic selection doesn't mean every workload has identical constraints. Production systems can have memory limits, device-specific batch behavior, or hard latency budgets. Scaling laws help as a starting point, but you still need to validate under your real deployment constraints.
Reviving dead clusters by splitting active ones
A dead cluster is a centroid that gets too few assignments to matter. In a large codebook, dead clusters waste effective capacity: if you're maintaining thousands of centroids that never get used, that's real lost representational power.
A common fix is to replace a weak centroid with a random point from the current batch. e-means takes a different approach: it revives dead clusters by sampling an active cluster in proportion to its size, then splitting it into two. The goal is to redistribute capacity from a region with many points, rather than inserting an arbitrary point that may not correspond to a useful region of the space.
In the reported ablation on one million vectors with k = 16,384 over five epochs, proportional splitting produced 47 dead clusters on DEEP, versus 2,343 when resampling from the batch and 4,313 with no revival. On SIFT, the counts were 10, 750, and 2,080, respectively. The study also reports lower test MSE and lower cluster imbalance for proportional splitting in those comparisons.
These numbers are specific to the reported datasets and setup, but they point to something important: cluster utilization should be measured alongside loss. You can get a decent average distance while still wasting a meaningful portion of your codebook.
What do the benchmarks mean for production use?
The reported performance is interesting because it claims improvements in both training speed and loss - something that's often hard to get at the same time in optimization. Across evaluated datasets and scales, e-means reached a test loss below Lloyd's 25-epoch result in 1.2 to 3 epochs. End-to-end comparisons report 10-45× lower wall-clock time on GPU against Flash-KMeans and 8-15× on CPU against Faiss.
The source also describes validation across 24 embedding and synthetic datasets and a broader comparison of 234 configurations across 16 datasets with 100,000 to 100 million points. That's stronger than a single toy example, but it still doesn't remove the need to check how closely your workload matches what was tested.
Separate algorithm gains from kernel gains
Fewer epochs don't automatically translate into lower total runtime if each epoch is expensive. e-means addresses part of this by using specialized, fused GPU kernels for nearest-centroid assignment. So the reported end-to-end speedups reflect a mix of algorithmic changes - fewer or cheaper updates - and implementation choices that turn those savings into wall-clock improvements.
If you're evaluating e-means, compare full pipelines on the hardware you'll actually run. Include data loading, memory movement, initialization, assignment, updates, finalization, and any host-device transfers. A benchmark that only measures the update kernel won't tell you much about production cost.
A practical adoption checklist
Before switching a clustering job, run a controlled comparison using the same data, preprocessing, initialization budget, hardware, and stopping criteria. Track:
- Test loss: Does e-means improve the objective on held-out or evaluation data?
- End-to-end runtime: Include setup and finalization, not just steady-state updates.
- Codebook utilization: Count clusters with negligible assignments and measure imbalance.
- Run-to-run variation: If downstream use is sensitive to consistency, compare several seeds.
- Downstream impact: Check the retrieval, compression, or model-quality metric the clustering step is meant to support.
- Operational fit: Make sure the implementation works with your memory limits, device behavior, and reproducibility needs.
e-means is most relevant when full-dataset Lloyd iterations are a major cost and when your workload has many vectors or clusters. It may matter less for smaller datasets where clustering is already cheap, or when the overhead of integrating specialized kernels offsets the time you save. The right comparison isn't just "new algorithm vs. old algorithm." It's system cost versus the quality your product actually needs.
Conclusion
e-means is a modern alternative to Lloyd's algorithm for static k-means workloads. Its core idea is to combine mini-batch updates and EMA statistics with automated scaling laws, scheduled updates, cluster revival, and a final exact-mean stage. In the reported evaluations, that combination achieved lower loss in fewer epochs and substantial end-to-end speedups compared with established implementations.
The strongest case is when clustering is a meaningful part of your compute budget and when codebook capacity matters. Treat published speedups as a reason to run your own tests, not as a guarantee. Benchmark on your own vectors, track both loss and cluster utilization, and measure the downstream outcome that makes the clustering step worth doing.