RAPIDS cuML 助力 GPU 实现 UMAP 的高速扩展

UMAP 是一种常用的降维算法，用于生物信息学、NLP 主题建模和 ML 预处理等领域。它的工作原理是创建 k 近邻（k-NN）图（在文献中称为全近邻图），以构建数据的模糊拓扑表示，用于将高维数据嵌入到较低维度中。

RAPIDS cuML 已经包含了加速的 UMAP，与最初基于 CPU 的 UMAP 相比，速度有了显著提升。正如我们在本文中演示的那样，还有改进空间。

在本文中，我们将探讨如何使用 RAPIDS cuML 24.10 中引入的新功能。我们还将深入探讨 nn-descent 算法和批处理流程的详细信息。最后，我们分享基准测试结果，以强调可能的性能提升。在本文结束时，我们希望您对 RAPIDS 更快速且可扩展的 UMAP 所带来的优势感到兴奋。

挑战 

我们面临的一个挑战是，所有邻居图形构建阶段需要很长时间，尤其是与 UMAP 算法中的其他步骤相比。

cuML UMAP 最初仅使用 强力方法 来计算 全近邻图 ，文献中通常将其称为全近邻图，因为它涉及对数据集中的所有向量进行详尽的向量搜索。

暴力法会详尽无遗地计算数据集中每对向量的距离，因此其扩展性往往较差。因此，与 UMAP 中的所有其他步骤相比，随着数据集中向量数的增加，此步骤花费的时间会呈二次增长（向量数的平方）。

图 1 显示了几个热门数据集在全邻图构建中花费的时间比例。在 1M 和 5M 向量规模下，全邻图构建所花费的比例迅速达到 99% 以上。

我们面临的第二个挑战是，与 cuML 中的许多算法一样，整个数据集必须适合 GPU 的内存。

在只有消费级 NVIDIA RTX GPU 可用于处理时，处理大型数据集（例如大小为数百 GB 的数据集）尤其具有挑战性。即使 NVIDIA H100 GPU 提供 80 GB 的内存，这也可能不足以处理 80 GB 的数据集，因为像 UMAP 这样的算法需要许多小的临时内存分配，这些内存分配可能会在算法过程中累加。

加速和扩展 UMAP 

我们使用新的分批式近似近邻（ANN）算法解决了这些挑战。虽然通用方法可以适用于搜索最近邻的任何算法能力，但我们使用了 RAPIDS cuVS 库中名为 最近邻下降 （nn-descent）的快速算法的 GPU 加速版本，该算法非常适合构建全近邻图。

ANN 算法通过权衡质量与速度来加速全邻居图形构建过程。一般来说，近似方法旨在减少寻找最近邻居所需的距离计算。由于此算法可以分段计算单个全邻居图形，因此我们可以在 RAM 内存中放置更大的数据集，并在需要时将所需数据拉入 GPU 内存。

正如我们在本文中演示的那样，我们的新方法以光速将 RAPIDS cuML 24.10 中的 UMAP 扩展为大规模数据集。更好的是，它是默认启用的，因此您无需对代码进行任何更改即可从中受益！

表 1 显示，UMAP 现在可以使用不适合在设备上运行的数据集 (50M、768 为 153GB) 运行。对于大型数据集，速度提升更为显著。过去在 GPU 上运行 10 小时的内容现在可以在 2 分钟内运行。

在 RAPIDS cuML 中使用更快速且可扩展的 UMAP 

如前所述，自 cuML 24.10 起，无需更改代码即可利用此新功能。

但是，为了进行更多控制，UMAP 估计器现在在初始化期间接受另外两个参数：

• build_algo: Algorithm to build the all-neighbors graph. It can be one of the following three values:
  • auto：在运行时根据数据集大小（大于 50K 个数据样本时使用 nn-descent）决定使用暴力或 nn-descent 进行构建。默认值。
  • brute_force_knn：使用暴力法构建全邻接图。
  • nn_descent：使用 nn-descent 构建所有邻居图。
