[Paper] mlx-snn：在 Apple Silicon 上通过 MLX 实现脉冲神经网络

发布: 2天前 (2026年3月4日 GMT+8 05:25)

7 分钟阅读

原文: arXiv

Source: arXiv - 2603.03529v1

Overview

本文介绍了 mlx‑snn，这是第一个直接基于 Apple 的 MLX 框架构建的脉冲神经网络（SNN）库。通过利用 Apple Silicon 的统一内存和惰性执行模型，mlx‑snn 使在配备 M 系列芯片的 Mac 上进行高性能 SNN 研究和开发变得对开发者更加友好。

原生 Apple Silicon 支持 – 首个在 MLX 上运行的 SNN 库，消除了 macOS 上基于 PyTorch 的后端需求。
丰富的神经元和编码工具箱 – 实现了六种流行的脉冲神经元模型（LIF、IF、Izhikevich、Adaptive LIF、Synaptic、Alpha）和四种代理梯度函数，另外提供四种脉冲编码方案（包括针对 EEG 的编码器）。
端到端训练流水线 – 提供完整的时序反向传播（BPTT）工作流，使用 mx.grad 进行自动微分，并可选通过 mx.compile 进行 JIT 编译。
性能提升 – 在 M3 Max GPU 上相较于领先的基于 PyTorch 的库（snnTorch）实现 2.0–2.5 倍更快的训练速度和 3–10 倍更低的内存消耗。
开源且易于安装 – 在 MIT 许可证下分发，已发布至 PyPI，并包含可直接运行的 MNIST 分类示例。

Source: …

mlx‑snn 利用 MLX 编程模型，该模型将张量视为不可变的、惰性求值的对象，仅在需要时才会实际化。库将每个 SNN 层构建为可组合的函数；在前向传播时，使用选定的神经元模型和代理梯度生成脉冲，而反向传播则由 mx.grad 自动处理。

由于 MLX 会自动将张量放置在最合适的内存（CPU、GPU 或统一内存）中，开发者无需手动管理设备间的数据传输。

准确率 – 在 MNIST 基准测试中，mlx‑snn 在五组超参数设置下达到了 97.28 % 的分类准确率，匹配或略微超出基于 PyTorch 的 SNN 库报告的最佳结果。
速度 – 在 M3 Max GPU 上运行时，每个 epoch 的训练时间比 snnTorch 低 2.0–2.5×，这归功于惰性求值和内核融合。
内存 – 峰值内存使用降至 PyTorch 基准的 10–30 %，从而在相同硬件上能够使用更大的批量或更深的 SNN。
后端灵活性 – 作者评估了三种后端（纯 MLX、编译的 MLX 和回退 CPU 模式），并观察到在 GPU 后端上持续的性能提升。

这些结果证实，Apple Silicon 可以作为 SNN 研究的一流平台，而不会牺牲速度或模型容量。

Mac‑Centric Development – 已经在 Mac 上进行机器学习的研究人员和工程师（例如用于原型设计或边缘设备仿真），现在可以在不启动 Linux 容器或外部 GPU 的情况下实验 SNN。
Energy‑Efficient Inference – 由于 SNN 本身稀疏且事件驱动，在功耗高效的 M 系列芯片上运行，可实现 macOS 笔记本或未来基于 Apple 的边缘设备上的低延迟、低功耗 AI 服务。
Rapid Prototyping – 统一的 API（兼容 NumPy 风格代码）以及能够编译完整训练循环的特性，使得在新颖神经元动力学或自定义编码器上迭代变得轻松。
Cross‑Disciplinary Projects – 引入针对 EEG 的专用编码器，为神经技术开发者在 macOS 上直接构建脑机接口原型提供了可能。
Open‑Source Ecosystem – 在 PyPI 上发布并采用 MIT 许可证，鼓励社区贡献、插件开发（如自定义可塑性规则），以及与其他 Apple 为中心的机器学习工具（如 Core ML）的集成。

基准范围 – 评估仅限于 MNIST；更大的视觉或神经形态数据集（例如 DVS‑CIFAR10）尚未测试。
硬件特定性 – 性能提升与 Apple Silicon 绑定；移植到其他 GPU（CUDA、ROCm）需要额外的后端支持。
训练算法 – 该库目前专注于 BPTT；其他学习规则（例如 STDP、无代理梯度方法）尚未支持。
生态系统集成 – 虽然 MLX 正在发展，但更广泛的 Python ML 生态系统（TensorBoard、Hugging Face）仍缺乏对 mlx‑snn 的无缝挂钩。

未来工作可以扩展基准覆盖范围，添加对生物启发可塑性机制的支持，并将 mlx‑snn 与跨平台运行时桥接，以扩大其在 Apple 硬件之外的适用性。