mlx-snn: Spiking Neural Networks on Apple Silicon via MLX

Dit paper introduceert mlx-snn, de eerste native bibliotheek voor spiking neural networks op Apple Silicon, die dankzij het MLX-framework aanzienlijk snellere training en lager geheugengebruik biedt dan bestaande PyTorch-gebaseerde oplossingen.

De kern

mlx-snn: De "Spitsige" Breinen voor de Apple Mac

Stel je voor dat je brein niet werkt als een strakke, continue stroom van elektriciteit (zoals in de meeste huidige computers), maar als een reeks van flitsen. Denk aan een kamer vol mensen die elkaar aansteken met een flitsapparaat: "Klik!" (een signaal), stilte, dan weer "Klik!". Dit noemen wetenschappers Spiking Neural Networks (SNN's) of "spitsige neurale netwerken". Het is energiezuiniger, sneller en werkt meer zoals een echt biologisch brein.

Maar tot nu toe was er een groot probleem: deze technologie werkte alleen goed op computers met speciale NVIDIA-chips (de "rode" chips). Gebruikers van Apple Macs (met hun eigen "M-chips") zaten met lege handen.

De Oplossing: mlx-snn
Jiahao Qin heeft mlx-snn bedacht. Dit is de eerste bibliotheek die het mogelijk maakt om deze "flitsende" breinen direct te laten draaien op Apple's eigen technologie (MLX). Het is alsof je eindelijk een sleutel hebt gevonden die precies past in het slot van je Mac, zodat je die speciale breinen kunt bouwen zonder een dure externe computer te kopen.

Hoe werkt het? (Met een paar creatieve vergelijkingen)

1. De Bouwstenen (Neuronen)
In een normaal computerprogramma zijn de onderdelen vaak statisch. In mlx-snn zijn er 6 verschillende soorten "flitsers" (neuronen) beschikbaar.

• Vergelijking: Stel je een orkest voor. Je hebt niet alleen trompetten (standaard flitsers), maar ook pauken, fluiten en harpen. Elk instrument (zoals de LIF of Izhikevich modellen) heeft zijn eigen unieke geluid en gedrag. Sommige zijn traag en leeglopen langzaam, andere reageren explosief. De gebruiker kan kiezen welk instrument het beste past bij het probleem.

2. Het Leerproces (De "Goocheltruc" voor de Mac)
Om deze netwerken slim te maken, moeten ze leren uit fouten. In de wiskunde is dit lastig omdat een "flits" (aan of uit) geen halve weg heeft; het is 0 of 1. Dat maakt het moeilijk om te berekenen hoe je het moet verbeteren.

• De truc: mlx-snn gebruikt een slimme "gok" (een zogenaamde surrogate gradient).
• Vergelijking: Stel je voor dat je een blindeman probeert te leren lopen. Hij kan niet zien waar hij gaat (geen gradiënt). De software doet alsof hij wel kan zien door een zachte, wazige bril op te zetten. Hij loopt dan op basis van die wazige indrukken, maar de echte stap blijft een harde "ja/nee". Zodra hij een stap zet, past hij zijn strategie aan op basis van de wazige bril. Dit werkt zo goed dat het Mac-computer niet merkt dat het een trucje is.

3. De Mac-voordeel: De "Gemeenschappelijke Tafel"
Normaal gesproken moet een computer gegevens heen en weer slepen tussen de processor (de chef) en het geheugen (de voorraadkast). Dat kost tijd en energie.

• Vergelijking: Apple's M-chips hebben een unified memory. Stel je voor dat de chef en de voorraadkast aan dezelfde grote tafel zitten. Ze hoeven niet te lopen om ingrediënten te halen; ze reiken ze gewoon door.
• Het resultaat: mlx-snn maakt hier slim gebruik van. Het is 2 tot 2,5 keer sneller en gebruikt 3 tot 10 keer minder geheugen dan de oude methoden op dezelfde Mac. Het is alsof je een raceauto rijdt op een weg waar je nooit hoeft te remmen voor stoplichten.

Wat heeft het opgeleverd?

De auteur heeft het getest met het herkennen van handgeschreven cijfers (MNIST-database).

• Snelheid: Het trainen van het brein ging als een speer.
• Efficiëntie: Het gebruikte zo weinig energie en geheugen dat het zelfs op een MacBook Pro (die je in je rugzak meeneemt) kon draaien, terwijl andere methoden vaak een zware server nodig hebben.
• Resultaat: Het haalde een nauwkeurigheid van 97,28%. Dat is bijna net zo goed als de beste methoden die op zware NVIDIA-chips draaien.

Waarom is dit belangrijk voor jou?

Vroeger moest je een dure, stroomverslindende computer kopen om te experimenteren met deze toekomstige technologie. Met mlx-snn kan elke onderzoeker, student of hobbyist met een Apple Mac (van de MacBook Air tot de Mac Studio) nu meedoen aan de revolutie van spitsige neurale netwerken.

Het is alsof je eindelijk de sleutel hebt gekregen om de "derde generatie" van kunstmatige intelligentie te openen, direct vanuit je eigen laptop, zonder dat je je rugzak hoeft te vervangen door een serverkast.

Kortom: mlx-snn maakt de toekomst van slimme, energiezuinige computers toegankelijk voor iedereen met een Apple-apparaat.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "mlx-snn: Spiking Neural Networks on Apple Silicon via MLX" in het Nederlands.

Probleemstelling

Spiking Neural Networks (SNN's), vaak beschouwd als de "derde generatie" neurale netwerken, bieden unieke voordelen zoals tijdsgecodeerde informatieverwerking, gebeurtenisgestuurde berekening en energie-efficiëntie. Ondanks de snelle groei van SNN-onderzoek, ontbreekt er een native oplossing voor gebruikers van Apple Silicon-hardware (M1, M2, M3-chips).

Bestaande populaire bibliotheken zoals snnTorch, Norse, SpikingJelly en Lava zijn allemaal gebouwd op PyTorch of op maat gemaakte backends die afhankelijk zijn van CUDA (NVIDIA GPU's). Dit dwingt Apple-gebruikers om ofwel hun code aan te passen, ofwel afhankelijk te zijn van inefficiënte CPU-omgevingen of externe cloud-infrastructuur, wat de toegankelijkheid voor SNN-onderzoek in het Apple-ecosysteem beperkt.

Methodologie en Architectuur

De auteurs introduceren mlx-snn, de eerste SNN-bibliotheek die volledig native is gebouwd op MLX, Apple's eigen array-computeframework. De architectuur volgt vier kernontwerpprincipes:

1. Native MLX-implementatie: Alle tensorbewerkingen gebruiken mlx.core. NumPy wordt uitsluitend gebruikt voor I/O.
2. Expliciete staat: Neuronentoestanden worden doorgegeven als Python-dicts in plaats van mutabele objecten. Dit maakt de bibliotheek compatibel met MLX's functionele transformaties en mx.compile.
3. snnTorch-compatibiliteit: De API (klassen, constructoren, forward-pass handtekeningen) is zo ontworpen dat deze sterk lijkt op snnTorch, wat een makkelijke migratie van bestaande code mogelijk maakt.
4. Onderzoeksvriendelijk: Alle componenten kunnen worden uitgebreid of overschreven.

Belangrijke technische componenten:

• Neuronmodellen: De bibliotheek implementeert zes modellen: Leaky Integrate-and-Fire (LIF), Integrate-and-Fire (IF), Izhikevich, Adaptive LIF (ALIF), Synaptic en Alpha.
• Surrogaatgradiënten: Omdat de stapfunctie van een spike geen gradiënt heeft, gebruikt mlx-snn surrogaatgradiënten. Vanwege beperkingen in MLX's huidige mx.custom_function API (die vormproblemen veroorzaakt bij matrixvermenigvuldiging), gebruiken de auteurs een Straight-Through Estimator (STE) patroon met mx.stop_gradient. Dit omzeilt de bug terwijl het numeriek equivalent blijft aan een aangepaste VJP (Vector-Jacobian Product).
• Spike Encoding: Vier methoden worden ondersteund, waaronder Poisson-rate coding, latency coding, delta-modulatie en een specifieke EEG-encoder voor medische signalen.
• Training: Het framework gebruikt Backpropagation Through Time (BPTT) via MLX's mx.grad. De tijdslussen worden "ontrold" (unrolled) om een volledig berekeningstaf te construeren voordat uitvoering plaatsvindt, wat MLX's lazy evaluation ten goede komt.

Belangrijkste Bijdragen

1. Eerste native SNN-bibliotheek voor Apple Silicon: Een volledig functionele SNN-stack die de voordelen van de gedeelde geheugenarchitectuur (Unified Memory) van Apple Silicon benut.
2. Innovatieve STE-patroon: Een oplossing voor surrogaatgradiënten die werkt binnen de huidige beperkingen van MLX, zonder de noodzaak van complexe, nog niet volledig ondersteunde custom VJP-definities.
3. Hoge API-compatibiliteit: Een interface die ontwikkelaars in staat stelt snnTorch-code met minimale wijzigingen te laten draaien op MLX.
4. Empirische validatie: Uitgebreide tests op MNIST-datasets die de prestaties vergelijken met snnTorch op verschillende backends.

Resultaten

De auteurs hebben mlx-snn getest op een Apple M3 Max (36 GB unified memory) en vergeleken met snnTorch op zowel de PyTorch MPS-backend (Apple GPU) als de CPU.

• Snelheid: mlx-snn is 2,0 tot 2,5 keer sneller per trainingsepoche dan snnTorch, zowel op de GPU als de CPU.
• Geheugenefficiëntie: mlx-snn gebruikt 3 tot 10 keer minder GPU-geheugen (Metal) dan snnTorch (MPS). Dit komt door de gedeelde geheugenarchitectuur en het ontbreken van data-overdrachten tussen CPU en GPU.
• Nauwkeurigheid: De beste configuratie van mlx-snn bereikte 97,28% nauwkeurigheid op MNIST. Dit ligt ongeveer 0,7% lager dan de beste resultaten van snnTorch (98,03%). De auteurs attribueren dit kleine verschil aan de gebruikte STE-surrogaatgradiëntmethode in plaats van aan de hardware of het framework zelf.
• Validatie: De dynamiek van alle zes neuronmodellen werd visueel gevalideerd en toont de verwachte biologische gedragingen (bijv. exponentiële stijging bij LIF, kwadratische stijging bij Izhikevich).

Betekenis en Toekomstperspectief

mlx-snn is een mijlpaal voor het SNN-onderzoek in het Apple-ecosysteem. Het democratiseert de toegang tot neuromorfe computing door onderzoekers in staat te stellen complexe SNN-experimenten uit te voeren op standaard MacBook Pro's en Mac Studio's zonder afhankelijk te zijn van NVIDIA GPU's of cloud-diensten.

De Unified Memory Architecture van Apple Silicon is bij uitstek geschikt voor SNN's, die vaak frequente updates van neuronstoestanden vereisen over vele tijdstappen; mlx-snn elimineert hiermee de traditionele bottleneck van data-overdracht.

Beperkingen en Roadmap:
Huidige beperkingen omvatten het ontbreken van mx.compile optimalisatie voor de traininglus (door Python-besturing) en nog niet geïmplementeerde loaders voor neuromorfe datasets (zoals N-MNIST). De toekomstige roadmap omvat versie 0.3.0 met Liquid State Machines (LSM), versie 0.4.0 met geoptimaliseerde compilatie en benchmarks, en een stabiele v1.0.0 release.

Kortom, mlx-snn opent de deur voor een nieuwe generatie energie-efficiënt en tijdsgecodeerd AI-onderzoek op Apple-hardware.