Flash-GMM: スケーラブルなソフトクラスタリングのためのメモリ効率的なカーネル

要旨

本稿では、Flash-GMMを提案する。これは、融合型Tritonカーネルであり、大規模データに対する混合ガウスモデル（GMM）の計算を、単一のGPUパスで効率的に行う。Flash-GMMは、完全な負担率行列をGPUメモリ上に具体化する必要をなくすことで、既存実装と比較して20倍の高速化を実現し、これまでは1台のデバイスでは不可能だった100倍以上のデータセットでの学習を可能にする。その効果を示すため、Flash-GMMを近似最近傍（ANN）探索用のIVF粗量子化器に統合した。この結果、ソフトGMMクラスタリングがk平均法のドロップイン置換として実用的であること、またGMMの負担率を活用して境界ベクトルを複数のクラスタに割り当てられることを示す。本アプローチでは、固定された再現率目標に対して最大1.7倍少ない距離計算で到達するか、あるいは同等の計算コストでrecall@10が2～12ポイント向上する。本カーネルをオープンソースプロジェクトとして公開する。

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.