---
title: 'mlx-snn: SNN Library for Apple Silicon'
url: https://www.emergentmind.com/topics/mlx-snn
type: topic
---

# mlx-snn: SNN Library for Apple Silicon

mlx-snn is a native spiking neural network (SNN) library designed for Apple Silicon and built on the Apple MLX array-computation stack. It addresses the gap in the SNN ecosystem where all major frameworks (such as snnTorch, Norse, SpikingJelly, and Lava) target PyTorch or custom backends, which are not optimized for the unified memory and compute characteristics of Apple hardware. mlx-snn offers a research-focused API with six neuron models, comprehensive surrogate gradient support, multiple spike encoding schemes (including EEG-specific), and a backpropagation-through-time (BPTT) training pipeline. Leveraging MLX features such as unified CPU/GPU memory, deferred/lazy evaluation, and composable functional transforms, mlx-snn provides efficient, low-overhead SNN experimentation and training on MacBook-class devices. Empirical evaluation on MNIST classification demonstrates up to 97.28% accuracy, with 2.0–2.5× speedup and 3–10× lower GPU memory footprint compared to snnTorch on the same Apple M3 Max hardware [2603.03529].

## 1. System Architecture and MLX Integration

mlx-snn is architected around the Apple MLX computation stack, which provides:

- **Unified Memory**: Both CPU and GPU share physical RAM, obviating the need for device placement operations such as `.to(device)`; tensors can fluidly move between compute units without explicit code intervention.
- **Lazy Evaluation**: MLX builds deferred computation graphs that are only executed upon explicit calls (e.g., `mx.eval(...)`). This allows global optimization over all temporal steps within an SNN sequence, amortizing kernel launches and reducing peak memory usage.
- **Composable Transforms**: Essential functionals such as `mx.grad` (automatic differentiation), `mx.compile` (JIT compilation), `mx.vmap`, `mx.jvp`, and `mx.linearize` enable transformation and vectorization of pure functions, enhancing both development flexibility and runtime performance in SNN contexts.

mlx-snn exploits these features as follows:

- All neuron and encoder operations utilize `mlx.core` tensors.
- Neuron states are implemented as immutable MLX arrays, stored in Python dictionaries (e.g., `{'mem':..., 'syn':..., ...}`), facilitating compatibility with pure-function transforms and JIT compilation.
- The training loop constructs a global computation graph covering all T time steps and executes in a single shot. This minimizes repeated kernel launches and peak memory requirements.
- Surrogate gradient computation follows a straight-through estimator (STE) pattern, employing `mx.stop_gradient` to circumvent MLX custom function VJP shape constraints while still supporting end-to-end differentiation.
- The library components are modular, grouped into `mlxsnn.neurons` (six neuron types), `mlxsnn.surrogates` (four surrogate gradients), `mlxsnn.encoding` (four encoding schemes), and `mlxsnn.training` (BPTT, built-in loss functions, and utility APIs).
- An API compatibility layer provides close alignment to snnTorch naming, constructor signatures, state management, and `forward` call patterns.

## 2. Neuron Models and Discrete-Time Update Equations

mlx-snn implements six neuron models with direct, explicit updates in discrete time ($t = 0, 1, \ldots, T-1$):

- **Leaky Integrate-and-Fire (LIF)**:
  
  $$
  U[t+1] = \beta \, U[t] + X[t+1] - S[t] \, V_{\mathrm{thr}}, \quad S[t]=\Theta(U[t] - V_{\mathrm{thr}})
  $$
  with decay $\beta$, threshold $V_{\mathrm{thr}}$.

- **Integrate-and-Fire (IF)**:
  
  $$
  U[t+1] = U[t] + X[t+1] - S[t]\, V_{\mathrm{thr}}
  $$
  (special case: $\beta = 1$).

- **Izhikevich (2D dynamical system)**:

  $$
  \frac{dv}{dt} = 0.04 v^2 + 5v + 140 - u + I, \quad \frac{du}{dt} = a(bv-u)
  $$
  Discretized via Euler; four presets (RS, IB, CH, FS).

- **Adaptive LIF (ALIF)**:
  
  $$
  U[t+1] = \beta U[t] + X[t+1] - S[t] V_{\mathrm{thr}}, \quad
  A[t+1] = \rho A[t] + S[t], \quad
  V_{\mathrm{eff}}[t] = V_{\mathrm{thr}} + b A[t]
  $$
  Threshold adapts with spike count.

