mlx-snn

基于 Apple MLX 构建的通用脉冲神经网络库。

mlx-snn 旨在提供一个高效、适合研究的 SNN 框架，充分利用 MLX 的统一内存架构和惰性求值机制。无论您是在探索神经元动力学、使用代理梯度训练分类器，还是通过 NIR 进行模型交换，mlx-snn 都提供了一个简洁、符合 Python 风格的 API，能够自然地融入 MLX 生态系统。

9 种神经元模型 · 6 种代理梯度 · 8 种脉冲编码 · 5 个神经形态数据集 · LSM 储层 · NIR 互操作性 · 403 个测试用例

亮点

统一内存 SNN SNN 在每个时间步都会存储每个神经元的状态——这在独立 GPU 架构上会成为内存瓶颈。而 Apple Silicon 的统一内存消除了 CPU↔GPU 之间的数据传输，使得我们可以构建更长的时间窗口和更大的储层，而不受显存限制。

能效提升 17–25 倍 M3 Max 在功耗仅为 Tesla V100 七分之一的情况下，训练 SNN 的速度比后者快 2.6–3.7 倍。递归脉冲动力学主要受限于延迟而非计算能力，因此 Apple Silicon 的高带宽统一架构相比数据中心的并行计算更具优势。

多尺度时间建模 MSLeaky 为并行的脉冲分支分配频率匹配的衰减率，从而在一个网络中同时捕捉 delta 到 gamma 频段的动力学特性。结合状态分离的分块 BPTT 方法，可以处理长时间的生物信号序列（如 EEG、fMRI），而不会导致内存爆炸式增长。

安装

pip install mlx-snn

需要 Python 3.9 或更高版本以及 Apple Silicon 处理器（M1/M2/M3/M4）。

快速入门

import mlx.core as mx
import mlx.nn as nn
import mlxsnn

# 构建一个脉冲网络
fc = nn.Linear(784, 10)
lif = mlxsnn.Leaky(beta=0.95, threshold=1.0)

# 将输入编码为脉冲序列并随时间运行
spikes_in = mlxsnn.rate_encode(mx.random.uniform(shape=(8, 784)), num_steps=25)
state = lif.init_state(batch_size=8, features=10)

for t in range(25):
    spk, state = lif(fc(spikes_in[t]), state)

print("输出膜电位:", state["mem"].shape)  # (8, 10)

功能特性

神经元模型

所有神经元都支持 learn_threshold 和可配置的复位机制（subtract / zero / none）。具有衰减常数的神经元支持 learn_beta；递归神经元则支持 learn_V。状态始终以显式字典形式表示，与 MLX 的函数式变换和 mx.compile 兼容。

模型	描述	状态变量
Leaky (LIF)	可配置衰减的漏积分发放神经元	mem
IF	不漏的积分发放神经元（理想积分器）	mem
Izhikevich	二维动力学，支持 RS/IB/CH/FS 预设	v, u
ALIF	自适应 LIF，动态阈值	mem, adapt
Synaptic	基于电导的双状态 LIF	syn, mem
Alpha	双指数突触模型	syn_exc, syn_inh, mem
RLeaky	可学习反馈的递归 LIF	mem, spk
RSynaptic	可学习反馈的递归 Synaptic	syn, mem, spk
MSLeaky	多尺度 LIF，每条分支有频率匹配的 beta	每分支 mem

代理梯度

所有神经元都支持通过 6 种代理梯度函数进行可微训练：

函数	代理梯度	属性
Arctan（默认）	α / (π(1 + α²x²))	BPTT 收敛稳定，局部性适中
Fast Sigmoid	scale / (2(1 + scale·|x|)²)	尾部较重，理性衰减
Sigmoid	slope · σ(slope·x) · (1 − σ(slope·x))	标准逻辑函数的导数
Triangular	max(0, 1 − |x|)	支撑紧凑，在阈值附近局部化
Straight-Through	1	最简单，处处梯度为 1
Custom	用户自定义	可插入任意可微近似

脉冲编码

8 种编码方法用于将连续信号转换成脉冲序列：

方法	输入	使用场景
Rate (Poisson)	静态值	图像、通用分类任务
Latency (TTFS)	静态值	高能效的时间编码
Delta Modulation	时间信号	变化检测、事件型编码
Direct	任意张量	直通预计算输入
Repeat	脉冲模式	将脉冲序列平铺到更长的时间窗口
Frequency-Band	EEG 信号	基于 FFT 的分解，划分为 delta/theta/alpha/beta/gamma 波段
Threshold-Crossing	时间信号	多级幅值穿越检测
EEG Encoder	原始 EEG	可配置速率/delta/阈值的生物信号编码

