E-means: 10x faster k-means at equal loss

E-means is a mini-batch alternative to Lloyd’s algorithm for clustering a fixed dataset. Rather than recomputing each cluster’s exact mean after a full pass over all points, it updates cluster sums and counts with exponential moving averages, allowing recent assignments to matter more than older ones. The method derives its batch size and update rate from scaling laws based on dataset size, cluster count, and training duration, so it requires no extra hyperparameter tuning. It also uses a cosine-scheduled update rate, improved initialization and dead-cluster recovery, a finalization step, and specialized parallel kernels to turn fewer updates into practical runtime savings.

Across varied datasets and scales, e-means reportedly reaches a lower test loss than Lloyd in 1.2 to 3 epochs, compared with 25 epochs, and cuts GPU runtime by 10-45 times versus Flash-KMeans and CPU runtime by 8-15 times versus Faiss. Toy experiments suggest mini-batch noise helps it escape poor local minima: across 200 initializations, it was substantially less likely than Lloyd to remain trapped in the tested configurations. The authors attribute the advantage to exploration from noisy updates, while Lloyd’s steadily decreasing loss can keep it in a bad solution. They also report that moving-average updates are more stable than gradient-based centroid updates at large learning rates, partly because they keep centroids within the convex hull of assigned points.