Approximate UMAP (aUMAP)

• The method introduces a k-NN regression-based projection that replaces iterative optimization, achieving sub-millisecond per-point embeddings on CPU.
• It leverages a sketch-based distributed approach, enabling aggregation of local embeddings for scalable visualization of massive, geo-distributed datasets.
• Deep parametric variants, like NUMAP, integrate neural spectral embedding to preserve global structure while providing instant out-of-sample generalization.

Approximate UMAP (aUMAP) refers to a class of methods for accelerating, scaling, or generalizing Uniform Manifold Approximation and Projection (UMAP), particularly in settings where standard UMAP’s computation or projection stages become prohibitively expensive. These approaches include efficient non-parametric out-of-sample mapping (aUMAP), sketch-based distributed variants for massive data (SnS-UMAP), and deep-learning–powered parametric approximations (NUMAP), each addressing distinct challenges in high-dimensional data visualization and embedding.

1. Core Concepts and Problem Motivation

UMAP is widely used for nonlinear dimension reduction and visualization based on the construction of a fuzzy simplicial set over high-dimensional data followed by cross-entropy optimization in a low-dimensional latent space. While UMAP provides both local and global structure preservation, it suffers from three principal bottlenecks in applied settings:

• Projection latency: Standard UMAP requires re-running local optimization for each new test point (“out-of-sample”) because no closed-form mapping exists.
• Computational and communication cost: For large or geo-distributed datasets, the quadratic scaling of pairwise similarities and global optimization renders standard workflows impractical.
• Generalizability: Vanilla UMAP, being non-parametric, does not provide an efficient mapping for new or streaming data.

Approximate UMAP methodologies tackle these technical challenges via algorithmic acceleration, distributed data summarization, and parametric mapping construction.

2. Algorithmic Foundations of Approximate UMAP

2.1 aUMAP: Fast Nonparametric Out-of-Sample Extension

Approximate UMAP (aUMAP), as introduced for online projection in high-dimensional data streams, maintains the original UMAP training regime but augments it with an auxiliary $k$-NN regression model. During inference, aUMAP projects new points $x^*$ by identifying their $k$ approximate nearest neighbors in input space and computing a weighted barycenter of their low-dimensional embeddings:

$$w_{i_j} = \frac{1/d_{i_j}}{\sum_{\ell=1}^k 1/d_{i_\ell}}, \qquad u^* = \sum_{j=1}^k w_{i_j} u_{i_j}$$

where $d_{i_j}$ is the Euclidean distance from $x^*$ to its $j$th nearest neighbor, and $u_{i_j}$ the neighbor’s low-dim coordinate. This replacement of gradient-based local optimization with a closed-form barycentric mapping enables sub-millisecond per-point projections using only CPU resources and minimal implementation overhead (Wassenaar et al., 2024).

2.2 Sketch-and-Scale (SnS) UMAP for Massive Geo-Distributed Data

The Sketch-and-Scale (SnS-UMAP) framework approximates UMAP embeddings by leveraging the linearity and compression properties of the Count Sketch data structure. Each distributed edge node summarizes its local $k$-NN UMAP neighborhood graph into a $O(s \times \ell)$ sketch array using hash and sign functions for streaming updates:

$$S[\ell, h_\ell(e)] \leftarrow S[\ell, h_\ell(e)] + g_\ell(e) \cdot \mu_{q|p}$$

where $\mu_{q|p}$ encodes the standard UMAP membership weight, and $(h_\ell, g_\ell)$ denote hash and sign maps for row $\ell$ (Wei et al., 2020).

At the central server, sketches from $m$ nodes are aggregated and heavy-hitters (largest edges by recovered weight) are extracted to build the global sparse graph for the final UMAP SGD embedding. This allows communication, memory, and computational costs to scale sublinearly relative to dataset size, making SnS-UMAP feasible for datasets with tens to hundreds of millions of points distributed across multiple datacenters.

2.3 Deep Parametric Approximate UMAP (NUMAP/GrEASE)

NUMAP extends UMAP with a generalizable, parametric architecture utilizing GrEASE—a neural-network-based spectral embedding framework. GrEASE minimizes the Rayleigh quotient loss over minibatches to learn an approximate leading eigenspace of the UMAP affinity Laplacian, supports fast eigenvector separation via postprocessing, and outputs a continuous $F_\theta(x)$ mapping for OOSE (Ben-Ari et al., 20 Jan 2025).

A shallow correction network $G_\phi$ is trained (with residual connection) using UMAP’s cross-entropy on top of the parametric spectral features, providing both global structure consistency and fine local embedding structure.

3. Mathematical Formulation and Pipeline Details

3.1 aUMAP Pseudocode and Formulation

Training:

• Fit UMAP on training set $X \to Y$ (low-dim).
• Build $k$-NN index on $X$.

