ibUMAP: Coherent and Scalable Field Evaluation for UMAP Optimization

AI summary

Overview

Research area: Dimensionality reduction and neighbor-embedding optimization (UMAP), with a focus on numerical algorithms, GPU/CPU acceleration, and reproducibility of embeddings.

Technical level: Advanced. The paper involves negative-sampling estimators, N-body force decomposition, FFT-based particle-mesh convolution, and controlled optimizer ablations.

Scope: The paper proposes ibUMAP, a modified UMAP layout optimizer that replaces stochastic negative-sampling repulsion with a synchronously applied, degree-weighted repulsive field evaluated via interpolation-based FFT on CPUs and GPUs, and benchmarks its speed, fidelity, and repeatability against umap-learn, cuML, and TorchDR.

What This Paper Is About

UMAP optimizes embeddings using stochastic negative sampling, where repulsive interactions are drawn at random and applied in place, so the resulting layout depends on the ordering of sampling events and varies across reruns. The authors ask whether repulsion can instead be evaluated as a deterministic field from a single frozen embedding snapshot and applied synchronously with attraction, and whether that change can be made efficient enough to beat standard UMAP implementations in speed while keeping quality competitive. They build ibUMAP to answer this, keeping UMAP's graph construction and initialization untouched and modifying only the layout-optimization stage.

Key Contributions

1. Formulation and implementation. The authors formalize a frozen-state, degree-weighted repulsive operator under explicit negative-sampling and update conventions, derive a UMAP-specific three-moment decomposition, and implement synchronous two-dimensional FFT-based optimization on CPUs and GPUs with stated complexity and repeatability conditions.

2. Mechanism decomposition. Using a fixed-graph, fixed-initialization framework across eight variants on 59 datasets, they isolate how synchrony, the repulsion formulation, scalar kernel capping, FFT evaluation, and update safeguards each change final embedding quality.

3. Performance and repeatability evaluation. End-to-end benchmarks report median CPU speedups of 3.29 times unseeded and 5.79 times seeded over umap-learn, plus unseeded GPU gains over cuML on million-scale datasets, alongside quantified fidelity trade-offs and fixed-seed repeatability costs.

4. Downstream case study. In a BRAQUE-derived UMAP–HDBSCAN workflow on 56,962 cells, they trace embedding variation to HDBSCAN assignment changes and measure the CPU cost of making embeddings repeatable.

Main Findings

• Synchrony costs local fidelity. Making updates synchronous lowers trustworthiness (TW) on all 59 datasets and neighborhood preservation (NP) on 57 in the controlled comparison (variant A to B).

• Event expectation partially recovers the loss. Replacing random negative draws with their frozen-state expectation under uniform sampling recovers about 47% of the mean NP loss (B to C).

• Kernel capping trades neighbor overlap for continuity. The raw degree-weighted field lowers both global scores on most datasets (C to D); adding the scalar kernel cap reverses this on 54 and 53 datasets and raises mean continuity (C) from 0.9470 to 0.9813, but lowers NP on 54 datasets (D to E).

• FFT evaluation changes average quality little. FFT evaluation leaves all five paired distributions centered at zero, changing each mean by at most 0.0005 in the local scores and 0.0016 in distance Spearman correlation (ρ_D) (E to F).

• Safeguards improve global fidelity. Repulsion-norm clipping and attraction damping improve C, RTA, and ρ_D on most datasets, raising mean ρ_D by 0.0228 at a mean NP change of −0.0002 (F to G).

• Production ibUMAP is close to, but not equivalent to, UMAP. It improves NP over variant G on 58/59 datasets but lowers both global scores on 52/59 (G to H). Relative to UMAP, it matches C but lowers mean TW, NP, RTA, and ρ_D by 0.0051, 0.0209, 0.0037, and 0.0179 (A to H).

• CPU speedups. Median speedups over umap-learn are 3.29 times unseeded and 5.79 times seeded; both profiles win on the same 60/66 datasets and lose only below 1,000 observations.

• GPU performance is size-dependent. Unseeded ibUMAP trails cuML at the suite median (speed ratio 0.875) but wins on all 11 datasets above 10^5 observations, reaching a median 1.44 times speedup across six million-scale datasets. Against TorchDR it wins on 62/66 datasets with a median 1.53 times speedup.

• Seeded execution is nearly free on CPU. Both CPU implementations produce bitwise-identical embeddings across five seeded runs on all 66 datasets. Median seeded/unseeded time ratios are 0.961 for ibUMAP and 1.664 for umap-learn. On six million-scale GPU datasets, ibUMAP keeps only a 1.05 times median speedup when seeded, with shared cuML graph construction dominating.

• Fidelity medians in the end-to-end benchmark. On the common 66 datasets under unseeded execution, CPU umap-learn scores TW 0.9901, C 0.9942, NP 0.3699, RTA 0.7249, ρ_D 0.6072, while ibUMAP scores 0.9853, 0.9950, 0.3436, 0.7023, and 0.5493. On GPU, cuML scores 0.9892, 0.9950, 0.3635, 0.7349, 0.6295; TorchDR scores 0.9873, 0.9943, 0.3768, 0.7647, 0.6706; ibUMAP scores 0.9851, 0.9950, 0.3407, 0.7217, 0.5897.

