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 擴展為大規模數據集。更好的是，它是默認啟用的，因此您無需對代碼進行任何更改即可從中受益！
| 矩陣大小 | 使用強力運行 UMAP | 使用 nn-descent 運行 UMAP |
| 1M x 960 | 214.4 秒 | 9.9 秒 (加速 21.6 倍) |
| 8 米 x 384 米 | 2191.3 秒 | 34.0 秒 (加速 54.4 倍) |
| 10 米 x 96 米 | 2170.8 秒 | 53.4 秒 (加速 40.6 倍) |
| 20 米 x 384 米 | 38350.7 秒 | 122.9 (加速 312 倍) |
| 59M x 768 | 錯誤：內存不足 | 575.1 |
表 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 UMAPdata = generate_data()# running default. Runs with NN Descent if data has more than 50K pointsumap = 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 hostumap = 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 clustersumap = 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 安裝指南 。
?
