deniz.in

Markets

Weather

Loading weather

· via Hacker News – Front Page (native)

E-means claims 10-45x faster k-means at equal loss via EMA mini-batch updates

E-means, a mini-batch reworking of k-means built on exponential moving average updates, claims 10-45x GPU speedups over the best Lloyd implementation at equal or better loss, with hyperparameters set automatically by scaling laws.

E-means claims 10-45x faster k-means at equal loss via EMA mini-batch updates

A new take on an old workhorse

A project called e-means has published a reworking of k-means clustering that claims order-of-magnitude speedups without sacrificing solution quality. According to the project's documentation, which recently reached the Hacker News front page, e-means arrives at a test loss better than what Lloyd's algorithm achieves after 25 epochs using only 1.2 to 3 epochs. Measured end to end, the project reports a 10-45x wall-clock reduction on GPU compared with Flash-KMeans, which it describes as the state-of-the-art Lloyd implementation, and 8-15x on CPU compared with Faiss. In clustering, speed and quality normally pull in opposite directions; the project's claim is that e-means improves both at once.

A different update rule

The project frames existing k-means variants as variations on one theme. Every method it surveys keeps, for each cluster, a running sum and count of assigned points, and defines the centroid as their ratio. Lloyd's algorithm recomputes those statistics from scratch over the full dataset each pass, typically for 20 to 50 iterations. MacQueen's 1967 update accumulates them indefinitely, so a centroid becomes the uniform mean of every point ever assigned to it, including points assigned when the centroids were still badly placed. E-means takes a third route: each mini-batch is blended into the statistics with an exponential moving average, so old assignments fade at rate α and centroids track recent data instead of the full history.

That style of update is normally used for moving distributions, such as codebook learning in VQ-VAEs. The project's finding is that it is also highly effective on static datasets, provided the rest of the method is designed around it.

No tuning, plus finishing touches

Two hyperparameters — batch size and step size — are not left to the user. A hyperparameter sweep reportedly shows the lowest losses lying along a diagonal band, where larger batches call for larger steps, while Lloyd, which corresponds to full batch and step 1, sits in a corner far from that optimum. The project fitted scaling laws to that sweep so e-means picks both values automatically. The method also anneals the step from α₀ down to a tenth of it on a cosine schedule, adds a new initialization, revives dead clusters by splitting large ones when they shrink below a threshold, and finishes training on each cluster's exact last-epoch mean for negligible extra runtime. Fused GPU kernels for the nearest-centroid assignment step are what convert the epoch savings into matching wall-clock gains.

Why not gradient descent

The project grew out of earlier work at Meta FAIR by Hugo Touvron and Hervé Jégou, who found that mini-batch k-means optimized by gradient descent could beat Lloyd. The e-means author says those results reproduce, but that gradient-descent updates turn unstable at large learning rates while moving-average updates do not. The suspected reason: an EMA update keeps a centroid inside the convex hull of its assigned points, whereas a gradient step can overshoot past it. At each method's best setting, gradient descent ended 1.6% above the moving average on the DEEP dataset and 0.4% on SIFT, which is why the EMA variant gives the method its name.

The documentation also reports that e-means escapes local minima more reliably than Lloyd, based on repeated runs over 200 seeds on toy datasets.

Caveats

All of the headline numbers come from the project's own benchmarks on its website; they have not been peer reviewed or independently replicated. The baselines are also specific — Flash-KMeans on GPU, Faiss on CPU — so the speedup a given workload sees will depend on its scale, dimensionality and cluster count. The project says its results hold across a large range of scales and datasets, but adopters should treat the figures as self-reported until third parties confirm them.

Why it matters

K-means sits underneath a lot of practical computing: quantizing embeddings, preparing data before model training, and clustering inside larger ML pipelines. As the project itself notes, it is one of the rare machine learning problems where full-batch optimization is still the norm. If a 10-45x speedup at equal or better loss holds up, clustering stops being a bottleneck for workloads that run it at scale, on both GPU and CPU. The broader lesson may travel further: optimization machinery from modern deep learning — moving averages, learning-rate schedules, hyperparameters derived from scaling laws — can be transplanted onto classic algorithms that never received that treatment.

  • #machine-learning
  • #clustering
  • #k-means
  • #optimization
  • #gpu

Related posts