Flash-GMM: 확장 가능한 소프트 클러스터링을 위한 메모리 효율적인 커널

초록

본 논문에서는 단일 GPU 패스에서 대규모 데이터에 대한 가우시안 혼합 모델(GMM)의 효율적 연산을 위한 융합 트리톤 커널인 Flash-GMM을 제시한다. 전체 책임 행렬을 GPU 메모리에 구체화할 필요를 없앰으로써, Flash-GMM은 기존 구현 대비 20배의 속도 향상을 달성하며, 기존에 한 장치에서 가능했던 것보다 100배 이상 큰 데이터셋에서의 학습을 가능하게 한다. 그 영향력을 입증하기 위해, Flash-GMM을 근사 최근접 이웃(ANN) 검색을 위한 IVF 조대 양자화기에 통합한다. 부드러운 GMM 클러스터링이 이제 k-평균을 대체할 수 있는 실용적인 대안이며, GMM 책임도를 활용하여 경계 벡터를 다중 클러스터에 할당할 수 있음을 보여준다. 본 접근법은 최대 1.7배 적은 거리 연산으로 고정된 재현율 목표에 도달하거나, 동등한 계산 비용에서 2~12포인트의 recall@10 향상을 제공한다. 해당 커널을 오픈소스 프로젝트로 공개한다.

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.