arXiv:2601.17136cs.DCcs.LG2026-01

突破单GPU内存限制,实现百万级数据的高效核聚类。

Communication-Avoiding Linear Algebraic Kernel K-Means on GPUs

  • 用分布式线性代数原语重构核聚类,减少通信开销。
  • 256卡下弱扩展效率79.7%,强扩展提速4.2倍,聚类时间从1小时缩至2秒。
  • 适合超大规模数据的聚类任务,尤其适用于多卡加速场景。

聚类是数据分析的重要工具,K-means因其简单和通用性广受欢迎,但无法处理非线性可分数据。核K-means解决了这一问题,但需构建大型核矩阵,导致计算与内存开销巨大。先前工作通过稀疏线性代数原语在单个GPU上加速核K-means,但受限于显存,仅能处理约8万样本以内的数据。本文提出一套面向多GPU系统的分布式内存并行算法,将核K-means中最耗时的组件映射到专为该任务设计的通信高效线性代数原语,实现百万级数据的高效聚类。核心在于设计分区策略,支持线性代数操作的低通信组合。1.5D算法表现最优,使核聚类规模提升一至两个数量级。在256个GPU上,几何平均弱扩展效率达79.7%,几何平均强扩展速度比为4.2倍。相比1D算法,1.5D在256卡下最高提速3.6倍,聚类时间由超过一小时降至不足两秒。结果表明,基于应用定制的线性代数范式可显著提升性能。

原文摘要 · Abstract (English)

Clustering is an important tool in data analysis, with K-means being popular for its simplicity and versatility. However, it cannot handle non-linearly separable clusters. Kernel K-means addresses this limitation but requires a large kernel matrix, making it computationally and memory intensive. Prior work has accelerated Kernel K-means by formulating it using sparse linear algebra primitives and implementing it on a single GPU. However, that approach cannot run on datasets with more than approximately 80,000 samples due to limited GPU memory. In this work, we address this issue by presenting a suite of distributed-memory parallel algorithms for large-scale Kernel K-means clustering on multi-GPU systems. Our approach maps the most computationally expensive components of Kernel K-means onto communication-efficient distributed linear algebra primitives uniquely tailored for Kernel K-means, enabling highly scalable implementations that efficiently cluster million-scale datasets. Central to our work is the design of partitioning schemes that enable communication-efficient composition of the linear algebra primitives that appear in Kernel K-means. Our 1.5D algorithm consistently achieves the highest performance, enabling Kernel K-means to scale to data one to two orders of magnitude larger than previously practical. On 256 GPUs, it achieves a geometric mean weak scaling efficiency of $79.7\%$ and a geometric mean strong scaling speedup of $4.2\times$. Compared to our 1D algorithm, the 1.5D approach achieves up to a $3.6\times$ speedup on 256 GPUs and reduces clustering time from over an hour to under two seconds relative to a single-GPU sliding window implementation. Our results show that distributed algorithms designed with application-specific linear algebraic formulations can achieve substantial performance improvement.

聚类多卡加速核方法分布式计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。