umap-metal 0.1.0

UMAP accelerated by Apple MLX on Mac MPS (Metal Performance Shaders)

umap-metal

UMAP accelerated by Apple MLX on Mac MPS (Metal Performance Shaders)

umap-metal is a drop-in replacement for umap-learn that runs all major compute phases on Apple Silicon GPUs via MLX. It is designed to be interface-compatible with umap-learn.UMAP so existing code requires only a one-line import change.

---

Requirements

Hardware	Mac with Apple Silicon (M1 / M2 / M3 / M4 family)
macOS	12.3 Ventura or later
Python	3.10+

---

Installation

pip install umap-metal

Note mlx is an Apple-only package and will only install on macOS with Apple Silicon.

---

Quick Start

The API mirrors umap-learn:

from umap_metal import UMAP
import numpy as np

X = np.random.standard_normal((10_000, 50)).astype("float32")

reducer = UMAP(n_neighbors=15, n_components=2, random_state=42)
embedding = reducer.fit_transform(X)   # shape (10_000, 2)

Drop-in replacement

# Before
from umap import UMAP

# After — nothing else changes
from umap_metal import UMAP

scikit-learn Pipeline compatibility

UMAP extends sklearn.base.BaseEstimator and TransformerMixin:

from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from umap_metal import UMAP

pipe = Pipeline([
    ("scaler", StandardScaler()),
    ("umap",   UMAP(n_neighbors=15, n_components=2)),
])
embedding = pipe.fit_transform(X)

---

Parameters

Parameter	Default	Description
n_neighbors	15	Size of local neighborhood
n_components	2	Embedding dimensionality
metric	"euclidean"	"euclidean" or "cosine"
n_epochs	None	SGD epochs (auto: 200 for large, 500 for small)
learning_rate	1.0	Initial learning rate
init	"spectral"	"spectral" or "random"
min_dist	0.1	Minimum distance between embedded points
spread	1.0	Scale of the embedding
n_negative_samples	5	Negative samples per positive edge
random_state	None	int, RandomState, or None
verbose	False	Print per-step timing
chunk_size	4096	Chunk size for exact MLX kNN (small n)
force_exact_knn	False	Force O(n²) exact kNN regardless of n
use_mlx_layout	True	Use MLX GPU layout optimization
nndescent_iters	5	NN-Descent refinement iterations (large n)
nndescent_trees	10	RP-trees for NN-Descent initialization (large n)

---

How It Works

umap-metal accelerates all three compute-heavy phases of UMAP on Apple MPS:

Phase 1 — k-Nearest Neighbor Search

n	Method	Complexity
≤ 15,000	MLX exact GEMM — chunked ‖xᵢ − xⱼ‖² via xᵢxⱼᵀ	O(n²·d)
> 15,000	MLX NN-Descent — RP-Forest init + neighbor-of-neighbor refinement	O(n·k²·d·iters)

The MLX NN-Descent builds binary random-projection trees (vectorized with np.einsum + np.bincount) to generate high-recall initial candidates, then refines with forward and reverse neighbor propagation. No Numba JIT is used, eliminating the 4–6 s cold-start penalty of pynndescent.

Phase 2 — Fuzzy Graph Construction

Smooth kNN distances (σ, ρ per point via binary search) and membership strength computation remain on CPU with Numba JIT.

Phase 3 — Layout Optimization (SGD)

A node-centric reformulation replaces the edge-centric SGD scatter pattern:

# Edge-centric (requires scatter_add — GPU-unfriendly)
for edge (i, j):
    force_i += attract(i, j)
    force_j -= attract(i, j)

# Node-centric (pure gather + reduce — GPU-optimal)
neighbors = emb[knn_indices[i, :]]          # (N, k, D) gather
attr_force = reduce(attract(i, neighbors))   # (N, D) no scatter needed
rep_force  = reduce(repel(i, random_neg))    # (N, D)
emb[i]    += lr * (attr_force + rep_force)

This achieves a consistent 4–4.5× speedup on the layout step alone, and requires no atomic operations or scatter kernels.

---

Benchmark

Tested on Apple M2 Max (96 GB unified memory), n_neighbors=15, dim=50, epochs=200.
umap-learn timings are after one Numba warm-up call (subsequent-call performance).

       n    umap-metal    umap-learn    speedup    kNN method
─────────────────────────────────────────────────────────────
   5,000        0.61 s        5.48 s      9.0×    MLX exact
  10,000        1.32 s        2.95 s      2.2×    MLX exact
  20,000        2.75 s        5.94 s      2.2×    MLX NN-Descent
  50,000        6.37 s       15.33 s      2.4×    MLX NN-Descent
 100,000       12.80 s       32.42 s      2.5×    MLX NN-Descent

Step-level breakdown for n = 50,000:

  kNN (MLX NN-Descent)   4.33 s   68%
  Layout (MLX SGD)       1.42 s   22%
  Graph construction     0.45 s    7%
  Spectral init          0.17 s    3%
  ─────────────────────────────────
  Total                  6.37 s

The 9× speedup at n=5,000 reflects umap-learn's first-call Numba compilation cost.
From n=10,000 onward, comparisons are against warm (post-compilation) umap-learn.

kNN quality (recall@15 vs exact brute-force)

Method	n=5,000	n=10,000
MLX exact	100%	100%
MLX NN-Descent	77%	68%
pynndescent (umap-learn)	81%	73%

---

Verbose Output

reducer = UMAP(n_neighbors=15, n_epochs=200, verbose=True)
reducer.fit_transform(X)

[umap-metal] n=50000  k=15  epochs=200  kNN=MLX NN-Descent
  kNN (MLX NN-Descent): 4.33s
  Graph construction: 0.45s
  Spectral init: 0.17s
  a=1.5769  b=0.8951  layout=MLX
  epoch 0/200  lr=1.0000
  ...
  Layout optimization: 1.42s

---

Project Structure

umap_metal/
├── __init__.py          # exposes UMAP
├── umap_metal.py        # UMAP class (sklearn-compatible)
├── knn.py               # dispatch: MLX exact or MLX NN-Descent
├── knn_nndescent.py     # RP-Forest init + NN-Descent (MLX, no Numba)
├── graph.py             # fuzzy simplicial set construction (Numba)
├── spectral.py          # spectral embedding initialization (scipy)
├── layout.py            # SGD layout, Numba CPU fallback
└── layout_mlx.py        # node-centric SGD layout on MLX GPU

---

Limitations

• macOS / Apple Silicon only — MLX does not run on Linux or Windows.
• transform() (out-of-sample projection) is not yet implemented.
• Approximate kNN (n > 15,000) achieves ~70–77% recall; embeddings are visually equivalent to umap-learn but may differ in fine structure.
• Very high-dimensional data (d > 200) may reduce the GPU efficiency advantage.

---

License

MIT

---

Citation

If you use this software, please also cite the original UMAP paper:

McInnes, L., Healy, J., & Melville, J. (2018).
UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction.
arXiv:1802.03426.