This article details a significant performance and scalability enhancement to the UMAP (Uniform Manifold Approximation and Projection) algorithm within the RAPIDS cuML 24.10 library. The authors identify the primary bottleneck in the previous GPU implementation: the brute-force, all-neighbors graph construction phase, which scaled quadratically and required the entire dataset to fit in GPU memory. To overcome this, they implemented a novel solution combining a fast, approximate nearest neighbor algorithm (nn-descent) with a new batching mechanism. This batching process uses k-means clustering to partition datasets larger than GPU memory, processing them in chunks from host RAM. The result is a dramatic acceleration of the UMAP algorithm—reducing a 10-hour process to 2 minutes in one case—and the ability to run on massive datasets (e.g., 150 GB) that previously caused out-of-memory errors, all without a significant loss in embedding quality. These improvements are enabled by default, making high-performance, large-scale dimension reduction more accessible to the data science community.
Accelerating and Scaling UMAP in RAPIDS cuML
The Problem with UMAP: The UMAP dimension reduction algorithm's performance was severely limited by its all-neighbors graph construction step. This step used a brute-force method, which was computationally expensive (quadratic scaling) and memory-intensive, requiring the entire dataset to fit on the GPU. For datasets with 1 million or more vectors, this step consumed over 99% of the total runtime.
Solution 1: Approximate Nearest Neighbors (ANN): To accelerate the process, the brute-force method was replaced with nn-descent, a fast approximate nearest neighbor algorithm from the RAPIDS cuVS library. This algorithm trades perfect accuracy for massive speed gains by cleverly reducing the number of distance calculations needed.
Solution 2: Scalable Batch Processing: To handle datasets larger than the GPU's memory, a new batching algorithm was introduced.
It first uses k-means clustering on a subsample of the data to define cluster centroids.
The full dataset, stored in host RAM, is then partitioned into overlapping batches based on each point's closest centroids.
Each batch is moved to the GPU to build a local k-NN subgraph, which is then merged into a global graph. This allows for processing datasets that are far larger than the available VRAM.
Massive Performance Gains: The new implementation provides dramatic speedups over the previous GPU version.
For a 20M x 384 dataset, the runtime was reduced from over 10 hours to just 2 minutes (a 312x speedup).
The algorithm can now process datasets that were previously too large, successfully running on a 59M x 768 (153 GB) dataset on an 80 GB GPU.
Quality is Maintained: Despite the use of approximation and batching, the quality of the final UMAP embedding is not sacrificed. Trustworthiness scores, which measure the preservation of local data structure, remain high and stable even when the data is split across many batches.
Easy to Use: These new features are enabled by default in RAPIDS cuML version 24.10. Users can benefit from the speed and scalability improvements without any code changes.
Advanced Control: For users needing more control, new parameters have been introduced:
build_algo: Explicitly select "nn_descent" or "brute_force_knn".
data_on_host=True: An option to run the batching algorithm on data stored in host RAM.
build_kwds: A dictionary to pass specific parameters to the nn-descent algorithm, such as nnd_n_clusters to control the number of batches.