IbUMAP speeds up and stabilizes UMAP data visualizations
ibUMAP: Coherent and Scalable Field Evaluation for UMAP Optimization
Machine LearningHuman-Computer Interaction
Summary
UMAP is a popular method to visualize complex data in simpler forms but can produce different results each time because of randomness in its calculations. The authors introduce ibUMAP, a new approach that calculates forces between data points all at once rather than in random order, making the results more consistent. They use clever math and computing tricks to do this quickly, even for very large datasets. Their tests show it runs faster and produces more stable visualizations compared to the usual method.
What this means in practice
- •For data engineers: Faster and more stable embedding computations for large-scale data visualization in machine learning pipelines.
- •For gpu software developers: Improve runtime efficiency when implementing embedding algorithms on GPUs for million-scale datasets.
Authors
Bin Chen, Yumeng Xue, Patrick Paetzold, Yunhai Wang, Oliver Deussen
Abstract
UMAP achieves scalable layout optimization through stochastic negative sampling. However, this stochasticity can lead to unstable embeddings across reruns and downstream reuse, as the estimated repulsive forces depend on the ordering of sampling events. We present ibUMAP, a coherent field-based alternative that evaluates attraction and repulsion from a shared embedding snapshot and applies them synchronously. Its degree-weighted repulsive field is motivated by the conditional expectation of negative sampling for a fixed embedding and represented by three scalar moments, which are evaluated efficiently on CPUs and GPUs using an interpolation-based FFT scheme. This formulation avoids explicit all-pairs computations while inducing optimization dynamics that differ from those of standard online UMAP. Controlled experiments show that synchrony and kernel capping alter the local-global fidelity trade-off, whereas FFT evaluation produces small average changes in final quality. End-to-end benchmarks show median speedups of 3.29x unseeded and 5.79x seeded over umap-learn on CPU, and 1.44x over cuML on million-scale datasets under unseeded GPU execution. These gains accompany greater run-to-run stability and measurable fidelity trade-offs.