卷积 SNN 层

利用空间池化构建深层脉冲卷积网络：

import mlxsnn

conv1 = mlxsnn.SpikingConv2d(in_channels=2, out_channels=64, kernel_size=3, padding=1)
pool1 = mlxsnn.SpikingMaxPool2d(kernel_size=2, stride=2)
lif1  = mlxsnn.Leaky(beta=0.95)
drop  = mlxsnn.SpikeDropout(p=0.2)  # 脉冲感知（无需重新缩放）
flat  = mlxsnn.SpikingFlatten()

此外还包含 6 种基于数学原理的时序算子，适用于卷积 SNN：TAC（时序聚合卷积）、TAC-TP（时序保持卷积）、L-TAC（可学习卷积）、FTC（傅里叶时序卷积）、IMC（信息最大化脉冲卷积）、TCC（时序坍缩卷积）。

液态机

基于脉冲神经元的储层计算——随机稀疏的递归连接，拓扑结构可配置：

import mlxsnn

lsm = mlxsnn.LSM(
    input_size=64, reservoir_size=500, output_size=10,
    connectivity=0.1, spectral_radius=0.9,
    topology="small_world",  # 也可选："erdos_renyi"、"scale_free"
    exc_ratio=0.8,           # 戴尔定律：80% 兴奋性
)
state = lsm.init_state(batch_size=32)

for t in range(num_steps):
    output, state = lsm(spikes[t], state)

训练工具

BPTT 变体：

• bptt_forward(model, spikes, state) — 标准的随时间反向传播
• chunked_bptt_forward(model, spikes, state, chunk_size) — 通过在块边界分离状态，在长序列上实现内存高效的训练
• detach_state(state) — 分离状态字典中的所有张量（用于截断 BPTT）

mx.compile 包装器：

• compiled_step(model) — 返回一个编译后的单步可调用函数
• compiled_forward(model, spikes, state) — 每个时间步都使用 mx.compile 的 BPTT 前向传播

损失函数（共 11 种）：

损失	方法
ce_rate_loss	对尖峰率（尖峰计数 / T）的交叉熵
ce_count_loss	对原始尖峰计数的交叉熵
mse_membrane_loss	对最终膜电位的均方误差
mse_count_loss	对尖峰计数与目标之间的均方误差
membrane_loss	对最终膜电位的交叉熵
rate_coding_loss	对尖峰计数的 log-softmax 的交叉熵
activity_reg_loss	惩罚与目标发放率的偏差
l1_spike_loss	对尖峰计数的 L1 稀疏惩罚
l2_spike_loss	对尖峰计数的 L2 正则化

实用函数：spike_rate, spike_count。

可学习参数： learn_threshold（所有神经元）、learn_beta（基于 LIF 的神经元）、learn_V（递归神经元）。支持标准 MLX 优化器（mlx.optimizers.Adam 等）。

神经形态数据集

5 个内置数据集加载器，具备自动下载、缓存以及事件到帧的转换功能：

数据集	模态	类别	样本数
DVS-Gesture	事件相机（手势）	11	1,342
CIFAR10-DVS	事件相机（自然图像）	10	10,000
N-MNIST	事件相机（数字）	10	70,000
SHD	音频（语音数字）	20	10,420
SSC	音频（语音命令）	35	100,000+

from mlxsnn.datasets import DVSGestureDataset, create_dataloader

dataset = DVSGestureDataset(root="./data", train=True, num_steps=16)
loader = create_dataloader(dataset, batch_size=16, shuffle=True)

可视化

需要运行 pip install mlx-snn[viz]（matplotlib）。

from mlxsnn.utils.visualization import plot_raster, plot_membrane, plot_firing_rate

plot_raster(spike_tensor, title="第1层尖峰")       # 尖峰随时间的点阵图
plot_membrane(state["mem"], title="膜电位")          # 膜电位轨迹
plot_firing_rate(spike_tensor, title="发放率")        # 每个神经元的发放率

NIR 互操作性

NIR 支持跨框架的 SNN 模型交换——可以导入/导出至 snnTorch、Norse、SpikingJelly 以及神经形态硬件。

pip install mlx-snn[nir]

# 导出
layers = [('fc1', nn.Linear(784, 128)), ('lif1', mlxsnn.Leaky(beta=0.9)),
          ('fc2', nn.Linear(128, 10)),   ('lif2', mlxsnn.Leaky(beta=0.9))]
