面向异构架构的可扩展K近邻图构建优化研究

《ACM Transactions on Parallel Computing》：Scalable KNN Graph Construction for Heterogeneous Architectures

本文系统介绍了PyRKNN库，这是一种针对分布式异构高性能计算（HPC）系统优化的K近邻（kNN）图构建新方法。作者通过设计新颖的GPU批处理内核（实现1.18倍至19倍加速）、优化通信开销（实现5倍加速）并支持CPU/GPU上的稠密与稀疏坐标，解决了大规模稀疏非整数数据集上构建精确kNN图的挑战。该工作为机器学习（ML）和科学计算中的聚类、密度估计等任务提供了高效解决方案。

　　

可扩展K近邻图构建的挑战与创新

在机器学习和科学计算的众多应用中，K近邻（kNN）图的构建是一个基础且关键的环节。尽管其应用广泛，但在分布式异构高性能计算（HPC）系统上，高效地构建大规模的全最近邻图仍然是一个艰巨的挑战，特别是对于大型稀疏非整数数据集。传统的直接计算方法需要计算所有点对之间的距离，其工作量随着数据点数量N的增长呈O(N²)级别增长，这在处理百万级甚至更多数据点时变得不切实际。虽然对于低维数据集，基于树的方法（如kd-树或八叉树）可以通过在距离计算前剪枝远处点来显著降低计算成本，但在高维情况下，这些方法的有效性会迅速下降，这就是所谓的“维度灾难”。

为了应对这一挑战，研究者们提出了两种主要的缓解途径。其一，尽管数据集可能嵌入在高维空间中，但其元素可能局部存在于低维流形上，算法可以设计为利用这种低内在维度结构。其二，放宽问题要求，只寻找每个点的近似最近邻，并接受一定的失败概率，可以使kNN图构建对于更高维度的问题变得可行。在众多近似最近邻搜索方法中，随机化分箱方法，如局部敏感哈希（LSH）、倒排文件索引（IVF）和随机投影树（RP-Trees）森林，因其能通过概率技术将可能彼此接近的点分组到同一子集中而成为一种有效的策略。

随机投影树的核心思想与实现

随机投影树（RP-Trees）为近似kNN图构建提供了一种高效的方法，其通过递归地将点划分为可能包含近邻的桶（bins）。该算法遵循一个迭代过程：在每次迭代中，通过沿随机超平面对点进行分割，将数据集重新分区，形成一个平衡的L层二叉树。在平衡树中，每个叶子节点总是包含大约m = N / 2^L个点。RP-Trees的理论保证是，随着迭代次数的增加，它们会收敛到真正的最近邻。如果数据集的倍增维度是d_o，那么为每个点构建精确图需要O(d_o^d_o + log N)的工作量。对于稀疏数据集，倍增维度可以表示为O(s + s log d)，其中s是数据集中所有元素非零坐标的最大数量。

在PyRKNN库的实现中，采用了多项优化措施。在树构建阶段，选择投影方向时，为了减少通信开销，采用了固定投影而非数据依赖的投影。具体来说，对树中每一层l的所有节点重复使用单个投影方向q_l ∈ R^d，这将唯一投影器的总数从2^L - 1减少到L。更重要的是，这允许通过单次矩阵乘法Y = XQ计算所有投影，其中X ∈ R^N×d按行存储数据集的坐标，Q = [q₁ … q_L] ∈ R^d×L存储投影方向。对于稠密坐标，这实现为稠密-稠密矩阵乘法；对于稀疏坐标，则变为稀疏-稠密乘法。通过预先计算所有投影，提高了内存局部性，并避免了在树构建过程中重复访问全维坐标。

批处理局部搜索的GPU内核优化

在RP-Tree方法中，计算成本最高的部分是在许多不相交的桶上进行穷举搜索。为了优化这一过程，PyRKNN设计了专门的GPU内核用于批处理精确搜索。研究比较了两种用于GPU上分块距离和合并计算的方法：行分区（Row Partitioning）和混合分区（Hybrid Partitioning）。

行分区将每个叶子的坐标按行划分为大小为b × d的水平块。对于一组叶子，计算每个叶子中第j个行块与该叶子中所有其他点之间的距离。这通过批处理乘积X_i,jX_i^T来实现。在计算距离之后，可以跨叶子选择所有块中每一行的kNN，并将其合并到近似全局图中。

混合分区则更好地利用了每个叶子距离矩阵D(X_i)的对称性。它将距离矩阵分解为捕获对角线的方形块和捕获剩余上三角部分的矩形块。对于每个块，计算其中的距离，并执行k选择以更新找到的近邻。对于非对角矩形块，其过程与行分区类似。给定一批叶子和一个活跃索引j，计算每个叶子距离矩阵中第j个行块所包含的距离。距离矩阵中的每个距离元素代表批处理中两个不同点（对应行和列索引）的潜在最近邻距离。为了更新这些点中每一个的最近邻，在块内对行和列执行单独的合并。水平合并从行中的m - jb + k个距离中选择k个近邻，而垂直合并从列中的b + k个距离中选择k个近邻。这两种合并操作可以并发执行，因为它们更新最近邻图的不重叠部分。这个过程对每个j重复，直到所有块都被处理完毕。

对于稠密距离计算，使用cuBLAS的cublasSgemmStridedBatched内核在活跃选择的块上操作。对于稀疏坐标之间的距离计算，由于供应商提供的内核不支持批处理SpGemm或当X和Y都以CSR格式存储时计算XY^T，因此实现了自定义的稀疏外积算法。对于距离计算，使用特定算法计算两组稀疏数据点X = {x_i}_i=1^n_x和Y = {y_j}_j=1^n_y之间的成对距离。为每个点x_i分配一个线程块来计算到所有点{y_j}_j=1^n_y的距离。每个线程处理距离矩阵的一个元素。x_i的非零元素的列索引保存在共享内存中，以便通过无分支二分查找实现与Y的坐标交集的快速查找。