Projection:

indices, dists = kNN.query(x_star, k) inv = 1.0 / dists w = inv / sum(inv) u_star = sum_j w[j] * Y[indices[j]] return u_star

No backpropagation or gradient descent is required for projection.

3.2 SnS-UMAP Pipeline

Edge node:

• For each local $x_p$, compute $k$-NN $N_k(p)$, calculate $\mu_{q|p}$, and update count sketch $S^{(i)}$.
• Transmit $S^{(i)}$ to central server.

Server:

• Compute $S = \sum_{i=1}^m S^{(i)}$.
• Extract top-$H$ edges by estimated median.
• Build global sparse graph $G$.
• Run vanilla UMAP SGD on $G$ for embedding $Y$.

3.3 NUMAP/GrEASE Workflow

• Affinity and Laplacian construction per minibatch.
• Neural net $F_\theta$ trained to minimize Rayleigh quotient loss; orthonormalization via QR.
• Postprocess on batch means to recover eigenvectors individually.
• Secondary network $G_\phi$ (residual) trained with UMAP loss on top of $F_\theta$ outputs.

4. Computational Complexity and Empirical Benchmarks

Method	Training Time	Projection Time (per point)	Storage/Deployment
Standard UMAP	$O(n\log n D + mnk)$	$O(\log n D + t k)$	Must store full embedding
aUMAP	Same as UMAP	$O(\log n D + k d)$	Adds $k$-NN index
Parametric UMAP	$O(E n p)$	$O(p)$	Neural net, larger infra
SnS-UMAP	$O(n_i (T_{ANN} + k\ell))$ per node	—	Distributed, communication-light
NUMAP (GrEASE)	$O(n m)$	$O(p)$	Fast OOSE, minibatch only

Empirically, aUMAP achieves sub-millisecond projection latency on CPU, matching UMAP’s embedding space with mean deviation $\ll 1\sigma$ for clusters. Training costs are identical to UMAP, outperforming parametric UMAP in speed unless GPU is available, where pUMAP closes the projection gap (Wassenaar et al., 2024). SnS-UMAP demonstrated one order-of-magnitude wall-time speedup and two orders of magnitude reduction in memory/communication, embedding 50–110M point datasets in under two hours with 99% cluster fidelity (Wei et al., 2020). NUMAP matches or exceeds parametric UMAP on global alignment metrics (e.g., Grassmann score) and yields instant out-of-sample embeddings (Ben-Ari et al., 20 Jan 2025).

5. Use Cases, Parameter Choices, and Practical Guidelines

aUMAP is highly suited for streaming applications such as BCI feedback loops, rapid prototyping of neural latent space trajectories, or any scenario demanding sub-millisecond, low-latency pointwise embedding updates. Recommended settings include using $k=15$ (UMAP default), identical distance metrics, and two or three output dimensions. For high-dimensional datasets ($D \gg 1000$) or large $n$, increase k-NN tree leaf size to balance query speed and neighbor fidelity.

SnS-UMAP is designed for large, geo-distributed, or privacy-constrained datasets, requiring partitioned edge nodes and moderate sketch parameters: depth $\ell=3$–$5$, width $s=5k$–$20k$, and heavy-hitter thresholds proportional to $nk$. NUMAP is appropriate in scenarios demanding global structure preservation and out-of-sample consistency, particularly when parametric deployment is desired.

6. Trade-offs and Limitations

aUMAP prioritizes projection speed and CPU efficiency at the expense of local embedding precision for outlier or boundary points, as no neighborhood optimization is performed for new samples. SnS-UMAP introduces minor approximation (count-sketch loss/variance) controlled by sketch size but achieves near-lossless results for practical configurations. NUMAP’s generalizability incurs a moderate initial training overhead due to neural network optimization, but once trained, offers instant embedding and stable OOSE.

A plausible implication is that use-case priority—latency, scalability, or structure preservation—should dictate choice among standard UMAP, aUMAP, and other approximate variants.

7. Extensions and Context within Embedding Research

Approximate UMAP variants represent a continuum of scalability and generalizability improvements within the nonlinear embedding literature. aUMAP serves as a near plug-and-play drop-in for deployment pipelines requiring low infrastructure changes. SnS-UMAP generalizes the compress-and-aggregate paradigm for large-scale manifold learning on distributed systems, extending ideas from streaming algorithms such as Count Sketch. NUMAP/GrEASE situates UMAP within the regime of neural spectral embeddings, connecting UMAP’s loss functional with recent advances in deep manifold learning and parametric out-of-sample generalization.

These approaches collectively extend the applicability of UMAP from static, in-memory datasets to streaming data, distributed systems, and online/real-time visualization scenarios (Wei et al., 2020, Wassenaar et al., 2024, Ben-Ari et al., 20 Jan 2025).