Flash-KMeans: Fast and Memory-Efficient Exact K-Means
重新审视经典 k-means 算法在现代 AI 系统中的定位,将其从离线预处理工具提升为在线系统的一等原语。论文发现现有 GPU 实现的瓶颈在底层系统约束(内存带宽、warp 调度)而非算法复杂度。通过重新设计 GPU kernel 的内存访问模式和并行策略,Flash-KMeans 在 50M 量级 embedding 数据上实现了数量级的加速,同时保持精确结果(非近似)。
INSIGHT将 k-means 作为在线原语而非离线工具,对搜索推荐系统意义重大:embedding 聚类可实时驱动索引分片、召回路由和用户分群。Flash-KMeans 的系统级优化思路——从 GPU 内存访问模式入手而非改变算法——对所有涉及大规模 embedding 操作的基础设施都有借鉴价值。该方法可直接应用于在线向量索引的动态聚类更新,降低向量检索系统的索引重建延迟。