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

e-means. Outperforming k-means (in one-tenth the epochs). By Swann Bessa and six others, from FAIR at Meta and the École Nationale des Ponts et Chaussées. Published on the sixth of November.

Lloyd's k-means algorithm has been the standard for clustering for a long time. It turns up in data analysis, in retrieval, in quantization, and in data curation. Being the default, though, is not the same as being fast or reaching the lowest loss, and the title already signals that the authors think epochs are being spent that need not be.

Their proposal is an exponential moving average optimization procedure for k-means, using mini-batch updates. They describe it as powered by modern deep learning machinery together with new insights. The resulting algorithm is named e-means.

A procedure like this brings its own settings, chiefly the decay rate and the batch size. Rather than leave those to be tuned, the authors establish scaling laws for how each affects final performance. From those laws comes a formulaic rule that fixes both across multiple problem sizes and embedding datasets, with no hyperparameter sweep needed. Someone running e-means on a new problem reads the settings off the rule instead of searching for them.

The comparison is set by what Lloyd k-means reaches after twenty-five epochs. Across a wide array of embedding datasets, e-means beats that loss within one point four to two point six epochs. In GPU wall-clock terms that comes to speedups of ten to forty-five times over state-of-the-art k-means implementations. It converges to lower losses than Lloyd k-means and frequently beats k-means++ as well.

The largest example is DINO10B, a dataset of ten billion DINO embeddings. e-means partitions it into sixteen million clusters in only two point five epochs, and ends below the loss Lloyd's algorithm reaches at twenty-five.

Two gains are claimed at once here: fewer epochs, and a lower loss at the end of them. The scaling laws matter as much as the update itself: a clustering algorithm with no sweep to run is one that can be pointed at ten billion points and simply run.