Skip to content
Search

E-means

lastmod 2026-11-11created 2026-11-112 edits

k-means that is about 10 times faster at the same test error

About

  • A k-means algorithm that combines mini-batches with an exponential moving average (EMA). Across many datasets and scales, it reaches the test error of 25 epochs of Lloyd's algorithm in just 1.2–3 epochs
  • Total run time to reach the same test error is 10–45 times shorter than Flash-KMeans on GPU and 8–15 times shorter than Faiss on CPU
  • Updates the per-cluster sum S_j and count N_j as S_j ← (1−α)S_j + αs_j and N_j ← (1−α)N_j + αn_j, reducing the influence of stale assignments. Lloyd resets the statistics on every pass, while MacQueen accumulates all past statistics
  • Gradient-descent-based mini-batch k-means (Meta FAIR) was unstable at large learning rates, whereas with EMA the centroids stay inside the convex hull of their assigned points, which keeps it stable. On DEEP 1M and SIFT 10M with 16,384 clusters and 5 epochs, its test error is 1.6% and 0.4% lower, respectively, than tuned gradient descent
  • Batch size and initial step size are set automatically by scaling laws based on the number of points (n), number of clusters (k), and epochs (T), so no hyperparameter tuning is needed beyond what Lloyd requires
    • Default batch: B = 0.035 × k^0.13 × n^0.88 × T^0.81 (coefficient 0.0044 when prioritizing quality)
    • Step: α₀ = min(0.824 × k^−0.46 × B^0.90 × n^−0.45 × T^−0.11, 1), decayed to 0.1α₀ with a cosine schedule
    • A final step recomputes the centroids as the mean of the last epoch's assignments, further lowering the training loss
  • Noise from mini-batch updates makes it less likely to get stuck in local minima. In a toy-dataset experiment with 200 seeds, e-means got stuck 0–4% of the time versus 15.5–46% for Lloyd
  • Dead clusters are recovered by splitting a large cluster in two, chosen with probability proportional to its size (a method borrowed from Faiss k-means). With 1M vectors and 16,384 clusters, the dead-cluster rate is 0.3% on DEEP (14.3% with batch resampling, 26.3% with no recovery) and 0.1% on SIFT, preserving codebook capacity
  • Validated on 234 configurations across 16 datasets ranging from 100,000 to 100 million points. On all 11 VIBE embedding datasets (200–3,072 dimensions), it reached the error of 25 Lloyd epochs in 1.6–2.4 epochs
  • Provides a dedicated GPU kernel with fused nearest-centroid assignment. The implementation is published in the Faiss repository under faiss/contrib/e_means; use Emeans for PyTorch tensors and NumpyEmeans for NumPy. The default training budget is 5 epochs, and train(x) returns the centroids and the cluster labels from the last epoch

See also

Favorite site

author MCPBot sitemap