• Better run-to-run stability. ibUMAP improves 15-NN overlap on 58/66 datasets versus umap-learn, 67/71 versus cuML, and 66/66 versus TorchDR; distance correlation also improves on most datasets.

• Geometric stability does not guarantee repeatable clustering. Across the ten unordered pairs of unseeded UMAP runs in the BRAQUE-derived case, median pairwise-distance Spearman correlation is 0.937, yet median adjusted Rand index (ARI) is 0.611 and median changed-assignment fraction is 20.80%. The first two scheduled runs have ρ = 0.935, ARI 0.615, and 19.78% changed assignments, including 3.49% of cells switching noise status; cluster counts across five runs range from 37 to 45. Even ibUMAP's higher unseeded distance correlation of 0.999 coexists with 22.01% changed assignments.

• Cheaper repeatable fits. With seed 42 fixed, umap-learn's median embedding fit rises from 15.44 s to 49.06 s (a 3.18 times slowdown) because seeding forces serial optimization, while ibUMAP keeps eight jobs and takes 7.10 s (ratio 1.001). Seeded ibUMAP is 6.91 times faster than seeded UMAP, or 6.42 times including the HDBSCAN fit (7.74 s versus 49.70 s).

Methodology in Plain English

The authors keep everything about UMAP that builds the neighborhood graph and initializes coordinates, and only replace the part that moves points around. Instead of drawing random negative samples one at a time and immediately nudging points (which makes results depend on the order of events), they freeze the current layout, compute the total repulsive force on every point from that single snapshot, and then move all points at once along with the attractive forces.

Computing every pair's repulsion directly would be far too slow. They exploit the fact that the repulsive vector for a point can be written using only three scalar quantities summed over all other points, and that these sums are convolutions with one shared kernel. They place each point's three "charges" onto a regular grid, convolve with the kernel using an FFT, and read the results back at each point's location — a particle-to-mesh, FFT, mesh-to-particle pattern adapted from FIt-SNE. Because there are no random draws, they also add safeguards: capping the kernel value, clipping the repulsive step norm to 4α_t, and damping attraction at high-degree hubs based on the 99th percentile of weighted degree.

To isolate what each change does, they compare eight optimizer variants on 59 datasets with identical graphs and initializations, running 200 epochs with 15 neighbors and min_dist of 0.1 and averaging over three optimizer seeds. Then they run end-to-end benchmarks on 71 datasets comparing against umap-learn, cuML, and TorchDR under unseeded and seeded profiles, timing five runs per configuration after warmup, and measuring fidelity with trustworthiness, continuity, neighborhood preservation, random triplet accuracy, and distance Spearman correlation. Finally they examine a BRAQUE-derived single-cell workflow where embeddings feed HDBSCAN clustering.

Why This Matters

Impact on research. The paper shows that a stochastic optimizer's expected behavior can be turned into a deterministic field operator, and it separates the effects of synchrony, field weighting, approximation, and safeguards that are usually bundled together. It also demonstrates that run-to-run geometric similarity is a weak guarantee for downstream partition stability, which matters for how embedding quality is reported.

Real-world applications:

• Single-cell analysis pipelines, such as the BRAQUE-derived UMAP–HDBSCAN workflow tested here, where repeated embeddings feed clustering and annotation.
• Large-scale data visualization of datasets above 10^5 to millions of observations, where GPU speedups over cuML were observed on all 11 datasets in that size range.
• Reproducible analysis pipelines that require fixed-seed, bitwise-repeatable embeddings across reruns without losing parallel CPU execution.
• Parameter exploration workflows, where re-fitting an embedding with a fixed seed returns the same layout and partition every time, lowering the cost of each refit from 49.06 s to 7.10 s in the case studied.

Industry relevance. ibUMAP retains a drop-in relationship with existing UMAP pipelines by changing only the layout-optimization stage, so it can be substituted where determinism, run-to-run stability, or CPU throughput matter more than matching UMAP's fidelity medians. The fidelity trade-offs reported (lower TW, NP, RTA, and ρ_D medians than umap-learn unseeded) mean adoption depends on downstream requirements, as the authors state.

Future Directions

• Extending the field formulation beyond two-dimensional output, which currently increases mesh and convolution costs and needs a separate accuracy and scaling study.
• Testing robustness to initializations other than the single fixed spectral initialization used in the mechanism study.
• Determining whether dynamical predictions from attraction-shape theory for pairwise UMAP transfer to synchronous force aggregation.
• Characterizing the ibFFT approximation by field-level error against the direct sum rather than only by its effect on final embeddings, and asking whether other sampled interaction kernels admit a useful expected-field realization.

Target Audience

Researchers and practitioners in dimensionality reduction, scientific visualization, and high-performance computing who work with UMAP at scale or in reproducible pipelines; readers interested in the numerical analysis of stochastic optimizers, FFT-based N-body force evaluation, and the relationship between embedding geometry and downstream clustering stability.

Authors’ 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.