Scalable Training of Continuous-Time Spiking Neural Networks with Differentiable Spike-Time Discretization
For researchers working on SNN training, this work provides a memory-efficient method that makes deeper continuous-time SNNs practical on limited hardware.
The paper tackles the memory bottleneck in training deep continuous-time spiking neural networks (SNNs). The proposed differentiable spike-time discretization (DSTD) reduces activation memory from O(N_out N_in) to O(N_out M) and achieves up to ~100x reduction in peak memory and ~20x reduction in training time, enabling training of 9-layer SNNs on CIFAR-10 and 20-layer SNNs on Fashion-MNIST on a single GPU.
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.