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

• The paper presents mlx-snn, a framework that delivers native spiking neural network performance on Apple Silicon by leveraging MLX's unified memory and lazy evaluation.
• It incorporates six neuron models and four surrogate gradient options, enabling efficient temporal unrolling and reduced memory overhead during training.
• Empirical results on MNIST show up to 97.28% accuracy with 2–2.5× faster training and significant memory savings compared to existing frameworks.

mlx-snn: Enabling Spiking Neural Network Research Natively on Apple Silicon via MLX

Introduction and Motivation

mlx-snn establishes the first comprehensive spiking neural network (SNN) framework built explicitly for Apple's MLX array computation library, targeting Apple Silicon hardware. Existing SNN libraries—such as snnTorch, Norse, SpikingJelly, and Lava—are tightly bound to PyTorch or bespoke backends, resulting in suboptimal performance and device support on non-NVIDIA architectures, particularly Apple Silicon. By leveraging MLX's unified memory and composable function transforms (e.g., mx.grad, mx.compile), mlx-snn eliminates common performance bottlenecks associated with cross-device memory transfers and fragmented computation graphs, streamlining SNN experimentation on Apple platforms (2603.03529).

Framework Design and Supported Features

mlx-snn is architected around four main principles: strict MLX-native tensor operations, explicit state handling for compatibility with functional APIs and JIT compilation, broad snnTorch API compatibility, and a research-centric approach that maximizes flexibility and extensibility.

Neuron Model Coverage: The framework includes six neuron models: Leaky Integrate-and-Fire (LIF), Integrate-and-Fire (IF), Izhikevich (with multiple preset regimes), Adaptive LIF (ALIF), Synaptic, and Alpha. Each model is implemented as an MLX module with explicit, externally managed state, which aligns well with backpropagation-through-time (BPTT) and supports composable research pipelines.

Figure 1: Membrane potential traces for all six neuron models with constant input current over 200 timesteps. Red vertical lines indicate spike events.

Surrogate Gradients: Four spike surrogate gradient options are provided, implemented through an STE pattern—crucial given MLX's interim limitations in custom VJP construction. Supported surrogates include the fast sigmoid, arctan, straight-through, and a user-defined option. This STE design ensures numerically stable gradient flow even under non-differentiable spiking nonlinearities.

Figure 2: Surrogate gradient functions. Blue: forward pass (Heaviside step function). Red dashed: backward pass (smooth gradient used for BPTT).

Spike Encoding: The framework offers four encoding methods: rate (Poisson), latency (time-to-first-spike), delta modulation, and an EEG-specific encoder optimized for medical signal integration.

Training Pipeline: mlx-snn features a complete BPTT pipeline with per-timestep unrolling, three loss definitions (including rate coding and membrane-based criteria), and seamless integration with MLX's optimizer suite. Its explicit state-dict pattern enables compatibility with function transforms and facilitates composable training procedures.

Apple Silicon Optimization Via MLX

A central innovation of mlx-snn is its deep integration with core MLX features:

• Unified Memory Architecture: By sidestepping explicit device transfers, mlx-snn exploits Apple Silicon's unified memory to optimize SNN workloads, which often involve rapid, repeated cross-timestep updates, typically a bottleneck in frameworks with discrete host-device memory models.
• Lazy Evaluation and Functional Transforms: mlx-snn capitalizes on MLX’s lazy computation graphs and composable transforms for JIT compilation and automatic differentiation, providing efficient temporal unrolling and batch updates.
• Immutable Array Semantics: All state updates eschew in-place mutation in favor of purely functional assignments, reducing aliasing errors and aligning with the analytical structure of SNN model updates.

Empirical Validation and Results

On MNIST digit classification, mlx-snn demonstrates up to 97.28% accuracy, with a 2.0–2.5× reduction in average epoch training time and 3–10× lower peak GPU memory usage compared to snnTorch on identical Apple M3 Max hardware. Training accelerates further for large batch sizes, illustrating the benefit of MLX’s memory architecture for temporally unrolled models. Across five hyperparameter configurations (varying LIF decay factor, hidden width, learning rate, and batch size), mlx-snn exhibits highly competitive accuracy—consistently within ~1 percentage point of the PyTorch/MPS implementation—with strong resource efficiency.

Figure 3: MNIST training curves for configuration C5 ($\beta{=}0.95$, $h{=}128$). Left: test accuracy vs. epoch. Right: training time per epoch.

The observed small accuracy gap is attributed to the current surrogate gradient implementation (STE), rather than backend or hardware differences. Memory reduction is especially pronounced, benefiting large-scale, multi-layer SNN models.

Implementation Constraints and Future Extensions

mlx-snn v0.2.1 still faces several limitations: mx.compile is not leveraged for the unrolled training loop due to current Pythonic control flow requirements, neuromorphic dataset loaders (e.g., N-MNIST, DVS-Gesture) are pending addition, and empirical validation is so far limited primarily to MNIST. The surrogate gradient workaround for current MLX custom VJP mechanics leads to minor accuracy penalties but is expected to be resolved with framework maturation.

Future goals include a Liquid State Machine backend, full mx.compile integration, broader dataset/applicability, optimized visualization utilities, and eventual v1.0.0 with comprehensive validation and documentation. Critically, the library lowers barriers for SNN research on Apple hardware, improving accessibility for researchers without CUDA-capable devices and enabling more energy-efficient on-device experimentation for neuroscience and neuromorphic computing.

Theoretical and Practical Implications

Practically, mlx-snn empowers AI researchers using Apple laptops and desktops to execute SNN training workflows natively and efficiently, bypassing the CUDA/NVIDIA dependency that dominates the field. Theoretically, it decouples core SNN research from specific hardware, permitting broader reproducibility and encouraging experimentation with novel neuron and surrogate models, temporal coding schemes, and custom losses in a highly modular, compositional Python environment. Direct, unified memory access and functional model definitions will benefit future work at the intersection of deep neuromorphic learning and hardware co-design.

Conclusion

mlx-snn delivers a feature-rich, Apple MLX-native SNN library, supporting a wide array of neuron and gradient models with a clean, research-focused API. Its high efficiency on Apple Silicon positions it as a foundational tool for SNN method development and benchmarking on non-traditional hardware. With ongoing enhancements and planned dataset/dynamic support, mlx-snn will play a central role in broadening the SNN software ecosystem and accelerating cross-platform neuromorphic research.