mlx-vis: GPU-Native Dimensionality Reduction

• mlx-vis is a GPU-native toolkit for dimensionality reduction that implements six state-of-the-art methods entirely on Apple Silicon.
• It provides a unified fit_transform API that leverages MLX arrays and JIT compilation to eliminate CPU bottlenecks in graph construction and optimization.
• The library features a GPU-accelerated renderer that produces static scatter plots and smooth animations with hardware H.264 encoding for efficient visualization.

Searching arXiv for the specified paper and closely related method papers to ground the article in current literature. mlx-vis is a Python library for dimensionality reduction and visualization on Apple Silicon that implements six dimensionality reduction methods—UMAP, t-SNE, PaCMAP, TriMap, DREAMS, and CNE—and the NNDescent k-nearest neighbor graph builder entirely in MLX, Apple’s array framework, so that they execute on the Metal GPU through a unified fit_transform interface (Xiao, 4 Mar 2026). In addition to embedding computation, it includes a GPU-accelerated circle-splatting renderer for static scatter plots and smooth animations without matplotlib, with frames piped directly to hardware H.264 encoding. The stated design objective is an end-to-end GPU-native pipeline that eliminates dependency sprawl and CPU bottlenecks by running PCA preprocessing, k-NN graph construction, dimensionality-reduction optimization, and rendering on GPU with unified memory (Xiao, 4 Mar 2026).

1. Scope, positioning, and software model

mlx-vis is described as a unified, GPU-native toolkit for dimensionality reduction and visualization on Apple Silicon. Its distinguishing characteristic is that the entire pipeline is implemented in MLX: there is no numba, Cython, or scipy in the core execution path, and the pipeline runs on the Metal GPU via MLX arrays with lazy evaluation and JIT compilation through @mx.compile (Xiao, 4 Mar 2026).

The library exposes a single API pattern across methods. Given $X \in \mathbb{R}^{n \times d}$, each estimator returns an embedding $Y \in \mathbb{R}^{n \times 2}$. The public interface includes:

$Y \in \mathbb{R}^{n \times 2}$35

Each class exposes fit_transform(X), where the input X is an MLX or NumPy array convertible to MLX, with shape (n, d), and the output is an n×2 embedding. An optional epoch_callback is invoked each iteration and receives the current Y as a NumPy array, enabling animation capture without altering the inner loop (Xiao, 4 Mar 2026).

The package depends only on MLX and NumPy, is released under the Apache 2.0 license, and is available at https://github.com/hanxiao/mlx-vis. GPU acceleration requires Apple Silicon with Metal GPU support. CPU-only fallback is explicitly not the focus of the library (Xiao, 4 Mar 2026).

2. End-to-end architecture on MLX and Metal

The architecture is organized around an end-to-end GPU execution model. All data structures—input data $X$, graph structures, embeddings $Y$, and rendered frames—remain in unified memory, so there are no explicit CPU↔GPU copies. The implementation relies on MLX lazy evaluation, where each epoch ends with mx.eval(), allowing MLX to fuse intra-epoch operations into larger GPU kernels (Xiao, 4 Mar 2026).

Hot loops are JIT-compiled with @mx.compile. The paper identifies compiled kernels for UMAP’s SGD step, t-SNE’s repulsive kernel, PaCMAP phase updates, and CNE losses. The intended effect is to reduce Python overhead while optimizing memory access and operator fusion (Xiao, 4 Mar 2026).

Several implementation techniques recur across the library. Pairwise distances use the GEMM identity

$\|a-b\|^2 = \|a\|^2 + \|b\|^2 - 2a^\top b,$

with precomputed norms and batched matrix multiplications. Top-$k$ selection uses mx.argpartition rather than full sorting. Random sampling for negatives and triplets is vectorized on GPU. Numerical stability is handled with saturating or soft kernels, small $\epsilon$ terms in denominators, and early stopping in NNDescent with $\delta = 0.001$ (Xiao, 4 Mar 2026).

The package layout reflects this decomposition. Methods reside in mlx_vis/_method/, NNDescent in mlx_vis/_nndescent/, and the renderer in render.py, with plotting and animation entry points in plot.py. The public API consists of the six dimensionality-reduction classes, NNDescent, and two visualization functions, scatter_gpu and animate_gpu (Xiao, 4 Mar 2026).

3. Dimensionality-reduction methods and mathematical objectives

mlx-vis reimplements six state-of-the-art dimensionality-reduction methods with objectives and schedules that follow the original methods. Benchmarks used matched iteration counts relative to reference implementations, and the paper states that embeddings visually match those references (Xiao, 4 Mar 2026).

