e-means hits 10–45× faster k-means on GPU with lower loss than Lloyd's
E-means: 10x faster k-means at equal loss
e-means swaps Lloyd's full-batch updates for mini-batch exponential moving averages, reaching Lloyd's test loss in 1.2–3 epochs instead of 25. End to end on GPU that is a 10–45× wall-clock speedup over Flash-KMeans and 8–15× on CPU over Faiss, with no hyperparameter tuning thanks to fitted scaling laws. It also escapes local minima far more often than Lloyd, and Lloyd turns out to be a suboptimal corner case of the same rule.
Lloyd is an extreme case of e-means without cosine scheduling on the learning rate.
- jlhawn
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
- Luke2642
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!
- real_faxenoff
I adapted the method to my IVF index (the code-daemon project) and achieved a 146x speedup on the GPU.
That's impressive!