nir.write('model.nir', mlxsnn.export_to_nir(layers))

# 导入
model = mlxsnn.import_from_nir(nir.read('model.nir'))
out, state = model(x, model.init_states(batch_size=32))

支持：nn.Linear ↔ nir.Affine/nir.Linear、Leaky ↔ nir.LIF、IF ↔ nir.IF、Synaptic ↔ nir.CubaLIF。

基准测试

所有实验采用相同的超参数：Adam（LR=1e-3）、泊松率编码、T=25 个时间步、批量大小 128、5 个随机种子、10 个 epoch、surrogate_fn="fast_sigmoid"。完整脚本位于 benchmarks/ 目录下，详细内容参见我们的 arXiv 论文。

训练准确度（在噪声范围内一致）

任务	mlx-snn (M3 Max)	snnTorch (V100)
FC SNN 在 MNIST 上	97.0%	97.2%
FC SNN 在 FashionMNIST 上	85.2%	85.2%
LSM 储存器在 MNIST 上	92.4%	93.6%

训练速度

任务	mlx-snn (M3 Max)	snnTorch (V100)	加速倍数
FC SNN（784→128→10）	7.4 s/epoch	19.4 s/epoch	2.6x
LSM 储存器（500 个神经元）	3.2 s/epoch	5.7 s/epoch	1.8x

推理吞吐量（样本/秒，T=25）

模型	批量	mlx-snn (M3 Max)	snnTorch (V100)	加速倍数
FC-SNN	128	23,875	6,552	3.6x
FC-SNN	512	95,735	26,058	3.7x
Conv-SNN	128	6,671	3,859	1.7x
Conv-SNN	512	6,434	4,252	1.5x

mx.compile 加速效果

模式	时间（25 步前向传播）
未编译	2.72 ms
编译	0.92 ms
加速倍数	2.9x

能效

M3 Max（TDP ~45 W）的 SNN 训练/推理速度比 V100（TDP 300 W）快 2.6–3.7 倍，估计能效提升 17–25 倍（每瓦性能）。

详细结果及复现脚本请参阅 benchmarks/ 目录和我们的 基准测试论文。

从 snnTorch 迁移

mlx-snn 的设计旨在让 snnTorch 用户感到熟悉：

	snnTorch (PyTorch)	mlx-snn (MLX)
导入	import snntorch as snn	import mlxsnn
创建	lif = snn.Leaky(beta=0.9)	lif = mlxsnn.Leaky(beta=0.9)
前向传播	spk, mem = lif(x, mem)	spk, state = lif(x, state)
状态	分离的张量（mem）	显式字典（state["mem"]）
张量	torch.Tensor	mx.array
梯度	autograd + 替代梯度	STE 模式 + mx.stop_gradient

关键设计差异：状态始终是显式的字典——传入即得到。没有隐藏的实例变量。这与 MLX 的函数式变换（mx.grad、mx.vmap、mx.compile）非常契合。

项目结构

mlxsnn/
├── neurons/       # Leaky、IF、Izhikevich、ALIF、Synaptic、Alpha、RLeaky、RSynaptic、MSLeaky
├── surrogate/     # arctan、fast_sigmoid、sigmoid、triangular、straight_through、自定义
├── encoding/      # 率编码、延迟编码、delta 编码、直接编码、重复编码、频带编码、阈值穿越编码、EEG
├── functional/    # 无状态纯函数、11 种损失函数、指标
├── layers/        # SpikingConv2d、MaxPool2d、AvgPool2d、Flatten、SpikeDropout
├── operators/     # TAC、TAC-TP、L-TAC、FTC、IMC、TCC
├── liquid/        # LiquidReservoir、LSM、拓扑生成器
├── datasets/      # DVSGesture、CIFAR10DVS、NMNIST、SHD、SSC
├── training/      # BPTT、分块 BPTT、mx.compile 包装器
├── utils/         # 可视化、状态管理
└── nir_*.py       # NIR 导出/导入工具

路线图

• v0.1 — 核心神经元（LIF、IF）、代理梯度、速率/延迟编码
• v0.2 — 扩展神经元（伊兹基维奇、ALIF、突触型、Alpha）、脑电图编码器、差分编码
• v0.3 — NIR 互操作性（导出/导入）
• v0.4 — 循环神经元、卷积/池化层、神经形态数据集、TAC 算子
• v0.5 — 直接/重复编码、活动正则化、SpikeDropout、可视化、SHD 数据集
• v0.6 — CI/CD、API 文档站点、完整示例
• v0.7 — LSM、MSLeaky 神经元、分块 BPTT、mx.compile、频带及阈值穿越编码
• v1.0 — 完整文档、全面基准测试、JOSS 论文

