E-means: Outperforming k-means in one-tenth the epochs
e-means.github.ioLloyd is not the fastest algorithm for k-means for very long time. I am curious how e-means compares with more advanced algorithms such as BBD-tree based implementation. SMILE has a k-means implementation with both BBD-tree and llyod algorithm. BBD-tree approach is often 10x faster. Wonder if e-means is better than BBD-tree.
See https://www.aihalo.dev/api/java/smile/clustering/KMeans and https://www.aihalo.dev/api/java/smile/clustering/BBDTree
Maybe we combine the ideas of e-means and BBD-tree?
BBD-tree is only for low dimensions because it suffers from the curse of dimensionality.
Note that BBD-tree is orthogonal to e-means, they can combined very straightforwardly just like it is combined with k-means. The 10x speedup should remain compared to BBD-tree k-means.
So it's not that e-means is better than BBD-tree, it's that you can combine them to make e-means even better, but only at low dimensions and on CPU.
I wonder if this could be used anywhere as an alternative to Lloyd's algorithm. I have a hobby project which uses Lloyd's to "relax" an icosohedral equal area mesh on the surface of a sphere [1]. I'm going to try to swap it out with an adaptation of this and see if it's much better or if I can afford more iterations with this method.
Edit: tried it, and it’s definitely not better for this use case. e-means wins by needing fewer passes over a huge dataset, but relaxing a mesh on a sphere has no dataset: each Voronoi cell’s centroid can be computed directly from its geometry. Mini-batches just swap that for a noisy sampled estimate of the same thing. It has the same objective and the same optimum, so at best it matches Lloyd. In practice it was much slower and left the grid noisy, and with small sample budgets it even tangled the mesh.
[1] https://webgcm.app/explain/grid#nudging-the-cells-into-shape
Hi, I'm Swann Bessa, the author of e-means. Thanks for trying e-means. If I may I'd like to come back on a few of the points you made.
"e-means wins by needing fewer passes over a huge dataset": the improvements of e-means are shown on datasets from 10k up to 10B points (you can try e-means against Lloyd at low problem sizes at https://e-means.github.io/playground.html) so I would not say huge datasets only. It depends on what you call huge obviously, and that's relative.
"It has the same objective and the same optimum, so at best it matches Lloyd." I strongly disagree with that. e-means has the same objective and the optimal clustering is set by the problem, but Lloyd doesn't find this optimal clustering in the vast majority of cases. Lloyd gets stuck in a local minimum, and what we find is that e-means reliably finds better clusterings than Lloyd in terms of MSE, but also other clustering metrics.
Regarding your usecase : I'd be curious to know what settings you used for e-means and a reproduction code. Don't hesitate to send me your code at swann@meta.com, I'm down to have a check because we're confident about e-means' performance against Lloyd.
Is E-means finding the true global minimum? Or is it just finding a better solution more consistently?
The answer is that e-means consistently finds a better solution than Lloyd. It's impossible to know if it's the global minimum or not. Finding the true global optimum is NP-hard, so no practical algorithm can guarantee that.
> if it's the global minimum or not
Not impossible, just very expensive, because k-means is NP-hard, and finding the strictly optimal solution may require exponential time (or worse).
Yes, right. That's more rigorous to say it like that. Note that at the scales of the paper it is in practice impossible with today's computational means.
Since it's basically a linear interpolation between Lloyd's algorithm and not moving any points it will be at least as good but gains might be limited without careful tuning.
Edit: in fact you might end up going faster by moving the other way, set alpha = -0.1 and see what happens!
Careful tuning is provided by the scaling laws, so you don't need to set the hyperparameters yourself. I am quite skeptical about setting alpha=-0.1 tbh, but I'm curious to see what happens too lol.
Your website looks incredibly interesting! Very exited to read it when I have time. When you mentioned icosahedral meshes, I immediately thought it would be related to weather/climate modeling :D.
On the topic of these meshes, I'm currently exploring if we can come up with better meshing strategies in data-driven weather models (https://tcumulus.dev/blog/thesis/).
I adapted the method to my IVF index (the code-daemon project) and achieved a 146x speedup on the GPU.
That's impressive!
Thank you for trying e-means, let me know if you have any questions !
i would like to see comparisons to k-means++. k-means is rarely the bottleneck in most modern algorithms, instead it's the multiple runs or seeding of kmeans to determine a stable 'good' solution that takes the lion's share of processing time. If this avoids the km++ seeding process while achieving similarly accuracy and stability, then this is would be a great practical improvement.
Thanks for this. I tried experimenting with e-means in retrieval engine im building, for that I had to port the CPU version from faiss to Rust
https://crates.io/crates/emeans
https://pypi.org/project/emeans
We use k-means to build the cells of our vector index, which today runs 5 Lloyd iterations on a sample of the data. So I ran a quick test on 500k LAION embeddings (768 dimensions, 4,096 cells) against our current code:
- Current setup of 5 iterations: about 4x faster, with slightly better recall at every probe count (for example 0.9518 vs 0.9473 at 32 probes) and fewer rows scanned per query
- Cells came out much more even. The biggest cell went from 6.6x the average size to 4.75x, and there were no empty cells (we had 7)
- With only 2 iterations it was about 10x faster and recall stayed close to what we have today
These numbers are from my mac. I'm rerunning on a Linux box with AVX-512 now. If it holds up, we'll try it in our index build
Some notes here: https://github.com/infino-ai/infino/issues/1279
Hi, I’m the author of the paper.
Thank you for trying e-means and detailing your results. This kind of results of e-means at 2 epochs compared with Lloyd is the typical kind of acceleration we want people to know about.
If you have any questions, I’ll be happy to answer.
As far as I can tell this isn't particularly new, people have been doing minibatch k-means for a while, including with an EWA update. I assume the contribution here is that they tried to measure how much better it is and provided code? I'm not sure I agree with their claims though, because they didn't really measure the things that matter–sure, it's better than Lloyd's algorithm, but that's a very low bar.
If you are claiming that your algorithm is "highly parallelizable" and comes with "wall-clock speedups", then the actual metrics that matter here is not number of steps or even basic wall clock time as your size remains fixed, but how your algorithm scales as more data is present or parallelism is required. This algorithm splits up the minibatch across workers and has them compute local updates, then allreduces it. As the amount of data goes up, the comms become much more expensive (since you have to use a worse transport layer), which means that you're basically locked in a desperate attempt to do more local work before you send your partial update to everyone else. If you attempt this (not explored in the paper) you will find that the improvement in clustering by doing this extra work has diminishing returns–you only reduce the number of updates by a few steps (which is what drives your wall clock time). At scale you end up being heavily comms limited and not super parallel at all.
Hi saagarjha, I'm Swann Bessa, the author of the paper. Thank you for your message on the work. If I may I'd like to address a few points.
"People have been doing minibatch k-means for a while, including with an EMA update" : EMA updates are indeed used in clustering, but to our knowledge there is no paper that uses pure EMA updates for k-means clustering on a static dataset. EMA updates are typically used for moving distributions because of their ability to forget the past (as in VQ-VAEs), but their advantage for static datasets had not been empirically demonstrated. If you take a look at appendix E of the paper, we compare EMA to other gradient descent methods and show that EMA is much more stable and has better quality, especially at large batches. Our method also integrates EMA updates in an optimization framework that reliably outperforms Lloyd and previous mini-batch methods.
"the contribution here is that they tried to measure how much better it is": The contribution here is to integrate EMA updates in an optimization framework to provide a method as robust and performant as we could. If you look at previous mini-batch k-means papers, the speedups and the performance is nowhere nearly as good as e-means, which beats Lloyd at 25 epochs in 2 to 3 epochs. In fact, mini-batch k-means is known in the literature as a fast method but with lesser convergence than Lloyd, and we compare to them in Figure 3. The e-means framework allows to be both faster and have better convergence than Lloyd, without the user having to tune the batch or the step.
"How your algorithm scales as more data is present": In the first figure of the paper, we show that e-means scales up to 10 billion points and 16 million clusters.
On communication: workers only exchange per-update statistics (K×d numbers with replicated centroids). You argue that bigger batches have diminishing returns, so that extra local work doesn't cut the number of updates. But that's precisely what our batch law addresses: it returns the largest batch whose loss stays within ~0.2% of the best loss at a fixed number of epochs. In other words, we carefully scale the batch size to avoid that our speedups are driven down by the diminishing returns you mention.
The multi-node regime is in the paper: the 10B run used 32 machines, so synchronization went over the inter-node network, and the bottleneck was still the distance computation, not communication. That's also why an IVF variant, which cuts distance computations, does it in ~23 GPU-hours.
> to our knowledge there is no paper that uses pure EMA updates for k-means clustering on a static dataset
So I'll be honest the reason I said this is that I know there are people who are using it and I just assumed it was common :P I will ask if any of them are relying on published work or what for where they got this from.
> In the first figure of the paper, we show that e-means scales up to 10 billion points and 16 million clusters.
Sorry, I should have clarified. I actually think this is what is saving you: with that many clusters, your arithmetic intensity is inherently really high, so you are truly super compute limited. In this case introducing a bit of comms is cheap relative to the bazillion matmuls you're doing for distance calculation. However if you have fewer centroids then the amount of work you do is much lower between updates. As an extreme example, if you're doing hierarchical k-means you might have as few as 2 centroids, which means you barely do any work, and adding any sort of comms just kills you. Of course you are doing very little math but even if you're just looking at memory bandwidth at that point (because your intensity is like 5 FLOPs/byte or something) it will be done too fast.
(I don't actually know how to solve this, which is why I was super excited to see someone else working on this. I agree that your thing works for more centroids, in fact I know it works because I just assumed it was what you were supposed to do, I just couldn't make it work in this situation.)
Thanks for your kind responses.
I'd be very interested if they have a reference. If you feel like it, you can always send me an e-mail at swann@meta.com with the references. I'll add them to my benchmarks.
The fewer centroids case is super interesting. Because we found in our experiments that basically if your n/k ratio (n is the dataset size, k the number of clusters) is high, running the clustering as is on your n points is a waste of resource. You should subsample your dataset such that there is about 250 points per cluster and you will get very similar loss at a fraction of the compute. This is why faiss k-means, by default, subsamples datasets to 256 points per cluster if n/k > 256. So I fully agree with your computational remarks, but I think it's a case that should never happen in practice because subsampling your dataset is the solution in those cases.
Yeah most of these tricks are already well established. And their own ablations show that combining them doesn’t bring any special synergy - it’s just putting together a lot of common techniques. The baselines also are tough to believe fully. Comparing against faiss is fair since many authors are faiss contributors so they know how to run it. But they don’t and should include a comparison to a well tuned mini batch kmeans like the one in scikit.
Hi oofbey, If you look at Figure 3 of our paper, the baseline "Sculley" is precisely the well tuned mini batch kmeans of scikit-learn. We swept on batch sizes and threshold for dead cluster reassignment (Batch size: 256, 512, 1024, 2048, 4096, 65536, 262144, 524288. Dead-cluster threshold : 0, 0.001, 0.01, 0.05, 0.1, 0.5). The reason e-means outperforms previous mini-batch methods is because it is a robust optimization framework powered by scaling laws, specializing techniques inspired from the deep learning literature for the k-means problem. Meanwhile mini-batch k-means methods rely on the MacQueen update rule introduced in 1967, which is known to be highly suboptimal and converge to a lesser solution than Lloyd (See Darken and Moody, 1990, Note on Learning Rate Schedules for Stochastic Optimization).
This looks awesome. Obvious in hindsight! The viz really explains it well.
Hope it gets added to scipy - it’s much more available in the envs many use than facebooks own faiss library.
The automatic hyperparameter laws don't seem intuitive to me yet. Do you have any intuition for it?
Nor me. This is one of the better things in ml I've seen in a while. Creative.
Thanks for the feedback, thatsabadlook. This is my first PhD work so I’m really happy to hear that.
If you can continue to look at known problems with fresh perspectives like this you have a very bright future ahead of you. But also if this is your biggest thing, you did good. Most grad students don't land on things this practical and useful.
These kinds of algorithms tend to have equivalent descriptions as expectation maximisation or simulated annealing. I wonder if it's possible to understand their algorithm in that sense.
Though I suppose it may be as simple as expectation maximization with a friction term, since Sj = sj - a (sj - Sj).
Importantly they note that this update rule keeps the centroid within the convex hull, which is pretty important.
Hi shiandow, I'm Swann Bessa, the author of the paper.
The fact that the update rule keeps the centroid within the convex hull is indeed very important. If you look at Appendix E, the gradient descent methods don't have this property and are less stable than EMA updates.
Regarding simulated annealing, I agree with the parallel. You can see mini-batches as introducing noise in the beginning of the optimization to espace local minima and the step size decreasing through the training acts like a temperature in simulated annealing. Simulated annealing was already used for k-means with theoretical proof that you can reach the global minimum in enough epochs, but in practice, those methods are really slow. A paper that comes to mind is "A simulated annealing algorithm for the clustering problem" (https://www.sciencedirect.com/science/article/pii/0031320391...).
The parallel with expectation maximization is also coherent, though less straightforward to me. K-means can be viewed as a hard-assignment, special case of the Expectation-Maximization algorithm applied to Gaussian Mixture Models (GMMs).
If you have any other questions, I'll be happy to answer.
Thanks for the extra detail!
I agree that simulated annealing is perhaps the closest. The missing link, I suppose, is knowing what distribution you get by running the algorithm for a fixed temperature (alpha). But the method works even if you don't know all that really matters is that it converges to the distribution you want, which is neat! I might use that trick sometime.
The reason simulated annealing sprang to mind is that I've done something similar once where I tuned the temperature curve for a simulated annealing algorithm, but I couldn't find a general enough rule. So seeing someone do something similar is really neat!
There are some theoretical guarantees to be had if the update step can be recast into a MCMC update, or equivalent. And I think there's an expression for the optimal temperature at any point in time (same way you found it for alpha). Any idea if there is a useful description for the steady state distribution at a given alpha?
Played around with changing our hierarchical balanced incremental clustering to this https://github.com/antflydb/antfly/pull/1607 bulk builds of existing indexes that need to be completely reindexed (like a shard split) were 2.5x faster, overall insertion times didn’t change that much because overall clustering incrementally as data is indexed is a small portion of the cpu time.
The one remaining hyperparameter is of course k. I vaguely remember there was some rule of thumb for picking k like the "elbow of a curve", but what is the state of the art now for picking k?
The 'Elbow method' [1], yeah. Sweep a bunch of different k values and pick the one where your target metric starts having diminishing returns. There are a lot of heuristics you could use.
There's still not one particularly great way to pick k for this type of clustering. There kind of can't be as clustering is sort of inherently subjective.
Other methods like HDBSCAN can find the number of clusters themselves. There's still a 'one key parameter' you have to pick yourself, but it can be more intuitive and less sensitive for some datasets, when you don't know how many clusters there are / how many you want.
Sometimes, picking k exactly is useful for practical reasons too. Like fixing the size of your codebook. Or partitioning a dataset into k equal-ish pieces, say to spread across a pool of servers of a known/fixed size.
Yeah, I guess you can't really avoid picking k without hierarchical clustering.
Yeah, it's not a well-defined problem, so there's no one true solution. Then only question left is "well what do you want?". Always have to pick something.
You could use X-means instead (https://www.cs.cmu.edu/~dpelleg/download/xmeans.pdf). The idea is to use an information criterion that trades the number of parameters against the reduction in within-cluster variance. The approach essentially splits a cluster when the BIC of the split is larger than without. I’d use this to select k and still run normal k-means though.
Why would you use run normal k-means if e-means gives lower loss and reliably goes much faster ? If you don't believe it you can try e-means at https://e-means.github.io/playground.html
There isn’t a well defined correct or ideal answer. It depends on what you’re doing with the clusters. Sometimes the raw data are too numerous to properly deal with and you just want a smaller set of points to reason about. In that case pick the number of points that are convenient to reason about. Sometimes you have a fixed budget for the cluster representation and that defines k. Sometimes just pick sqrt(N) so the average cluster size is k.
This is a great example of why p=np is sometimes the wrong question. For this np hard problem approximate solutions are good enough, and if you get them 45x faster that's even better!
The best class of algorithms in this category are those where you know for certain that one answer is impossible and the other is highly probable. Like Bloom Filters, and if you squint, consistent hashing. I wish we had more algorithms of this sort and books about them. If I were a professor I’d be tempted to propose teaching a class.
When I was just starting out, client-server was still a bit of a specialty, mostly reserved for gaming and large institutions. Most of my education was in single task computing. If the university hadn’t offered a class in distributed computing I wouldn’t have gotten much in that direction except in independent study.
In the time since then, we seem to have focused much less on seconds per task and more on tasks per second, which results in more concern about statistical spreads, and queuing mechanisms to manage the variability between a hit and a miss. It would be better if this sort of computation were more organized into a discipline.
My personal favourite is compressive sensing, sparse signals beat the Nyquist sampling limit, by what seems like an unbounded amount 1000x, 1000000x, etc, because all real world stuff we care about lives on a low dimensional manifold, so unlike the noise.
Sampling above the nyquist sampling limit then representing it as something else is not beating the nyquist sampling limit, it's just compression.
You may enjoy randomised algorithms (https://www.cl.cam.ac.uk/teaching/2425/RandAlgthm/)
The entire field of operations research is essentially founded on the premise that NP hard is just a starting place.
If anything, NP hard problems are usually the ones worth tackling.
Here's one of my own efforts on density-aware locality-sensitive hashing for embeddings. It uses k-means at one stage, I'll give this approach a shot.
Hi, I'm the author of e-means. If you can send them, I would love to hear about the results you got !
So this is a different clustering algorithm, of which there are many, maaaany?
That makes it... not very interesting, because it's not a drop-in replacement then, and potentially many other clustering methods would have worked in particular applications as well. Or does this one beat most of the others as well?
It's not a new kind of clustering. It solves exactly the same problem as k-means (minimize the mean squared distance to k centroids), returns the same thing (k centroids and an assignment), and has the same inputs (k and a number of epochs). The batch size and step size are set automatically, so there's nothing extra to tune. So it is a drop-in replacement for k-means specifically: in Faiss it's Emeans(d, k).train(x) where you'd call Kmeans(d, k).train(x). The code is in faiss at https://github.com/facebookresearch/faiss/tree/main/contrib/...
so the epoch reduction is cool but the hyperparameter laws feel like they were fit post-hoc
Hi bashhq_r, I'm Swann Bessa, the author of the e-means paper.
If you look at our paper on arxiv (https://arxiv.org/abs/2611.09501), Table 2, the hyperparameters laws have been validated on 4 datasets outside of the fit, but also on scales outside of the fit. There the hyperparameters they picked made the loss only 0.07% worse than the best hyperparameters possible on average. The robust results in the paper, which are on 26 different datasets where e-means beats Lloyd by a wide margin, are also a proof of the generalization of our scaling laws.
Can someone explain in layman’s terms the contexts this is used?
Let's say you have a way of encoding the properties of something using a vector of numbers. But two things might be very similar and yet still produce vectors that are not identical to each other. Indeed, the vectors are nearly guaranteed to always be at least slightly different. Also, assume there are a large number of dimensions.
You want to assign things to a set of clusters so that things that share similar properties are grouped closely together. But, we also want those clusters to be as big as possible. We don't want to leave little islands that only have a few things in them.
Take, for example, classification of email messages. Let's say you have a billion email messages and you want to find a way to cluster them automatically into 10 categories. You don't know what the categories are ahead of time. You just have a way of making sure that an email can be represented by a vector of numbers such that two emails that are close together in their meaning and characteristics will have similar numbers.
K-Means clustering solves that problem but is intractable to calculate for large datasets having high dimensionality (such as the email example). My understanding is that E-Means is away to achieve this spectacularly more efficiently than the best alternative previously known.
It's clustering. You have a bunch of data points and you want to group them into k clusters such that all the data points in a cluster are similar. Like dividing customers into high/medium/low spenders, or whatever (usually you'd have more dimensions than 1D though).
This is what I was looking for. I appreciate all the responses.
Faster vector calculations. What help for example, helps with ivfflat or prism index building. It cut CPU (or GPU) usage, and time needed to build. Seen index build times drop by 50 to 80%, and at the same time CPU usage drop by 50%.
Impressive results ...
I found this unhelpful, sorry. "Faster vector calculations" sounds like a faster way to add or multiply vectors, which TTBOMK k-means doesn't help with at all. There are hundreds of types of things called an "index"; I'm familiar with ordinary RDBMS indices but haven't heard of ivfflat or prism indices. Gemini tells me PostgreSQL uses them to support approximate nearest neighbour queries on geometric data -- e.g., finding the nearest train station to your current position.
Well, thanks for the downvote.
what's the common use case for e-means/k-means nowadays?
Minimizing Mean Squared Error between a dataset and a given number of centroids. In practice it can be used in deduplication, data curation, compression, Vector Search, and many others. It remains one of the most used Machine Learning algorithms in the world.
What is most impressive about this excellent work is not just the staggering time savings but the robustness of the technique in its assignment of single centroids where legacy may create two due to artefacting or worse one centroid for two blobs. Their explanation that: mini-batches provide more noisy updates that are more likely to explore the optimization landscape and escape local minima makes perfect sense and just goes to show it is all not settled hash. Incredible work with real world benefits of real marginal utility.
Very cool, saving this to play around with later.
Hi, I'm Swann Bessa, the author of the paper.
If you want, you can also play with it now in your browser at : https://e-means.github.io/playground.html
Feel free to ask any questions, I'll be happy to answer.
I'd want to see final cluster quality, not just inertia, before treating "one-tenth the epochs" as a drop-in win. On badly imbalanced or overlapping clusters, Lloyd's often spends early iterations bouncing centroids around. If the interpolation mainly damps that motion, you can look converged on the objective while the partition is still poor. A plot of ARI or NMI versus epoch budget on messy data would make the claim a lot more convincing.
Hi, I'm Swann Bessa, the author of the e-means paper. FYI, we get very good results on ARI and NMI on ImageNet and also synthetic data. We will release them in a later revision. If you have another distribution you have in mind for us to test, I would be glad to run it as well.
I feel like they cover that, a bit, with both explicit examples and the quality curves, right?
Their example is terrible, there are clearly 3 clusters and trying to find 200 clusters doesn't make any sense. They show their algorithm performs better on some degenerate case but what about when there are actually k clusters?
EDIT: actually I see that they have a section about "The k-means problem" and that's where I think they get it wrong. What they describe is not the problem itself but how to solve it. The real problem (for most people at least, they seem to have a different problem) is to find the clusters in a dataset where you already know that there are clusters and how many there are.
Hi Thomashuet, I'm Swann Bessa, the author of e-means.
The goal of e-means is to be a drop-in replacement for k-means. It solves the exact same problem, with the exact same inputs and delivers better centroids and quality in a fraction of the time. So I'm not sure what we get wrong by working on a problem that hundreds of papers also worked on. Is the whole k-means literature getting it wrong then ?
Picking the k is a separate problem, on which also a body of literature has worked too. e-means is a drop-in replacement for k-means so you can use any method you would use for choosing the k but with e-means instead of k-means.
e-means performs better on a very wide array of datasets, definitely not degenerate cases, don't hesitate to check out our paper at https://arxiv.org/abs/2611.09501.
If you have any question or objection, I'll be happy to answer.