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

우리는 대규모 데이터에 대해 단 한 번의 GPU 패스(pass)로 가우시안 혼합 모델 (Gaussian Mixture Models, GMMs)을 효율적으로 계산하기 위한 융합된 Triton 커널인 \textbf{Flash-GMM}을 제시합니다. GPU 메모리에 전체 책임 행렬 (responsibility matrix)을 구현할 필요를 제거함으로써, Flash-GMM은 기존 구현 대비 \textbf{20\times}의 속도 향상을 달성하며, 단일 장치에서 이전에 가능했던 것보다 \textbf{100\times} 이상 더 큰 데이터셋에 대한 학습을 가능하게 합니다. 그 영향력을 입증하기 위해, 우리는 Flash-GMM을 근사 최근접 이웃 (approximate nearest-neighbor, ANN) 검색을 위한 IVF coarse quantizer에 통합합니다. 우리는 소프트 GMM 클러스터링이 이제 $k$-means의 실행 가능한 교체 수단이 되었으며, GMM 책임 (responsibilities)을 활용하여 경계 벡터 (border vectors)를 여러 클러스터에 할당할 수 있음을 보여줍니다. 우리의 접근 방식은 최대 $1.7\times$ 적은 거리 계산으로 고정된 재현율 (recall) 목표에 도달하거나, 동일한 계산 비용에서 $+2$--$12$ recall@10을 달성합니다. 우리는 이 커널을 오픈 소스 프로젝트로 공개합니다.