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

TL;DR

mlx-snn is a native Apple Silicon library for spiking neural networks that offers diverse neuron models, efficient training, and high accuracy, filling a gap in existing SNN tools for Apple hardware.

Contribution

It introduces mlx-snn, the first SNN library built directly on Apple's MLX framework, optimized for Apple Silicon with multiple neuron models and training capabilities.

Findings

Achieves up to 97.28% accuracy on MNIST

Provides 2-2.5x faster training than snnTorch

Uses 3-10x less GPU memory

Abstract

We introduce mlx-snn, the first spiking neural network (SNN) library built natively on Apple's MLX framework. As SNN research grows rapidly, all major libraries -- snnTorch, Norse, SpikingJelly, Lava -- target PyTorch or custom backends, leaving Apple Silicon users without a native option. mlx-snn provides six neuron models (LIF, IF, Izhikevich, Adaptive LIF, Synaptic, Alpha), four surrogate gradient functions, four spike encoding methods (including an EEG-specific encoder), and a complete backpropagation-through-time training pipeline. The library leverages MLX's unified memory architecture, lazy evaluation, and composable function transforms (mx.grad, mx.compile) to enable efficient SNN research on Apple Silicon hardware. We validate mlx-snn on MNIST digit classification across five hyperparameter configurations and three backends, achieving up to 97.28% accuracy with 2.0--2.5 times faster training and 3--10 times lower GPU memory than snnTorch on the same M3 Max hardware. mlx-snn is open-source under the MIT license and available on PyPI. https://github.com/D-ST-Sword/mlx-snn