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.
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
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.
saagarjha
Lloyd 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.
Maybe we combine the ideas of e-means and BBD-tree?
haifeng
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.
wood_spirit
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.
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.
comments (9)
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
jlhawn
That's impressive!
real_faxenoff
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
gilfyole
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.
saagarjha
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?
haifeng
Hope it gets added to scipy - it’s much more available in the envs many use than facebooks own faiss library.
wood_spirit
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.
shiandow
Let me know if you hit any issues.
haifeng
leecarraher