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

我们推出了 mlx-snn，这是首个原生构建于苹果 MLX 框架之上的脉冲神经网络（SNN）专用库。随着 SNN 研究的迅猛发展，当前所有主流 SNN 库——包括 snnTorch、Norse、SpikingJelly 和 Lava——均以 PyTorch 或自定义后端为基础，导致 Apple Silicon 用户长期缺乏原生支持方案。mlx-snn 提供六种神经元模型（LIF、IF、Izhikevich、自适应 LIF、突触型、Alpha），四种替代梯度函数，四种脉冲编码方法（含专为脑电图 EEG 设计的编码器），以及一套完整的、支持时间维度反向传播（BPTT）的训练流程。该库充分利用 MLX 框架的统一内存架构、惰性求值机制，以及可组合的函数变换能力（如 mx.grad 和 mx.compile），从而在 Apple Silicon 硬件上高效支撑 SNN 的研究与开发。我们在 MNIST 手写数字分类任务上对 mlx-snn 进行了系统性验证，覆盖五组超参数配置及三种后端环境；实验结果表明，在相同的 M3 Max 硬件平台上，相比 snnTorch，mlx-snn 最高可达 97.28% 的测试准确率，训练速度提升 2.0–2.5 倍，且 GPU 显存占用降低至其 1/3–1/10。mlx-snn 采用 MIT 开源许可证，已发布至 PyPI，源代码托管于 GitHub：https://github.com/D-ST-Sword/mlx-snn

苹果硅芯片（Apple Silicon）用户缺乏原生支持的脉冲神经网络（SNN）研究框架，现有主流SNN库（如snnTorch、Norse、SpikingJelly、Lava）均基于PyTorch或自定义后端，无法充分利用MLX框架的统一内存、lazy evaluation与函数式编译特性，导致在M系列芯片上训练效率低、显存开销大、开发体验割裂。该问题在SNN研究快速发展的背景下日益突出，属首次系统性填补MLX生态中SNN基础设施空白的新问题。

构建首个完全原生于MLX框架的轻量、可微、模块化SNN库（mlx-snn），深度耦合MLX核心范式——包括mx.grad自动微分、mx.compile图编译、零拷贝设备内存管理——实现真正的‘one-memory, one-graph’脉冲计算；通过将脉冲动力学建模为可微状态转移函数，并结合时间维度上的显式BPTT展开与MLX原生vmap/jit融合，突破传统SNN库对PyTorch动态图/autograd的依赖。

• 提供6种标准神经元模型（含EEG适配的Alpha与Adaptive LIF）、4种 surrogate梯度、4种编码器（含首创EEG-aware spike encoding）；• 完整BPTT训练流水线，支持时序展开、梯度截断与状态缓存优化；• 在M3 Max上对比snnTorch：准确率最高97.28%（MNIST），训练速度提升2.0–2.5×，GPU内存占用降低3–10×；• 全开源（MIT）、PyPI可安装、GitHub托管（https://github.com/D-ST-Sword/mlx-snn），含完整文档、教程与复现实验脚本；• 值得深入方向：MLX-native SNN在边缘实时推理（如Vision Pro头显）、跨模态脉冲编码（EEG+video）、以及与MLX生态其他组件（如mlx-vlm、mlx-llm）的协同微调。

• snnTorch: A PyTorch-based library for spiking neural networks (2021); • Norse: A Python library for simulating spiking neural networks on GPUs (2022); • SpikingJelly: A deep learning framework for SNNs with PyTorch backend (2022); • Lava: A framework for neuromorphic computing developed by Intel (2021); • MLX: A framework for machine learning on Apple silicon (2023, Apple AI Research)