- **Synaptic (filtered input current)**:
  
  $$
  I_{\mathrm{syn}}[t+1] = \alpha I_{\mathrm{syn}}[t] + X[t+1], \quad
  U[t+1] = \beta U[t] + I_{\mathrm{syn}}[t+1] - S[t] V_{\mathrm{thr}}
  $$

- **Alpha (dual-exponential synapse)**:
  
  $$
  I_{\mathrm{exc}}[t+1] = \alpha I_{\mathrm{exc}}[t] + X[t+1], \quad
  I_{\mathrm{inh}}[t+1] = \alpha I_{\mathrm{inh}}[t] + I_{\mathrm{exc}}[t+1], \quad
  U[t+1] = \beta U[t] + I_{\mathrm{inh}}[t+1] - S[t] V_{\mathrm{thr}}
  $$
  Models rapid rise and slower decay dynamics.

All updates are implemented using MLX arrays and permit batched, vectorized computation for scalable simulation.

## 3. Surrogate Gradient Techniques and Training Dynamics

Spiking neurons implement hard threshold nonlinearity, resulting in zero gradients almost everywhere; thus, training via backpropagation-through-time (BPTT) requires surrogate gradients. In mlx-snn, the straight-through estimator pattern is employed:

$$
\mathrm{output} = \mathrm{stop\_grad}(\Theta(x) - \tilde{\sigma}(x)) + \tilde{\sigma}(x)
$$

This yields identity forward output ($\Theta(x)$), with gradient replaced by a smooth surrogate $\tilde{\sigma}(x)$ in backpropagation.

Supported surrogates include:

- **Fast sigmoid**:
  $$
  \tilde{\sigma}(x) = \frac{kx}{2(1 + k|x|)} + \frac{1}{2}
  $$
- **Arctan**:
  $$
  \tilde{\sigma}(x) = \frac{1}{\pi}\arctan(\alpha x) + \frac{1}{2}
  $$
- **Straight-through (piecewise linear)**:
  $$
  \tilde{\sigma}(x) = \mathrm{clip}(s x + 0.5, 0, 1)
  $$
- **Custom**: Arbitrary user-supplied smooth function, wrapped for STE relaxation.

The BPTT loop builds a full temporal graph, allowing gradients to be propagated through all intermediate neuron states. The workflow is fully compatible with MLX's autograd and lazy evaluation, making it possible to combine SNNs with arbitrary differentiable models.

## 4. Spike Encoding Schemes

mlx-snn provides four spike encoding methods for mapping continuous or time-series data into spike trains:

- **Rate Coding (Poisson)**: Each scalar feature $x \in [0,1]$ is interpreted as the probability of emitting a spike at every time step, $S[t] \sim \mathrm{Bernoulli}(x)$.
- **Latency Coding (Time-to-First-Spike)**: $x$ is mapped to a spike time $t_{\mathrm{spike}}$, either linearly or exponentially; a spike is emitted only once per sequence.
- **Delta Modulation**: Emits a spike when $|x[t] - x[t-1]| > \theta$, suitable for temporal and streaming data.
- **EEG Encoder**: Specialized encoding for multi-lead EEG data; supports rate, delta, and threshold-crossing on a per-channel basis, producing tensors of shape $[T, B, C]$.

Encoders are designed for both prototypical vision experiments (e.g., MNIST) and bio-signal pipelines, and integrate seamlessly with subsequent network layers.

## 5. Empirical Benchmarking and Performance Analysis

mlx-snn was benchmarked on MNIST digit classification using an Apple M3 Max MacBook Pro (36 GB unified memory), running MLX 0.29.3 and Python 3.9.18. Baselines included snnTorch 0.9.4 with PyTorch 2.8.0 on both Apple MPS GPU and CPU backends.

The network topology was a 2-layer LIF SNN (784 inputs → $h$ hidden → 10 output, $T = 25$ steps), with five hyperparameter configurations:

