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

TL;DR

mlx-vis is a GPU-accelerated library for dimensionality reduction and visualization on Apple Silicon, enabling fast embedding and rendering of high-dimensional data with minimal dependencies.

Contribution

It introduces a GPU-based implementation of eight dimensionality reduction methods and a built-in GPU renderer for Apple Silicon, achieving significant speed improvements.

Findings

Seven of eight methods embed data in 2.0-4.7 seconds.

800-frame animations are rendered in 1.4 seconds.

Library depends only on MLX and NumPy.

Abstract

mlx-vis implements eight dimensionality reduction methods -- UMAP, t-SNE, PaCMAP, LocalMAP, TriMap, DREAMS, CNE, MMAE -- and NNDescent k-NN graph construction entirely in MLX for Apple Silicon Metal GPU. A built-in GPU renderer produces scatter plots and smooth animations via hardware H.264 encoding. On Fashion-MNIST (70K points, M3 Ultra), seven of eight methods embed in 2.0-4.7s and render 800-frame animations in 1.4s. The library depends only on MLX and NumPy and is available at https://github.com/hanxiao/mlx-vis.