| title | Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering |
|---|---|
| authors | Gal Bloch, Ariel Gera, Matan Orbach, Ohad Eytan, Assaf Toledo |
| date | 2026-06-09 |
| venue | arXiv (preprint, not peer reviewed) |
| id | arXiv:2606.10896 |
| affiliation | IBM Research |
| upvotes | 8 (Hugging Face Daily Papers, 2026-11-01) |
| links | abs pdf hf code |
| abstract | We present Flash-GMM, a fused Triton kernel for efficient computation of Gaussian Mixture Models (GMMs) over large-scale data in a single GPU pass. By eliminating the need to materialize the full responsibility matrix in GPU memory, Flash-GMM achieves a 20times speedup over existing implementations and enables training on datasets more than 100times larger than previously feasible on one device. To demonstrate its impact, we integrate Flash-GMM into the IVF coarse quantizer for approximate nearest-neighbor (ANN) search. We show that soft GMM clustering is now a viable drop-in replacement for k-means, and that GMM responsibilities can be leveraged to assign border vectors to multiple clusters. Our approach reaches fixed recall targets with up to 1.7times fewer distance computations, or equivalently, yields +2--12 recall@10 at matched computational cost. We release the kernel as an open-source project. |
cite
@misc{bloch2026flash,
title={Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering},
author={Gal Bloch and Ariel Gera and Matan Orbach and Ohad Eytan and Assaf Toledo},
year={2026},
eprint={2606.10896},
archivePrefix={arXiv},
url={https://arxiv.org/abs/2606.10896}
}