Method	Core formulation	Optimization notes
t-SNE	KL divergence between $P$ and $Q$	FFT-accelerated $Y \in \mathbb{R}^{n \times 2}$0 repulsion, optional early exaggeration
UMAP	Cross-entropy between fuzzy graphs	Negative sampling, compiled SGD kernel
PaCMAP	Weighted near, mid-near, and further pair loss	Three-phase schedule $Y \in \mathbb{R}^{n \times 2}$1
TriMap	Weighted triplet loss	Triplet sampling from GPU k-NN
DREAMS	t-SNE KL plus PCA anchor penalty	Dual-stage annealing of $Y \in \mathbb{R}^{n \times 2}$2
CNE	InfoNCE-style contrastive neighbor loss	Compiled contrastive terms

For all methods, initialization is configurable, typically PCA or random. The paper also states that UMAP may use spectral initialization and that learning-rate schedules and iteration counts mirror the reference implementations (Xiao, 4 Mar 2026).

t-SNE

With $Y \in \mathbb{R}^{n \times 2}$3, $Y \in \mathbb{R}^{n \times 2}$4, and $Y \in \mathbb{R}^{n \times 2}$5, the high-dimensional conditional probabilities are

$Y \in \mathbb{R}^{n \times 2}$6

where $Y \in \mathbb{R}^{n \times 2}$7 is found by binary search to match the user-specified perplexity. Symmetrization yields

$Y \in \mathbb{R}^{n \times 2}$8

The low-dimensional affinities use the Student-$Y \in \mathbb{R}^{n \times 2}$9 kernel with one degree of freedom,

$X$0

The objective is

$X$1

with gradient

$X$2

mlx-vis performs k-NN graph construction and perplexity search in MLX, supports an FFT-accelerated $X$3 repulsive-force variant following FIt-SNE, and uses standard t-SNE scheduling with learning-rate, momentum, and optional early exaggeration (Xiao, 4 Mar 2026).

UMAP

For each point $X$4, UMAP finds $X$5 neighbors via NNDescent and defines local connectivity $X$6 and bandwidth $X$7 through a monotonic search so that smoothed k-NN distances meet a target entropy:

$X$8

The fuzzy simplicial set is symmetrized as

$X$9

In the low-dimensional space, with $Y$0,

$Y$1

mlx-vis fits the parameters $Y$2 and $Y$3 by Gauss–Newton on GPU rather than scipy curve_fit. The objective is the cross-entropy

$Y$4

Optimization uses stochastic updates with negative sampling, and the attractive and repulsive updates are fused in a compiled SGD kernel. The neighborhood size is n_neighbors, and min_dist controls the curve shape via $Y$5 and $Y$6 (Xiao, 4 Mar 2026).

PaCMAP

PaCMAP uses three pair sets: near pairs $Y$7, mid-near pairs $Y$8, and further pairs $Y$9. A commonly used loss is

$\|a-b\|^2 = \|a\|^2 + \|b\|^2 - 2a^\top b,$0

where $\|a-b\|^2 = \|a\|^2 + \|b\|^2 - 2a^\top b,$1 and $\|a-b\|^2 = \|a\|^2 + \|b\|^2 - 2a^\top b,$2 is a small constant for numerical stability. The three-phase schedule adjusts $\|a-b\|^2 = \|a\|^2 + \|b\|^2 - 2a^\top b,$3 across iterations to emphasize near structure, then mid-near structure, and finally repulsion. Benchmarks used the standard 100/100/250 schedule (Xiao, 4 Mar 2026).

TriMap

TriMap uses relative comparisons over sampled triplets $\|a-b\|^2 = \|a\|^2 + \|b\|^2 - 2a^\top b,$4, where $\|a-b\|^2 = \|a\|^2 + \|b\|^2 - 2a^\top b,$5 is closer than $\|a-b\|^2 = \|a\|^2 + \|b\|^2 - 2a^\top b,$6 to $\|a-b\|^2 = \|a\|^2 + \|b\|^2 - 2a^\top b,$7 in high-dimensional space. With

$\|a-b\|^2 = \|a\|^2 + \|b\|^2 - 2a^\top b,$8

the triplet loss is

$\|a-b\|^2 = \|a\|^2 + \|b\|^2 - 2a^\top b,$9

Triplets are generated from GPU k-NN by choosing inliers from the k-NN graph and outliers from non-neighbors, using a mix of random and far points. Loss and gradients are batched and compiled (Xiao, 4 Mar 2026).

DREAMS

DREAMS is described as a hybrid that combines t-SNE’s local-neighborhood preservation with a global regularizer based on PCA. Let $k$0 be the two-dimensional PCA projection. A standard formulation is

$k$1