分布式内存树的算法与优化

对于无法在单台机器上处理的大规模数据集，PyRKNN将其扩展到多GPU和多节点配置，使用MPI进行通信。分布式方法包括三个阶段：数据集的重新分布、重复的局部搜索以及局部结果的分布式合并。这些阶段迭代进行，构建和合并多个分布式树，直到达到收敛。

在树构建阶段，坐标最初以一维块行分区分布在P个进程中。在每次迭代中，形成一个分布式的P层RP-Tree，其中每个叶子由单个进程拥有。形成每个分布式树涉及在所有进程之间划分和重新分布整个数据集。为了减轻集体通信的巨大开销，PyRKNN采用了四项主要优化：延迟坐标重新分布、使用随机化共享状态和打破平局、将树构建与局部搜索重叠以隐藏通信时间，以及重复局部搜索。

PyRKNN使用固定投影，而不是在构建树时发送完整坐标，仅通信必要的重新洗牌的全局索引和投影。当L << d对于稠密坐标时，这具有显著优势。对于稀疏坐标，这避免了在树构建的每一层将数据拆分并重新打包成CSR格式的需要。投影Y = X_lQ在树构建之前通过矩阵乘法计算，可选的正交化步骤复杂度为O(d log²(P))。为了避免通信，使用共享随机种子在每个进程中冗余生成并正交化相同的随机投影向量。

重复局部搜索是另一个关键优化。在每个进程内执行多次局部树搜索，以查找进程内数据的近似近邻。由于通信昂贵，在执行下一次分布式树搜索之前，执行多次局部搜索以提取部分信息非常重要。增加局部搜索次数可以显著提高召回率并减少时间。然而，需要选择局部搜索的次数以避免收益递减。

性能评估与对比分析

研究对PyRKNN库在CPU和GPU配置下针对稀疏和稠密数据集的性能进行了全面评估。在批处理精确kNN内核的性能测试中，对于稠密数据，使用GSKNN内核作为参考CPU实现；对于稀疏数据，则使用Intel MKL的sparse_s_syrkd进行距离计算，并在过滤唯一候选者后使用std::nth_element进行近邻选择。评估显示，提出的GPU内核在所有情况下均优于CPU参考实现，其中混合分区内核在距离计算上表现更好，因为它利用对称性节省了大约一半的计算量。对于更高维的数据集，如GAUSS256，这种优势尤为明显，因为计算距离是主导成本。

在单节点端到端性能评估中，PyRKNN与常用的近似最近邻搜索库进行了比较，包括FAISS、NMSlib和GGNN。结果显示，PyRKNN-GPU在大多数数据集上实现了与FAISS-GPU竞争的性能，并且在许多数据集上更快。例如，在Embed200数据集（包含嵌入200维空间的6维高斯斑点和球体，并添加了均匀噪声）上，PyRKNN比FAISS-CPU快至少3倍。与GGNN的比较表明，在SIFT1M数据集上，PyRKNN在低目标命中率下保持竞争力，尽管在高命中率下GGNN可能快达3倍，但PyRKNN具有更简单的参数空间和易于迭代细化的优势。

在分布式内存设置下的性能评估表明，PyRKNN-CPU在每次迭代的速度、可扩展性以及对每个进程核心数的敏感性方面均显著优于GOFMM。PyRKNN在节点之间交换的数据量比GOFMM少，并且每次分割的工作量更少，因为它只在树的叶级别通信点的坐标。重叠树构建和局部搜索可以在扩展到更多GPU时加快端到端执行时间并减少可见的通信开销。此外，适度的局部迭代次数（每次分布式迭代5到10次）可以实现接近最佳的性能。

相关工作的比较与展望

在相关工作方面，近似kNN搜索方法主要分为分箱方法（如LSH、IVF、RP-Trees）和图遍历方法（如HNSW、NN-Descent）。分箱方法提供可调的精度与工作量权衡、负载平衡的分区以及收敛保证。在图遍历方法中，搜索方法（如HNSW）并不直接寻找数据集的kNN图，而是通过预处理初始图来构建一个具有改进路由效率和连通性的图索引。在查询阶段，从一组入口点贪婪地遍历索引以找到目标的近似近邻。搜索方法在实践中表现非常好，以较少的距离比较实现高精度。然而，其对于一般数据集的理论精度保证仍然是一个悬而未决的问题，并且它们通常有许多超参数需要调整索引构建。

PyRKNN的贡献在于提供了一个针对分布式异构架构优化的可扩展kNN图构建框架。通过从内核级别到分布式层的精心设计，该方法有效地平衡了计算成本、精度和内存使用。虽然随机投影树不是新方法，并且通常比一些现代图遍历方法具有更慢的每次查询时间，但由于其构建成本低、精度与工作量之间易于调整的权衡、负载平衡的分区以及收敛保证，它们对于全最近邻图构建来说是一个有竞争力的解决方案。在分布式设置中，当达到目标精度所需的迭代次数少于进程数时，RP-Tree迭代尤其有效。

未来的工作方向包括开发自适应数据驱动技术来调整叶子大小和子树迭代次数等参数，探索量化技术在搜索期间压缩数据集的方法，以及进一步优化分块精确搜索内核中的操作融合。PyRKNN框架为未来的研究和实际应用奠定了坚实的基础，有望成为研究人员和从业者在大规模数据分