论文精选

Flash-GMM:单GPU实现百倍规模软聚类,加速20倍

Flash-GMM: A Memory-Efficient Kernel for Scalable Soft Clustering

精选理由

做大规模聚类或近似最近邻搜索的团队,Flash-GMM 让你在单卡上处理百倍数据量,速度还快 20 倍,直接替换 k-means 就能提升召回率,值得一试。

AI 摘要

Flash-GMM 是一个融合的 Triton 内核,用于在单次 GPU 传递中高效计算高斯混合模型(GMM)。它通过避免在 GPU 内存中实例化完整的责任矩阵,实现了比现有实现快 20 倍的加速,并支持在单个设备上训练比以往大 100 倍以上的数据集。该内核被集成到 IVF 粗量化器中用于近似最近邻搜索,表明软 GMM 聚类可以替代 k-means,并利用 GMM 责任将边界向量分配到多个簇。在达到固定召回率目标时,Flash-GMM 减少了最多 1.7 倍的距离计算,或在相同计算成本下召回率提升 2-12 个点。该项目已开源。

AI 翻译 · 中文

Flash-GMM 是一个融合的 Triton 内核,用于在单次 GPU 传递中高效计算高斯混合模型(GMM)。它通过避免在 GPU 内存中实例化完整的责任矩阵,实现了比现有实现快 20 倍的加速,并支持在单个设备上训练比以往大 100 倍以上的数据集。该内核被集成到 IVF 粗量化器中用于近似最近邻搜索,表明软 GMM 聚类可以替代 k-means,并利用 GMM 责任将边界向量分配到多个簇。在达到固定召回率目标时,Flash-GMM 减少了最多 1.7 倍的距离计算,或在相同计算成本下召回率提升 2-12 个点。该项目已开源。

arXiv cs.LGWe 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