出版物

• mlx-snn v0.1: 通过 MLX 在 Apple Silicon 上实现脉冲神经网络（arXiv，2026 年）
• mlx-snn v0.4: 在 Apple Silicon 上训练脉冲神经网络：跨框架基准测试（准备中）

引用

如果您在研究中使用 mlx-snn，请引用以下内容：

@misc{qin2026mlxsnn,
  title         = {mlx-snn: 通过 {MLX} 在 Apple Silicon 上实现脉冲神经网络},
  author        = {Jiahao Qin},
  year          = {2026},
  eprint        = {2603.03529},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2603.03529}
}

贡献

欢迎贡献！请在 GitHub 上提交问题或拉取请求。

许可证

GPL-3.0

mlx-snn 快速上手指南

mlx-snn 是一个基于 Apple MLX 框架构建的通用脉冲神经网络（SNN）库。它利用 Apple Silicon 的统一内存架构和惰性求值特性，为研究神经元动力学、使用代理梯度训练分类器以及模型交换提供了高效且符合 Python 习惯的 API。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：macOS
• 硬件要求：必须使用 Apple Silicon 芯片设备（M1, M2, M3, M4 系列）。该库不支持 Intel Mac 或非 Apple 设备。
• Python 版本：Python 3.9 或更高版本。
• 前置依赖：无需手动安装 MLX，mlx-snn 会自动处理相关依赖。

提示：国内开发者若遇到 PyPI 下载缓慢的问题，建议使用国内镜像源进行安装。

安装步骤

使用 pip 进行安装。推荐使用国内镜像源以加速下载：

pip install mlx-snn -i https://pypi.tuna.tsinghua.edu.cn/simple

如果需要安装可视化功能（依赖 matplotlib）或 NIR 互操作性支持，可使用以下命令：

# 安装可视化组件
pip install "mlx-snn[viz]" -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装 NIR 互操作组件
pip install "mlx-snn[nir]" -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

以下是一个最简示例，展示如何构建一个简单的脉冲神经网络，将输入编码为脉冲序列，并随时间步运行网络。

1. 构建网络与前向传播

import mlx.core as mx
import mlx.nn as nn
import mlxsnn

# 1. 定义网络层：全连接层 + 漏积分发放 (LIF) 神经元
fc = nn.Linear(784, 10)
lif = mlxsnn.Leaky(beta=0.95, threshold=1.0)

# 2. 数据编码：将随机输入转换为脉冲序列 (Rate Encoding)
# 形状：(Batch_Size, Features) -> 编码为时间步序列
spikes_in = mlxsnn.rate_encode(mx.random.uniform(shape=(8, 784)), num_steps=25)

# 3. 初始化神经元状态
state = lif.init_state(batch_size=8, features=10)

# 4. 随时间步运行网络 (Time-Loop)
for t in range(25):
    # 输入当前时间步的脉冲，获取输出脉冲和更新后的状态
    spk, state = lif(fc(spikes_in[t]), state)

# 查看最终膜电位状态形状
print("Output membrane:", state["mem"].shape)  # 输出：(8, 10)

2. 核心特性速览

• 神经元模型：支持 LIF, IF, Izhikevich, ALIF 等 9 种模型，均支持 learn_threshold 和可配置的复位机制。
• 代理梯度：内置 Arctan, Fast Sigmoid 等 6 种代理梯度函数，支持端到端微分训练。
• 脉冲编码：提供 Rate, Latency, Delta Modulation 等 8 种编码方式，可直接处理静态图像或时序信号（如 EEG）。
• 训练工具：支持标准 BPTT 及显存优化的 chunked_bptt，兼容 mx.compile 进行加速。
• 数据集：内置 DVS-Gesture, CIFAR10-DVS, SHD 等 5 种神经形态数据集加载器。

3. 使用卷积 SNN 层（可选）

构建深层脉冲卷积网络同样简单：

import mlxsnn

conv1 = mlxsnn.SpikingConv2d(in_channels=2, out_channels=64, kernel_size=3, padding=1)
pool1 = mlxsnn.SpikingMaxPool2d(kernel_size=2, stride=2)
lif1  = mlxsnn.Leaky(beta=0.95)
drop  = mlxsnn.SpikeDropout(p=0.2)  # 脉冲感知的 Dropout

现在您已经完成了环境的搭建并运行了第一个脉冲神经网络。更多高级用法（如液体状态机 LSM、NIR 模型导出等）请参考官方文档。