Flash-GMM: Эффективное по памяти ядро для масштабируемой мягкой кластеризации

Аннотация

Мы представляем Flash-GMM — объединённое ядро Triton для эффективного вычисления моделей гауссовых смесей (GMM) над крупномасштабными данными за один проход GPU. Благодаря исключению необходимости материализации полной матрицы ответственности в памяти GPU, Flash-GMM достигает 20-кратного ускорения по сравнению с существующими реализациями и позволяет обучать модели на наборах данных, более чем в 100 раз превышающих ранее возможные пределы на одном устройстве. Для демонстрации его влияния мы интегрируем Flash-GMM в грубый квантор IVF при поиске приближённых ближайших соседей (ANN). Мы показываем, что мягкая кластеризация на основе GMM теперь является жизнеспособной прямой заменой k-средних, а ответственности GMM можно использовать для назначения граничных векторов нескольким кластерам. Наш подход достигает заданных порогов полноты, требуя до 1,7 раза меньше вычислений расстояний, или, что эквивалентно, даёт прирост полноты@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.