The weight $k$2 is decreased during optimization in a dual-stage schedule: it is large initially to place clusters globally, then reduced to allow local refinement. The additional gradient term is $k$3 (Xiao, 4 Mar 2026).

CNE

CNE reframes neighbor embedding as a contrastive learning problem. For each anchor $k$4, neighbors $k$5 are positives and sampled non-neighbors $k$6 are negatives. The similarity $k$7 is a monotone decreasing function of distance, with typical choices

$k$8

The per-anchor InfoNCE-style loss is

$k$9

and the total loss is $\epsilon$0. The attractive and repulsive components are extracted and compiled for operator fusion (Xiao, 4 Mar 2026).

The paper’s textual summary of results further differentiates the methods qualitatively: UMAP and t-SNE emphasize local cluster purity; PaCMAP and TriMap preserve more global relationships; DREAMS places global layout closer to PCA while maintaining local separation; and CNE provides contrastive neighbor structure akin to UMAP and t-SNE, with competitive clustering. This is a summary of observed behavior rather than a new metric (Xiao, 4 Mar 2026).

4. Approximate k-NN graph construction with NNDescent

NNDescent is the common graph-building primitive used across the library. The algorithm initializes each point with $\epsilon$1 random neighbors and iteratively refines the graph by examining neighbors-of-neighbors and optionally random samples. Distances are computed on GPU, merged with current neighbors, and truncated to the top $\epsilon$2 entries using mx.argpartition. Convergence is determined by the update rate, defined as the fraction of neighbor entries changed, and iteration stops when this rate falls below $\epsilon$3 (Xiao, 4 Mar 2026).

The high-level pseudocode is:

• Input: $\epsilon$4, $\epsilon$5
• Precompute rowwise norms $\epsilon$6
• Initialize $\epsilon$7 with $\epsilon$8 random neighbors and their distances
• Repeat:
  • for each $\epsilon$9 in parallel, form a candidate set from $\delta = 0.001$0 and neighbors-of-neighbors
  • compute

  $\delta = 0.001$1

  for candidates $\delta = 0.001$2, batched on GPU - update $\delta = 0.001$3 to the top-$\delta = 0.001$4 indices in the candidate set - compute update_rate

• stop when update_rate < δ

The paper characterizes the practical complexity as near-linear scaling in $\delta = 0.001$5 for fixed $\delta = 0.001$6, with candidate evaluation benefiting from batched distance computations and good GPU occupancy (Xiao, 4 Mar 2026). Parallelization is pointwise, neighbor indices and distances are stored in contiguous MLX arrays, and the termination metric is also computed on GPU.

A plausible implication is that the library’s uniform performance across six embedding methods depends substantially on sharing a common GPU-resident graph-construction substrate rather than maintaining separate neighbor-search implementations for each estimator.

5. GPU rendering, animation, and video encoding

mlx-vis includes a circle-splatting renderer implemented entirely on GPU. For each point $\delta = 0.001$7, the renderer rasterizes a filled disk of radius $\delta = 0.001$8 pixels around the mapped screen coordinate. For each pixel offset with radius $\delta = 0.001$9, the falloff weight is

$P$0

The premultiplied per-pixel color contribution is

$P$1

and contributions are accumulated into an RGBA framebuffer using atomic scatter-add: $Y \in \mathbb{R}^{n \times 2}$36 A final pass normalizes RGB by accumulated alpha, masking division by zero, and composites over the background (Xiao, 4 Mar 2026).

The alpha blending model is standard source-over destination:

$P$2

Because accumulation uses premultiplied alpha, the paper states that this yields correct order-independent blending for an opaque background. Renderer antialiasing is achieved through the linear falloff $P$3, which provides soft edges without branching (Xiao, 4 Mar 2026).

Animation is driven by snapshots captured through epoch_callback. Three implementation optimizations are identified:

1. Hold frames reuse buffers when positions are unchanged.
2. mx.async_eval() overlaps GPU rendering of frame $P$4 with I/O of frame $P$5 through double buffering.
3. Constants such as splat kernel offsets are converted to MLX tensors once and reused.

Frames are piped to ffmpeg using the macOS hardware encoder h264_videotoolbox. Resolution, frame rate, codec, and output filename are configurable (Xiao, 4 Mar 2026).

The public visualization entry points are illustrated by scatter_gpu for static rendering and animate_gpu for videos. A typical scatter call accepts the embedding, colors, width, height, point radius, and alpha; a typical animation call additionally accepts fps, codec="h264_videotoolbox", and an output path (Xiao, 4 Mar 2026).

6. Performance characteristics, usage patterns, and limitations

The benchmark dataset is Fashion-MNIST with $P$6 and $P$7, evaluated on an Apple M3 Ultra with 512 GB unified memory. Reported results are means $P$8 standard deviation over 10 runs, with iteration counts matched to reference implementations (Xiao, 4 Mar 2026).

