mlx-vis: GPU-Native Dimensionality Reduction on Apple Silicon

Han Xiao

cs.LGMar 4, 2026

mlx-vis: GPU-Native Dimensionality Reduction on Apple Silicon

Authors: Han Xiao

Organizations: Elastic

Abstract

Dimensionality reduction is a foundational tool for visualizing high-dimensional data, yet its reference implementations span a fragmented stack of CPU-bound Python packages that leaves the Metal GPU on Apple Silicon entirely unused. We present mlx-vis, a library that reimplements seven widely used dimensionality reduction methods and k-nearest neighbor graph construction in pure MLX, with every stage -- from PCA preprocessing through embedding optimization to a circle-splatting renderer -- executing on GPU. On Fashion-MNIST 70K, all seven methods embed in 2.1--4.6 s on an M3 Ultra, achieving 3--13x speedups over CPU baselines while reducing the entire dependency stack to MLX and NumPy. The same pipeline scales to ten million points on a single workstation. Code at https://github.com/hanxiao/mlx-vis

Figures & tables

Panel of Figure 1 — Figure 1: Fashion-MNIST 70K embeddings produced by mlx-vis (MMAE shown for completeness; quantitative results in Sections 3–3.3 cover the seven validated methods).

Panel of Figure 2 — Figure 2: The mlx-vis pipeline. All stages inside the shaded region execute on Metal GPU through MLX.

Method	Time (s)	Mem (GB)	Power (W)
UMAP	$2.60\pm 0.03$	2.5	58
t-SNE	$4.59\pm 0.00$	3.2	74
PaCMAP	$3.97\pm 0.03$	3.0	70
LocalMAP	$4.18\pm 0.02$	3.0	68
TriMap	$2.11\pm 0.03$	2.6	67
DREAMS	$4.63\pm 0.01$	3.2	73

Table 1: Embedding performance on Fashion-MNIST (70K points, 784 dimensions), M3 Ultra. Mean

\pm

std over 5 runs.

Backend	Time (s)	vs. MLX
MLX ( @mx.compile )	$0.42\pm 0.01$	$1.0\times$
PyTorch/MPS (eager)	$1.07\pm 0.04$	$2.6\times$
PyTorch/MPS ( torch.compile )	$0.49\pm 0.03$	$1.17\times$

Table 2: MLX vs. PyTorch/MPS on the identical UMAP optimization loop (Fashion-MNIST 70K, 500 epochs, M3 Ultra, mean

\pm

std over 5 runs). Same graph, init, and gradient; only the array framework differs.

	Trust		Cont
Method	mlx	ref	mlx	ref
UMAP	0.979	0.979	0.982	0.986
t-SNE	0.990	0.989	0.985	0.985
PaCMAP	0.973	0.972	0.984	0.984
TriMap	0.974	0.865	0.987	0.900

Table 3: Embedding quality vs. reference implementations on Fashion-MNIST 10K (

k{=}15

, matched hyperparameters and seed). Trust: trustworthiness; Cont: continuity. mlx (this work) vs. ref (CPU package).

Method	kNN-P	Trust	Cont
UMAP	0.118	0.998	0.999
t-SNE	0.143	0.998	0.999
PaCMAP	0.082	0.998	0.998
LocalMAP	0.122	0.998	0.999
TriMap	0.084	0.998	0.998
DREAMS	0.049	0.998	0.998

Table 4: Embedding quality on Fashion-MNIST 70K (

k{=}10

). All ranks computed from exact pairwise Euclidean distances on GPU. kNN-P:

k

-NN preservation; Trust: trustworthiness; Cont: continuity.

Figure 3: Scaling behavior on GloVe-Twitter (

n

from 70K to 10M points, 200-d) on a single M3 Ultra Mac Studio with 512 GB unified memory [ 2 ] . Top: wall-clock embedding time. Bottom: peak GPU memory. Error bars (mean

\pm

std over 3 runs) are smaller than the markers.

Dataset	Method	Time (s)	Mem (GB)
MNIST	UMAP	$2.42\pm 0.04$	2.4
60K	t-SNE	$4.07\pm 0.01$	3.1
784-d	PaCMAP	$3.28\pm 0.02$	3.0
CIFAR-10	UMAP	$2.54\pm 0.03$	3.0
50K	t-SNE	$5.69\pm 0.02$	3.9
3072-d	PaCMAP	$3.73\pm 0.02$	3.0

Table 5: Cross-domain performance at fixed scale (mean

\pm

std over 3 runs, M3 Ultra). Domains span images (MNIST, CIFAR-10) and pretrained word vectors (GloVe-Twitter).

Dimensionality Reduction

Dimensionality reduction maps high-dimensional data to compact representations for visualization, prediction, compression, and scientific modeling.

See all 145 papers

FastUMAP: Scalable Dimensionality Reduction via Bipartite Landmark Sampling
May 12, 2026Hongmin LiUnsupervised Learning Nyström Method
School of Life Science and Technology, Institute of Science Tokyo · 2-12-1 Ookayama, Meguro-ku, Tokyo 152-8550, Japan 2Department of Computational Biology and Medical Sciences Graduate School of Frontier Sciences, The University of Tokyo
VkSplat: High-Performance 3DGS Training in Vulkan Compute
Apr 30, 2026Jingxiang Chen, Mohamed Ibrahim, Yang Liu3D Gaussian Splatting 3DGS Acceleration
Huawei Canada
Lightweight Gaussian Process Inference in C++ on Metal and CUDA
May 18, 2026Yu-Hsueh FangGPU Kernel Optimization GPU Acceleration
Department of Information Management, National Taiwan University · 2H. Milton Stewart School of Industrial and Systems Engineering, Georgia Institute of Technology

May 12, 2026J/K move · Enter open · S save

FastUMAP: Scalable Dimensionality Reduction via Bipartite Landmark Sampling

Hongmin Li

School of Life Science and Technology, Institute of Science Tokyo · 2-12-1 Ookayama, Meguro-ku, Tokyo 152-8550, Japan 2Department of Computational Biology and Medical Sciences Graduate School of Frontier Sciences, The University of Tokyo

Unsupervised Learning Nyström Method Manifold Learning Dimensionality Reduction

Open paper AskPDF arXiv

Exploratory analysis of high-dimensional data rarely stops at a single embedding. In practice, analysts rerun dimensionality reduction after changing preprocessing, subsets, or hyperparameters, and standard nonlinear methods can quickly become the bottleneck. We introduce FastUMAP (Bipartite Manifold Approximation and Projection), a landmark-based method designed for this repeated-use setting. FastUMAP builds a sparse point-landmark fuzzy graph, computes a Nystrom spectral warm start from the induced landmark affinity, and then refines all sample coordinates with a UMAP-style objective on the bipartite graph. The landmark ratio r = m/n provides a direct way to trade runtime against fidelity. On 9 benchmark datasets spanning 178 to 70,000 samples, FastUMAP has the lowest runtime on 7 datasets in our reported default-implementation comparison on one workstation. On MNIST and Fashion-MNIST (n=70000), it runs in about 4.6 seconds, compared with about 73--75 seconds for Barnes--Hut t-SNE, while reaching 91.4% mean kNN accuracy versus 94.6% for the strongest accuracy baseline. FastUMAP is therefore best viewed as a fast option for repeated exploratory embedding, rather than as a replacement for accuracy-first methods.

Abstract

Dimensionality Reduction

Explore similar work

FastUMAP: Scalable Dimensionality Reduction via Bipartite Landmark Sampling

VkSplat: High-Performance 3DGS Training in Vulkan Compute

Lightweight Gaussian Process Inference in C++ on Metal and CUDA