Flash-GMM让大模型软聚类更快更省显存，训练数据量扩大100倍。
Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering
- 用融合Triton核计算GMM，不存全责任矩阵，节省显存。
- 相比旧方法快20倍，单卡可处理超100倍更大的数据集。
- 可用于近似最近邻搜索，召回率提升2%-12%，或减少1.7倍计算量。
我们提出Flash-GMM，一种用于大规模数据上高斯混合模型（GMM）高效计算的融合Triton内核，可在单个GPU遍历中完成计算。通过避免在GPU内存中完整存储责任矩阵，Flash-GMM相比现有实现获得20倍加速，并使单设备训练的数据规模超过以往的100倍。为展示其效果，我们将Flash-GMM集成到IVF粗量化器中进行近似最近邻（ANN）搜索。结果表明，软GMM聚类现在可作为k-means的直接替代方案，且GMM责任值可用于将边界向量分配至多个聚类。该方法在固定召回率目标下，最多减少1.7倍距离计算；或在相同计算成本下，提升2%–12%的recall@10。代码已开源。
原文摘要 · Abstract (English)
We present \textbf{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 \textbf{20$\times$} speedup over existing implementations and enables training on datasets more than \textbf{100$\times$} 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.7\times$ fewer distance computations, or equivalently, yields $+2$--$12$ recall@10 at matched computational cost. We release the kernel as an open-source project.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。