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

mlx-vis is een Python-bibliotheek die zes methoden voor dimensiereductie en een k-nabuur-grafiek-algoritme volledig op Apple Silicon-gpu's implementeert via MLX, en tevens een versnelde renderer biedt voor het genereren van scatterplots en animaties.

De kern

MLX-Vis: De "Supersnelle Foto-App" voor Je Data op Apple-apparaten

Stel je voor dat je een enorme doos hebt vol met duizenden verschillende knikkers. Sommige zijn rood, sommige blauw, sommige groot, sommige klein. Ze liggen allemaal door elkaar in een 3D-doos (of zelfs in 1000 dimensies, wat voor ons onzichtbaar is). Je wilt weten: Welke knikkers lijken op elkaar? Welke groepen vormen ze?

Om dit te zien, moet je die knikkers op een plat stuk papier leggen (2D), zodat je de patronen kunt zien. Dit heet in de tech-wereld dimensiereductie.

Het probleem is dat dit normaal gesproken heel lang duurt, alsof je die knikkers één voor één met je vingers probeert te sorteren. En als je een video wilt maken van hoe die knikkers zich langzaam in groepjes verdelen, duurt het nog langer.

Wat is MLX-Vis?
Deze nieuwe software, genaamd mlx-vis, is als een superkrachtige robotarm die speciaal is gebouwd voor de nieuwe Apple-chips (zoals de M3 in je MacBook of iPhone). In plaats van met je vingers te werken, pakt deze robot de hele doos knikkers, schudt ze razendsnel door elkaar en legt ze perfect in groepjes op je scherm.

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. De "Unieke" Werkplek (Apple Silicon)

Normaal gesproken werken computers met twee aparte werkplekken: een voor het denken (CPU) en een voor het tekenen (GPU). Data moet heen en weer reizen tussen deze plekken, wat tijd kost.

• De Analogie: Stel je voor dat je een chef-kok bent (de CPU) en een assistent (de GPU). De chef moet elke ingrediënt naar de assistent brengen, die het snijdt, en dan terugbrengen. Dat is traag.
• De Oplossing: Apple Silicon heeft unified memory. Het is alsof de chef en de assistent aan dezelfde grote tafel zitten en alle ingrediënten direct binnen handbereik liggen. Er is geen reizen meer. MLX-Vis maakt hier slim gebruik van.

2. De 6 Manieren om te Sorteren

De software bevat zes verschillende "strategieën" om de knikkers te sorteren. Denk hieraan als zes verschillende manieren om een feestje te organiseren:

• Sommige methoden (zoals UMAP en t-SNE) zorgen dat vrienden die dicht bij elkaar staan, ook dicht bij elkaar blijven zitten.
• Andere methoden (zoals PaCMAP) kijken ook naar de hele zaal om te zorgen dat grote groepen niet in de war raken.
• MLX-Vis doet dit allemaal in één pakketje, zonder dat je verschillende gereedschappen hoeft te kopen.

3. De "Magische Projector" (GPU Rendering)

Dit is misschien wel het coolste deel. Normaal gesproken maakt software een plaatje, slaat het op, maakt het volgende plaatje, slaat dat op... en pas daarna zet je ze samen tot een video. Dat is als een fotograaf die elke foto eerst in een album plakt voordat hij de volgende maakt.

• De Analogie: MLX-Vis gebruikt een GPU-native projector. Het is alsof je een magische projector hebt die de beelden direct op het scherm projecteert terwijl de robot de knikkers verplaatst.
• Het gebruikt een techniek die "circle-splatting" heet. In plaats van lijntjes te tekenen, "spat" het de punten als verf op het scherm.
• Het Resultaat: Je kunt een video maken van 800 beelden (een vloeiende animatie van hoe je data zich ordent) in slechts 1,4 seconden. Dat is sneller dan het duurt om een kop koffie te zetten!

4. Hoe snel is het echt?