• build_kwds: Python dictionary type for passing parameters related to all-neighbors graph building, with the following parameters:
  • nnd_graph_degree：使用 nn-descent 构建 k-nn 时的图形度。默认值：64。
  • nnd_intermediate_graph_degree：使用 nn-descent 构建 k-NN 时的中级图形度。默认值：128。
  • nnd_max_iterations：运行 nn-descent 的最大迭代次数。默认值：20。
  • nnd_termination_threshold：提前停止 nn-descent 迭代的终止阈值。默认值：0.0001。
  • nnd_return_distances：是否返回 nn-descent 的距离。应将其设置为 true，以便将 nn-descent 与 UMAP 结合使用。默认值：True.。
  • nnd_n_clusters：用于批处理方法的集群数量。在使用更大的数据集运行时，集群数量越多，内存占用率就越低。默认值：2。
  • nnd_do_batch：应设置为 True 进行批处理。默认值：False。

您还可以选择将数据放在主机上，而不是使用 data_on_host 选项将整个数据放在设备上。该选项默认为 False。这仅与 build_algo=”nn_descent” 兼容，不支持使用 brute-force 算法进行构建。

我们建议您将数据放在主机上，以便充分利用我们的大型数据集批处理算法。

from cuml.manifold.umap import UMAP

data = generate_data()

# running default. Runs with NN Descent if data has more than 50K points
umap = UMAP(n_neighbors=16)
emb  = umap.fit_transform(data)

# explicitly set build algo. Runs with this regardless of the data size. Data can be put on host
umap = UMAP(n_neighbors=16, build_algo="nn_descent", build_kwds={"nnd_graph_degree": 32})
emb = umap.fit_transform(data, data_on_host=True)

# batching NN Descent with 4 clusters
umap = UMAP(n_neighbors=16, build_algo="nn_descent", build_kwds={"nnd_do_batch": True, "nnd_n_clusters": 4})
emb = umap.fit_transform(data, data_on_host=True)

为什么近似近邻？

Brute-force 是一种精确且详尽的算法。相比之下，ANN 算法并不能保证找到精确的近邻，但它们可以高效地浏览搜索空间，以更快地构建到最近邻的近似值，代价是牺牲一些准确性以换取搜索速度。

近邻下降（nn-descent）是一种 ANN 算法，可以直接近似计算全邻图。该算法首先为每个数据点随机初始化最近邻，然后通过探索每个点的近邻的近邻来迭代改进最近邻近似。

如原始论文所述 ，nn-descent“通常会收到 90%以上的召回率，而每个点的召回率平均仅占整个数据集的几个百分点。” 简而言之，ANN 算法通常会找到减少必须计算的距离数的巧妙方法。

我们使用 NVIDIA cuVS 库 中的 nn-descent 为 UMAP 构建全邻接图。对于大型数据集，此方法将全邻接图构建过程加速数百倍，同时仍保持功能等效的结果。

使用批处理扩展所有邻居图构建

通过将大型数据集保留在主机上并在设备上批量处理来管理大型数据集似乎很简单。然而，在构建 k-NN 子图形时，使用数据集的特定子集会遇到一个关键挑战：具有相似索引的数据样本不一定靠近。因此，您无法简单地将数据集分成批次。

我们采用了一种受热门 DiskANN 算法相关文献启发的批处理方法，从而解决了这一问题。我们首先在数据集的子样本上执行平衡的 k-means 聚类，以便为预定义数量的聚类提取质心。然后，我们使用这些信息，根据最近的聚类将数据集分割成批处理。

此方法可确保每个批量中的数据点更有可能相互接近，从而提高在同一批量中找到最近邻点的可能性。本节的其余部分将详细说明批处理过程的每个步骤：

• 提取集群核心
• 为每个集群查找数据点
• 构建集群数据点的子图形
• 合并 k-NN 子图形与全局所有邻居图形

