Flash-GMM : un noyau économe en mémoire pour le clustering flou passant à l'échelle

Résumé

Nous présentons Flash-GMM, un noyau Triton fusionné conçu pour le calcul efficace de modèles de mélanges gaussiens (GMM) sur des données à grande échelle en un seul passage sur GPU. En évitant de matérialiser l'intégralité de la matrice des responsabilités dans la mémoire GPU, Flash-GMM atteint une accélération d'un facteur 20 par rapport aux implémentations existantes et permet l'entraînement sur des ensembles de données plus de 100 fois plus volumineux qu'auparavant sur un seul dispositif. Pour démontrer son impact, nous intégrons Flash-GMM dans le quantificateur grossier IVF pour la recherche approximative des plus proches voisins (ANN). Nous montrons que le clustering GMM doux constitue désormais un remplacement direct viable des k-moyennes, et que les responsabilités issues du GMM peuvent être exploitées pour affecter les vecteurs frontières à plusieurs clusters. Notre approche atteint des cibles de rappel fixes avec jusqu'à 1,7 fois moins de calculs de distance, ou, de manière équivalente, offre un rappel@10 amélioré de +2 à 12 pour un coût de calcul comparable. Nous publions le noyau en tant que projet open-source.

English

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.