De makers hebben het getest op een dataset van 70.000 punten (zoals 70.000 foto's van kledingstukken).

• De oude manier (op een normale computer): Het kon 8 tot 60 seconden duren om de knikkers te sorteren.
• De MLX-Vis manier (op een Apple M3 Ultra): Het duurt slechts 2 tot 4 seconden.
• De totale tijd: Van ruwe data tot een prachtige, vloeiende video: minder dan 5 seconden.

Waarom is dit belangrijk voor jou?

Vroeger moest je voor dit soort dingen zware software installeren met veel verschillende onderdelen (zoals nummers, C-extensions, etc.) die vaak vastliepen of traag waren.

• MLX-Vis is licht, snel en werkt puur op de kracht van je eigen Apple-apparaat.
• Het is gratis en openbaar (zoals een openbare bibliotheek).
• Het maakt data-analyse niet alleen sneller, maar ook mooier. Je kunt nu live zien hoe je data zich beweegt, net als een dansend balletje, in plaats van wachten op een statisch plaatje.

Kortom: MLX-Vis is de snelste, meest efficiënte manier om je complexe data op je Apple-apparaat om te zetten in prachtige, vloeiende beelden die je direct kunt begrijpen. Het is alsof je een traag, statisch landkaartje hebt omgezet in een live, 3D-videospelletje van je eigen gegevens.

Technische samenvatting

Probleemstelling

Dimensionaliteitsreductie is een fundamentele techniek in exploratieve data-analyse om hoogdimensionale data om te zetten in 2D-visualisaties. Bestaande methoden zoals t-SNE, UMAP, PaCMAP, TriMap, DREAMS en CNE zijn echter verspreid over verschillende Python-pakketten met heterogene afhankelijkheden (zoals numba, Cython, scipy en sklearn).
De belangrijkste beperkingen van deze bestaande oplossingen zijn:

1. CPU-gebonden uitvoering: Zelfs wanneer GPU-hardware beschikbaar is, voeren deze bibliotheken gradient-updates en zoekopdrachten naar naburige punten (k-NN) uit op de CPU.
2. Onbenutte capaciteit op Apple Silicon: Op Apple Silicon-machines (M-chips) delen CPU en GPU hetzelfde unified memory. De huidige CPU-gebaseerde aanpak laat aanzienlijke rekenkracht onbenut en vereist onnodige data-overdrachten.
3. Fragmentatie: Het ontbreken van een uniforme interface en de afhankelijkheid van zware dependencies maken de integratie complex.

Methodologie

De auteurs introduceren mlx-vis, een Python-bibliotheek die zes dimensionaliteitsreductiemethoden en een k-nearest neighbor (k-NN) algoritme volledig opnieuw implementeert in MLX, Apples array-framework voor Metal GPU.

Architectuur en Implementatie:

• Pure MLX Implementatie: Alle operaties, inclusief PCA-voorverwerking, k-NN-graafconstructie, gradient-gebaseerde optimalisatie en rendering, worden uitgevoerd op de Metal GPU.
• Algoritmen: De bibliotheek ondersteunt UMAP, t-SNE, PaCMAP, TriMap, DREAMS en CNE, evenals NNDescent voor k-NN-zoekopdrachten.
• Geoptimaliseerde Kernen:
  • NNDescent: Geïmplementeerd met GPU-matrixvermenigvuldiging voor afstandsberekening en mx.argpartition voor top-k-selectie zonder volledige sortering.
  • Specifieke Aanpassingen: UMAP gebruikt Gauss-Newton-optimatie in plaats van scipy-curve fitting; t-SNE maakt gebruik van een FFT-versnelde variant voor afstotende krachten (geïnspireerd door FIt-SNE); CNE compiles contrastieve loss-functies tot statische methoden voor operatorfusie.
• GPU-Native Rendering Pipeline: In plaats van matplotlib te gebruiken, implementeert de bibliotheek een "circle-splatting" renderer in MLX.
  • Punten worden omgezet in pixels met lineaire vervaging.
  • Kleurcontributies worden via atomische "scatter-add" bewerkingen op een framebuffer opgeteld.
  • Frames worden direct via ffmpeg met hardware H.264-encoding (VideoToolbox) naar video gepipeerd.
  • Optimalisaties: Dubbele buffering overlapt GPU-rendering met I/O, en mx.async_eval() zorgt voor parallelle verwerking.

API-Design:
De bibliotheek biedt een uniforme API (fit_transform) die NumPy-achtige arrays accepteert. Het gebruikt lazy evaluation en JIT-compilatie via @mx.compile om Python-overhead te elimineren en operaties binnen een epoch te fuseren.

Belangrijkste Bijdragen

1. Eerste volledige GPU-acceleratie voor DR op Apple Silicon: mlx-vis is de eerste bibliotheek die dimensionaliteitsreductie en visualisatie volledig op de GPU uitvoert via MLX, zonder CPU-GPU data-overdrachten.
2. Unificatie en Minimalisme: Het biedt een enkele, dependency-minimale bibliotheek (alleen MLX en NumPy) die zes populaire methoden en een k-NN-algoritme combineert, verwijderend de noodzaak voor scipy, sklearn, numba en Cython.
3. GPU-Native Visualisatie: Een uniek, op de GPU gebaseerd rendering-systeem dat scatterplots en soepele animaties genereert zonder CPU-interventie, wat interactieve exploratie mogelijk maakt.
4. Open Source: De code is beschikbaar onder de Apache 2.0-licentie.

Resultaten

De prestaties zijn getest op een Apple M3 Ultra met 512 GB unified memory, gebruikmakend van de Fashion-MNIST dataset (70.000 punten, 784 dimensies).

• Snelheidswinst (Embedding):
  • t-SNE: 15,5x sneller dan de CPU-versie van openTSNE (3,78s vs 58,62s).
  • TriMap: 6,0x sneller dan trimap.
  • UMAP: 2,6x sneller dan umap-learn.
  • PaCMAP: 3,1x sneller dan pacmap.
  • Alle methoden voltooien de embedding in 2,1 tot 3,8 seconden.
• Rendering en Totale Pipeline:
  • Het genereren van een animatie van 800 frames (1000x1000 resolutie) duurt slechts 1,4 seconden.
  • De totale tijd van ruwe data naar gerenderde video ligt tussen 3,6 en 5,2 seconden.
• Kwaliteit: De kwaliteit van de embedding is vergelijkbaar met de referentie-implementaties, aangezien de doelstellingen en optimalisatieschema's onveranderd zijn gelaten.

Betekenis

mlx-vis markeert een doorbraak in de efficiëntie van data-analyse op Apple Silicon-apparaten. Door de volledige rekenpijplijn (van data-voorbereiding tot visualisatie) op de GPU te verplaatsen, elimineert het de bottleneck van CPU-verwerking en data-overdracht. Dit stelt onderzoekers en data scientists in staat om:

• Interactief te werken met grote datasets (70K+ punten) in real-time.
• Snel iteratieve experimenten uit te voeren met verschillende algoritmen.
• Publicatie-kwaliteit visualisaties en animaties te genereren zonder zware afhankelijkheden.

De bibliotheek demonstreert het potentieel van het MLX-ecosysteem voor het versnellen van machine learning workflows die traditioneel CPU-gebonden waren, en opent de deur voor toekomstige ontwikkelingen in GPU-native data-exploratie.