Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Jiahao Qin

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

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

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
Paper Structure (35 sections, 6 equations, 3 figures, 3 tables)

This paper contains 35 sections, 6 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. SNN Libraries.
  4. Apple MLX.
  5. Architecture and Design
  6. Design Principles
  7. Neuron Models
  8. Leaky Integrate-and-Fire (LIF).
  9. Integrate-and-Fire (IF).
  10. Izhikevich.
  11. Adaptive LIF (ALIF).
  12. Synaptic.
  13. Alpha.
  14. Surrogate Gradients
  15. STE pattern.
  16. ...and 20 more sections

Figures (3)

  • Figure 1: MNIST training curves for configuration C5 ($\beta{=}0.95$, $h{=}128$). Left: test accuracy vs. epoch. Right: training time per epoch.
  • Figure 2: Membrane potential traces for all six neuron models with constant input current over 200 timesteps. Red vertical lines indicate spike events.
  • Figure 3: Surrogate gradient functions. Blue: forward pass (Heaviside step function). Red dashed: backward pass (smooth gradient used for BPTT).