提取集群核心 

我们首先从数据集中提取集群重心。由于我们假设大型数据集不适合 GPU 设备，因此我们将数据保留在主机内存中，并随机对一组点进行了二次采样，以确保子集适合 GPU 设备内存。通常，数据集的 10%是足够大的子样本，可以找到一组可用的重心。

使用用户提供的 nnd_n_clusters 参数，我们在采样子集上运行平衡的 k-means，以识别指定数量的集群中心。

为每个集群查找数据点 

接下来，我们为每个数据点确定了前两个最近的集群中心，然后反转索引以查找属于每个集群的数据点。该过程的结果是将每个数据点分配给两个单独的集群。

这种方法可确保每个聚类的邻域都有重叠，从而增加最终邻域至少包含可接受数量的近邻的可能性，就像我们计算出确切结果时所期望的那样。

构建集群数据点的子图形 

当我们知道属于每个集群的数据点后，我们继续在每个集群的数据点上迭代构建子图形。这意味着，对于每个集群，我们在 GPU 的内存中收集该集群的数据点，并在此子集上运行 NN-descent 以构建该集群的全邻图形。

合并 k-NN 子图形与全局全邻图形

在构建集群的全邻图之后，我们将该 k-NN 子图与全局全邻图合并。为此，我们使用了一个自定义的 CUDA 内核，该内核在不分配额外设备内存的情况下合并了这两个子图。

以这种方式迭代所有集群后，返回全局全邻图作为最终结果。由于该图通常比输入数据集小很多，因此即使输入数据集过大无法容纳，也可以安全地将其复制到 GPU 的内存空间中。

性能提升 

我们评估了使用 cuML UMAP 和新的批量所有邻居图构建方法对性能的影响。

在这些实验中，我们使用了具有 80 GB 内存的 NVIDIA H100 GPU。这些比较与 GPU 版本的 UMAP 进行比较，因此这些加速不是来自于 CPU 到 GPU 的比较，而是对现有 GPU 实现的改进。

图 2 展示了 UMAP 在 cuML 中的总运行时间，并将新的 NN-descent 策略与强力多邻图构建策略进行了比较。对于具有 2000 万个点和 384 个维度的数据集，我们使用 NN-descent 将速度提高了 311 倍，将 UMAP 在 GPU 上的总运行时间从 10 小时缩短到仅 2 分钟！

图 2 采用对数比例，因为加速效果非常显著。

我们还观察到，即使数据集的大小为 150 GB（远超 GPU 中的内存量），现在也可以在 GPU 上运行规模高达 5000 万个点、维度高达 768 的数据集的 UMAP 算法。

通过将数据集划分为 5 个集群，批处理算法实现了这一壮举。相比之下，暴力全邻图形构建算法由于一次性尝试将整个数据集加载到设备上，因此会出现内存不足的情况。

虽然这种新技术可以提高 UMAP 的速度和可扩展性，但我们需要保持质量，以确保低维嵌入可以有效使用。为了衡量质量，我们使用 信任度评分 。信任度评分是介于 0 和 1 之间的分数，表示在运行 UMAP 之前，原始向量的近邻结构与低维 UMAP 嵌入空间中的局部近邻结构相比，保留程度如何。在此度量标准中，越高越好。

图 3 显示，这些显著的加速和优势不会影响 UMAP 嵌入结果的质量。我们可以看到，随着批次数量的增加，可信度评分没有显著变化。

结束语 

我们很高兴与数据科学社区分享这些性能结果。鉴于 UMAP 在各个领域的受欢迎程度，我们相信 RAPIDS cuML 中的这些新功能将显著加速工作流程，并帮助计算科学家发现只有通过在 GPU 上处理大规模数据集才能实现的见解。

要开始使用 cuML 并安装 conda 和 pip 软件包以及现成的 Docker 容器，请参阅 RAPIDS 安装指南 。