$e$-means : Outperforming k-means (in one-tenth the epochs)

Lloyd's k-means algorithm is a long-running standard for clustering, with applications in data analysis, retrieval, quantization and data curation. In this paper, we introduce an exponential moving average optimization procedure for k-means with mini-batch updates, powered by modern deep learning machinery and new insights. We also establish scaling laws on how decay rate and batch size impact final performance. This gives us a formulaic rule to fix them for multiple problem sizes and embedding datasets, without the need for any hyperparameter sweep. The resulting clustering algorithm, coined $e$-means, is much faster than Lloyd k-means and converges to lower losses, even when compared to k-means++. It consistently beats the loss of Lloyd k-means at convergence in 1.2 to 3 epochs, which translates to 10-45$\times$ GPU wall-clock speedups compared to state-of-the-art k-means implementations. As an example, $e$-means partitions DINO10B, a dataset of 10 billion DINO embeddings, into 16 million clusters in only 2.5 epochs, reaching a loss below Lloyd's at a tenth of its epochs.