Meta FAIR Releases e-means: A K-Means Alternative with 10x Speedup and Lower Loss

• Meta FAIR developed e-means, a k-means alternative that achieves approximately 10x speedup in epochs and wall-clock time with lower loss compared to Lloyd's k-means, without requiring additional hyperparameters.
• The algorithm is based on a mini-batch update rule using exponential moving averages (EMA), where centroids are calculated as the ratio of two running averages (point count and embedding sum).
• Implementation is available in the faiss library under the contrib/e_means module, supporting CPU (NumPy/PyTorch) and GPU (PyTorch/Triton/CuTe) backends.
• Default configuration requires only nepochs (default 5); batch size and initial step size are automatically derived via empirical scaling laws.

Performance and Backends

• GPU execution supports Hopper (H100/H200) via CuTe kernels (requiring CUTLASS DSL 4.7.1+, QuACK 0.6.5+, cuda-python 13.x) and other NVIDIA GPUs (A100+) via Triton.
• Backend selection is automatic or manual (assignment_backend), with CuTe being the fastest, followed by Triton, and PyTorch as a fallback.
• Data precision supports FP16, BF16, FP8, and INT8 for nearest-centroid search on GPU, while centroids and returned distances remain in FP32.

Configuration and Scaling

• batch_law defaults to "auto" (selecting "quality" for d≤16, otherwise "throughput"); "quality" uses approximately 1/8 the batch size of "throughput" for lower final MSE.
• Scaling laws determine batch size $B$ and initial step size $\alpha_0$ based on vector count $n$, cluster count $k$, and passes $T$.
• max_points_per_centroid limits the training set size; if exceeded, a random subset is used for training and remaining vectors receive label -1.

Advanced Features

• Spherical e-means supports cosine similarity clustering by normalizing vectors and centroids to unit norm.
• Multi-GPU support includes DistributedEmeans (replicated centroids) and EmeansShardedCentroids (sharded centroids) for large-scale training.
• Cluster revival splits busy clusters to revive those with insufficient points, disabled by default for d ≤ 2.