| Configuration         | $\beta$      | $h$     | Acc (mlx-snn/MLX) | Time (s)   | GPU Mem (MB) | Acc (snnTorch/MPS) | Time (s) | GPU Mem (MB) | Acc (snnTorch/CPU) | Time (s) |
|-----------------------|--------------|---------|------------------|------------|--------------|---------------------|----------|--------------|---------------------|----------|
| C1                    | 0.85         | 256     | 97.28%           | 4.0        | 61           | 98.00%              | 8.8      | 241          | 98.01%              | 12.8     |
| C2                    | 0.9          | 256     | 97.02%           | 4.3        | 65           | 98.03%              | 8.9      | >241         | 97.97%              | 13.5     |
| C3                    | —            | 256*    | 96.91%           | 2.4        | —            | 98.03%              | 4.8      | —            | 98.17%              | 16.7     |
| C4                    | —            | 128     | 96.90%           | 4.3        | —            | 97.84%              | 9.0      | —            | 97.74%              | 10.9     |
| C5                    | 0.95         | 128     | 94.98%           | 4.4        | —            | 97.09%              | 9.0      | —            | 97.00%              | 11.1     |

(\*Batch size and learning rates varied per configuration.)

Peak GPU memory usage for mlx-snn ranged 61–138 MB, while snnTorch on MPS required 241–1453 MB. Per-epoch training was 2.0–2.5× faster for mlx-snn compared to snnTorch GPU; memory usage reduced by 3–10×.

Surrogate gradient performance (C5, 10 epochs):

| Surrogate         | Accuracy (%) | Training Time (s) |
|-------------------|-------------|-------------------|
| Fast sigmoid      | 93.65       | 44.7              |
| Arctan            | 92.44       | 43.8              |
| Straight-through  | 46.28       | 46.3              |

## 6. Usage Paradigms and Implementation

mlx-snn is distributed via PyPI and installed using:

```
pip install mlx-snn
```

A typical SNN definition, using MLX and mlx-snn modules, follows snnTorch conventions:

```python
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
import mlxsnn

class SpikingMLP(nn.Module):
    def __init__(self, num_steps=25, beta=0.9):
        super().__init__()
        self.fc1 = nn.Linear(784, 128)
        self.lif1 = mlxsnn.Leaky(beta=beta)
        self.fc2 = nn.Linear(128, 10)
        self.lif2 = mlxsnn.Leaky(beta=beta, reset_mechanism="none")
        self.num_steps = num_steps

    def __call__(self, spikes_in):
        s1 = self.lif1.init_state(spikes_in.shape[1], 128)
        s2 = self.lif2.init_state(spikes_in.shape[1], 10)
        for t in range(self.num_steps):
            x = self.fc1(spikes_in[t])
            spk, s1 = self.lif1(x, s1)
            x = self.fc2(spk)
            _, s2 = self.lif2(x, s2)
        return s2["mem"]   # final membrane potentials
```

To optimize training performance:

- Wrap the step function in `mx.compile`.
- Build the full T-step computation graph before any `mx.eval`.
- Use large batch sizes to amortize compute overhead.
- Avoid in-place operations to ensure graph purity and compatiblity with JIT transforms.

## 7. Best Practices, Limitations, and Development Roadmap

Recommended practices for high accuracy and efficiency:

- $\beta \approx 0.85\text{–}0.9$ for hidden layer decay.
- Hidden size $h = 256$ yields top-1 accuracy (97.28%).
- Adam optimizer, learning rate $1 \times 10^{-3}$.
- Fast-sigmoid surrogate ($k=25$) optimally balances accuracy and speed.

Deployment on Apple Silicon is streamlined: explicit device transfers are unnecessary; MLX automatically manages memory placement. The use of `mx.profile` aids in analyzing memory and compute overheads.

Limitations in version 0.2.1:

1. JIT compilation of temporal SNN loops is not yet supported (`mx.compile` cannot process Python `for` loops).
2. No integrated neuromorphic dataset support (e.g., N-MNIST, DVS-Gesture, SHD).
3. Published benchmarks are restricted to MNIST.
4. The STE surrogate introduces a roughly 1% accuracy penalty versus native custom VJPs.

Future development aims include:

- v0.3.0: Liquid State Machine (LSM) support, reservoir computing modules, EEG classification.
- v0.4.0: Full `mx.compile` optimization of forward passes, neuromorphic data loaders, visualization suite, benchmarks on complex datasets (CIFAR, DVS).
- v1.0.0: Stable API, extensive validation, release of comprehensive documentation, and a stable PyPI package.

In summary, mlx-snn enables SNN research, development, and production on Apple hardware, delivering competitive accuracy and resource efficiency by exploiting MLX features while maintaining a familiar and portable API for the broader spiking neural network community [2603.03529].

Source: https://www.emergentmind.com/topics/mlx-snn