Flash-GMM: Een geheugenefficiënte kernel voor schaalbare zachte clustering

Samenvatting

We presenteren Flash-GMM, een gefuseerde Triton-kernel voor efficiënte berekening van Gaussian Mixture Models (GMM's) over grootschalige data in één enkele GPU-passage. Door de noodzaak om de volledige verantwoordelijkheidsmatrix in het GPU-geheugen te materialiseren weg te nemen, behaalt Flash-GMM een 20-voudige versnelling ten opzichte van bestaande implementaties en maakt het training mogelijk op datasets die meer dan 100 keer groter zijn dan voorheen haalbaar op één apparaat. Om de impact ervan te demonstreren, integreren we Flash-GMM in de IVF-grove quantizer voor approximate nearest-neighbor (ANN) search. We laten zien dat zachte GMM-clustering nu een haalbare directe vervanging is voor k-means, en dat GMM-verantwoordelijkheden kunnen worden benut om grensvectoren aan meerdere clusters toe te wijzen. Onze benadering bereikt vaste recall-doelen met tot 1,7 keer minder afstandsberekeningen, of equivalent, levert +2--12 recall@10 bij gelijke rekenkosten. We brengen de kernel uit als een open-sourceproject.

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.