提出新方法让脉冲神经网络训练更省内存更快
Scalable Training of Continuous-Time Spiking Neural Networks with Differentiable Spike-Time Discretization
- 将不规则脉冲映射到固定时间点的可微加权事件，降低内存开销
- 在单卡上训练出9层CIFAR-10和20层Fashion-MNIST的SNN模型
- 适合做类脑计算、神经形态硬件部署的研究者参考
连续时间脉冲神经网络（SNN）为时序计算、计算神经科学和类脑硬件提供了事件驱动框架。然而，深度连续时间SNN的训练受限于精确脉冲时间计算所需的内存，该过程需在由前突触脉冲顺序决定的时间区间内评估并保留候选放电时间。本文提出一种基于可微脉冲时间离散化（DSTD）的高效训练框架，适用于具有通用膜和突触时间常数的漏积分-放电（LIF）神经元。DSTD将不规则前突触脉冲映射到固定时间点的可微加权事件，用$M$个固定时间区间替代依赖输入的候选维度，同时准确近似连续时间膜电位动态。这使基于首次脉冲时间（TTFS）编码的激活内存从$O(N_{\mathrm{out}}N_{\mathrm{in}})$降至$O(N_{\mathrm{out}}M)$，其中$N_{\mathrm{in}}$和$N_{\mathrm{out}}$分别为前突触与后突触神经元数量。此外，引入受同步放电链启发的时间正则化，组织层间放电窗口，缓解神经元失效问题，并实现类似流水线的处理。在密集型LIF层中，相比精确脉冲时间计算，DSTD使峰值内存消耗减少约100倍，训练时间缩短约20倍。结合上述方法，成功在单个GPU上训练了9层卷积SNN（CIFAR-10）和20层卷积SNN（Fashion-MNIST）。
原文摘要 · Abstract (English)
Continuous-time spiking neural networks (SNNs) provide an event-driven framework for temporal computation, computational neuroscience, and neuromorphic hardware. However, training deep continuous-time SNNs is severely constrained by the memory required for exact spike-time computation, which evaluates and retains candidate firing times over intervals determined by presynaptic spike ordering. Here we introduce a memory-efficient training framework based on differentiable spike-time discretization (DSTD) for leaky integrate-and-fire neurons with general membrane and synaptic time constants. DSTD maps irregular presynaptic spikes onto differentiable weighted events at fixed time points, replacing the input-dependent candidate dimension with $M$ fixed time intervals while accurately approximating continuous-time membrane-potential dynamics. This reduces candidate-related activation memory from $O(N_{\mathrm{out}}N_{\mathrm{in}})$ to $O(N_{\mathrm{out}}M)$ in the case of time-to-first-spike (TTFS) coding, where $N_{\mathrm{in}}$ and $N_{\mathrm{out}}$ denote the numbers of presynaptic and postsynaptic neurons, respectively. We further introduce synfire-chain-inspired temporal regularization that organizes layer-wise firing windows, mitigates dead-neuron failures, and enables pipeline-like processing. In dense LIF layers, DSTD reduced peak memory consumption by up to approximately 100-fold and training time by up to approximately 20-fold compared with exact spike-time computation. Together, these methods allowed us to train 9-layer convolutional SNNs on CIFAR-10 and 20-layer convolutional SNNs on Fashion-MNIST on a single GPU.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。