Task	Result
UMAP embedding, 500 epochs	$P$9 s vs $Q$0 s, $Q$1 speedup
t-SNE embedding, 500 iterations	$Q$2 s vs $Q$3 s, $Q$4 speedup
PaCMAP embedding, 450 iterations	$Q$5 s vs $Q$6 s, $Q$7 speedup
TriMap embedding, 500 iterations	$Q$8 s vs $Q$9 s, $Y \in \mathbb{R}^{n \times 2}$00 speedup
DREAMS embedding, 500 iterations	$Y \in \mathbb{R}^{n \times 2}$01 s
CNE embedding, 500 iterations	$Y \in \mathbb{R}^{n \times 2}$02 s
Rendering, 800 frames at $Y \in \mathbb{R}^{n \times 2}$03	$Y \in \mathbb{R}^{n \times 2}$04 s
End-to-end pipeline	$Y \in \mathbb{R}^{n \times 2}$05–$Y \in \mathbb{R}^{n \times 2}$06 s from raw $Y \in \mathbb{R}^{n \times 2}$07 to rendered MP4

The abstract summarizes the same benchmark at a coarser level, stating that on Fashion-MNIST with 70,000 points, all methods complete embedding in $Y \in \mathbb{R}^{n \times 2}$08–$Y \in \mathbb{R}^{n \times 2}$09 seconds and that 800-frame animations render in $Y \in \mathbb{R}^{n \times 2}$10 seconds on an M3 Ultra, with the full pipeline finishing in $Y \in \mathbb{R}^{n \times 2}$11–$Y \in \mathbb{R}^{n \times 2}$12 seconds (Xiao, 4 Mar 2026).

Typical hyperparameters are also specified. For t-SNE, perplexity $Y \in \mathbb{R}^{n \times 2}$13–$Y \in \mathbb{R}^{n \times 2}$14 is suggested for $Y \in \mathbb{R}^{n \times 2}$15–$Y \in \mathbb{R}^{n \times 2}$16, with $Y \in \mathbb{R}^{n \times 2}$17–$Y \in \mathbb{R}^{n \times 2}$18 iterations and optional early exaggeration such as $Y \in \mathbb{R}^{n \times 2}$19 for the first $Y \in \mathbb{R}^{n \times 2}$20–$Y \in \mathbb{R}^{n \times 2}$21 iterations; FFT repulsion is recommended for $Y \in \mathbb{R}^{n \times 2}$22. For UMAP, n_neighbors in $Y \in \mathbb{R}^{n \times 2}$23 and min_dist in $Y \in \mathbb{R}^{n \times 2}$24 are described as typical, with negative_sample_rate in $Y \in \mathbb{R}^{n \times 2}$25. For PaCMAP, the standard 100/100/250 schedule and modest n_neighbors such as $Y \in \mathbb{R}^{n \times 2}$26–$Y \in \mathbb{R}^{n \times 2}$27 are recommended. For TriMap, around $Y \in \mathbb{R}^{n \times 2}$28 inlier neighbors and $Y \in \mathbb{R}^{n \times 2}$29–$Y \in \mathbb{R}^{n \times 2}$30 outliers per anchor are described as effective. For CNE, $Y \in \mathbb{R}^{n \times 2}$31–$Y \in \mathbb{R}^{n \times 2}$32 negatives per anchor is described as a common balance between speed and quality. PCA initialization is noted to often stabilize and accelerate convergence across methods (Xiao, 4 Mar 2026).

The limitations are explicit. Apple Silicon with Metal GPU is required for GPU acceleration. Very large $Y \in \mathbb{R}^{n \times 2}$33 may become memory-bound even with unified memory, and graph structures scale with $Y \in \mathbb{R}^{n \times 2}$34. The renderer targets 2D embeddings; although 3D could be supported by projection, that is not the library’s target. Reproducibility is not strict because approximate k-NN uses random initialization, negative sampling is stochastic, GPU parallel reductions and atomic scatter-adds may be nondeterministic, and asynchronous overlaps can introduce run-to-run variance. Fixing random seeds where exposed and avoiding asynchronous overlaps is suggested for more reproducible runs, though minor numerical drift due to GPU execution order should still be expected (Xiao, 4 Mar 2026).

The paper also identifies future directions: interactive viewers using the same GPU renderer, native 3D embeddings and camera control, streaming or online dimensionality reduction, and additional methods in the same MLX style, including LargeVis and FIt-SNE variants. This suggests that mlx-vis is designed less as a single-method implementation and more as a common execution substrate for high-throughput embedding and visualization workflows on Apple Silicon (Xiao, 4 Mar 2026).