mlx-vis: GPU-Accelerated Dimensionality Reduction and Visualization on Apple Silicon

Abstract: mlx-vis is a Python library that implements six dimensionality reduction methods and a k-nearest neighbor graph algorithm entirely in MLX, Apple's array framework for Apple Silicon. The library provides UMAP, t-SNE, PaCMAP, TriMap, DREAMS, CNE, and NNDescent, all executing on Metal GPU through a unified fit_transform interface. Beyond embedding computation, mlx-vis includes a GPU-accelerated circle-splatting renderer that produces scatter plots and smooth animations without matplotlib, composing frames via scatter-add alpha blending on GPU and piping them to hardware H.264 encoding. On Fashion-MNIST with 70,000 points, all methods complete embedding in 2.1-3.8 seconds and render 800-frame animations in 1.4 seconds on an M3 Ultra, with the full pipeline from raw data to rendered video finishing in 3.6-5.2 seconds. The library depends only on MLX and NumPy, is released under the Apache 2.0 license, and is available at https://github.com/hanxiao/mlx-vis.

• The paper demonstrates a unified GPU-native implementation of six dimensionality reduction algorithms on Apple Silicon using MLX, achieving significant speedups and maintaining embedding quality.
• It employs innovative optimizations such as lazy evaluation and JIT compilation to fuse computational phases and eliminate CPU-bound bottlenecks.
• Empirical tests on Fashion-MNIST show up to 15.5× speedup over CPU implementations, enabling rapid generation of high-resolution, animated visualizations.

Unified GPU-Accelerated Dimensionality Reduction and Visualization on Apple Silicon

Overview and Motivation

mlx-vis presents a consolidated Python library for dimensionality reduction and visualization, purpose-built for Apple Silicon ecosystems via MLX, the Metal-native array programming framework. The library implements six major dimensionality reduction algorithms—UMAP, t-SNE, PaCMAP, TriMap, DREAMS, and CNE—alongside fast $k$-nearest neighbor graph construction (NNDescent), all fused and executed entirely on the Metal GPU. This integration addresses two longstanding deficits in the dimensionality reduction ecosystem: fragmented reference implementations with heterogeneous dependencies and CPU-bound operation, even when substantial GPU resources are present on Apple Silicon. mlx-vis achieves operational unification and performance gains by refactoring all algorithmic phases—including $k$-NN search, iterative embedding optimization, and visualization rendering—into pure MLX, leveraging unified memory and JIT kernel fusion.

Algorithmic Implementation and Pipeline Architecture

The design philosophy emphasizes modularity layered atop MLX’s capabilities. Each dimensionality reduction algorithm is encapsulated as a class, ingesting hyperparameters and exposing a fit_transform API. The interface is standardized: users pass an $n \times d$ data matrix and receive an $n \times 2$ embedding. The NNDescent implementation supplies approximate $k$-NN graphs via iterative refinement, with distance computations leveraging matrix multiplication primitives and top-$k$ selection performed with partial sorting on GPU.

Critical MLX-specific optimizations are leveraged throughout:

• Lazy evaluation confines GPU kernel dispatch to epoch boundaries, enabling operator fusion across optimization phases.
• JIT compilation via @mx.compile is applied to per-algorithm “hot loops” (SGD steps, repulsion kernels, triplet-loss phases), reducing Python overhead and maximizing throughput.

The visualization subsystem is GPU-native. Embeddings are rendered by circle-splatting scatter-add accumulation and alpha blending on the Metal GPU, bypassing matplotlib entirely. For animation, per-epoch snapshots are pipelined directly to hardware H.264 encoding via ffmpeg, with asynchronous evaluation and double-buffering minimizing latency.

Empirical Performance and Embedding Fidelity

mlx-vis exhibits high-performance on canonical tasks, dramatically outpacing reference CPU-bound implementations. On Fashion-MNIST (70,000 points, $d=784$), embedding time for all six methods consistently falls in the 2.1–3.8 second range, while GPU-native rendering of 800-frame videos completes in under 1.5 seconds. End-to-end pipeline times (data ingestion, embedding, rendering, animation) range from 3.6–5.2 seconds—achieved on an Apple M3 Ultra with unified memory—establishing a new practical baseline for dimensionality reduction tasks at scale.

The numerical benchmarks are summarized as follows:

• UMAP: 2.6$\times$ GPU speedup (mlx-vis vs. umap-learn)
• t-SNE: 15.5$\times$ GPU speedup (mlx-vis vs. openTSNE)
• PaCMAP: 3.1$\times$ GPU speedup (mlx-vis vs. pacmap)
• TriMap: 6.0$\times$ GPU speedup (mlx-vis vs. trimap)

Quality of embeddings generated by mlx-vis is maintained, as the library reproduces every published algorithmic objective and optimization schedule faithfully, evidenced by the visualization outputs. The circle-splatting renderer introduces no artifacts, yielding sharp cluster delineation and clear continuity across all methods.

Figure 1: Fashion-MNIST 70K embeddings produced by the six methods in mlx-vis, rendered by the GPU circle-splatting pipeline.

Practical and Theoretical Implications

The consolidation and GPU-acceleration of dimensionality reduction tools directly enhance exploratory data analysis workflows on Apple Silicon. Elimination of CPU-GPU data transfers, dependency minimization, and acceleration of both embedding generation and visualization address bottlenecks prevalent in real-world tasks (e.g., single-cell transcriptomics, imaging, high-dimensional clustering). The homogeneity of API conventions facilitates rapid prototyping and deployment, particularly in settings where interactive or animated embedding visualization is desired.

Theoretical implications pertain to the architectural viability of MLX/Metal for scientific computation tasks beyond neural network training, demonstrating that high-level algorithmic translation (neighbor embedding, triplet optimization, contrastive losses) can be efficiently vectorized and compiled on domain-specialized hardware. The pipeline further suggests that diffusion- and matrix-exponential-based methods (e.g., PHATE, StarMAP) could be supported, contingent on future MLX abstractions for spectral operations.

Outlook and Future Directions

mlx-vis establishes an efficient, unified baseline for dimensionality reduction and visualization on Apple Silicon, with implications for adoption in environments such as macOS-powered research clusters and interactive notebook interfaces. Future developments may target:

• Expansion to additional methods (diffusion, MDS, centroid-driven), as MLX adds primitives for matrix exponentiation and spectral analysis.
• Integration with MLX-based LLM inference and diffusion workflows, driving end-to-end GPU-native exploratory analysis.
• Extension of the rendering pipeline to higher-order visualization (e.g., density fields, interactive selection, batch-animated embeddings).

Broader theoretical directions include systematic benchmarking against CUDA/TensorFlow/PyTorch pipelines, analysis of memory scaling with unified GPU architectures, and investigation of fused optimization loss landscapes for multi-modal embeddings.

Conclusion

mlx-vis delivers a dependency-minimal, GPU-native framework for dimensionality reduction and visualization on Apple Silicon, spanning UMAP, t-SNE, PaCMAP, TriMap, DREAMS, CNE, and NNDescent. Its MLX-based architecture achieves substantial acceleration, maintains embedding quality, and introduces a performant pipeline for high-resolution, publication-quality visualization and animation. The streamlined workflow, modular API, and architectural choices provide a foundation for high-throughput, exploratory data analysis and visualization, with clear avenues for expansion as the MLX and Metal ecosystems evolve.