[Paper Review] Towards Memory- and Time-Efficient Backpropagation for Training Spiking Neural Networks

Qingyan Meng, Mingqing Xiao+4 more|arXiv (Cornell University)|Feb 28, 2023

Advanced Memory and Neural Computing4 citations

TL;DR

This paper proposes Spatial Learning Through Time (SLTT), a memory- and time-efficient backpropagation method for training Spiking Neural Networks (SNNs). By selectively ignoring unimportant temporal routes in the computational graph during backpropagation, SLTT reduces scalar multiplications and achieves constant memory usage independent of sequence length. On ImageNet, SLTT achieves state-of-the-art accuracy with over 70% lower memory cost and 50% faster training than BPTT with surrogate gradients.

ABSTRACT

Spiking Neural Networks (SNNs) are promising energy-efficient models for neuromorphic computing. For training the non-differentiable SNN models, the backpropagation through time (BPTT) with surrogate gradients (SG) method has achieved high performance. However, this method suffers from considerable memory cost and training time during training. In this paper, we propose the Spatial Learning Through Time (SLTT) method that can achieve high performance while greatly improving training efficiency compared with BPTT. First, we show that the backpropagation of SNNs through the temporal domain contributes just a little to the final calculated gradients. Thus, we propose to ignore the unimportant routes in the computational graph during backpropagation. The proposed method reduces the number of scalar multiplications and achieves a small memory occupation that is independent of the total time steps. Furthermore, we propose a variant of SLTT, called SLTT-K, that allows backpropagation only at K time steps, then the required number of scalar multiplications is further reduced and is independent of the total time steps. Experiments on both static and neuromorphic datasets demonstrate superior training efficiency and performance of our SLTT. In particular, our method achieves state-of-the-art accuracy on ImageNet, while the memory cost and training time are reduced by more than 70% and 50%, respectively, compared with BPTT.

Motivation & Objective

• To address the high memory and training time costs of backpropagation through time (BPTT) with surrogate gradients in Spiking Neural Networks (SNNs).
• To reduce the number of scalar multiplications during backpropagation by identifying and pruning unimportant temporal routes in the computational graph.
• To enable online, incremental gradient computation in SNNs to eliminate dependence on storing full sequence states, thus achieving constant memory usage regardless of sequence length.
• To develop SLTT-K, a variant that performs backpropagation only at K selected time steps, further reducing computational complexity without performance loss.
• To achieve state-of-the-art SNN performance on benchmark datasets while significantly improving training efficiency.

Proposed method

• The method decomposes gradients from BPTT into spatial and temporal components to explicitly analyze temporal contributions to error backpropagation.
• It introduces a pruning strategy that removes unimportant temporal paths in the computational graph during backpropagation, reducing scalar multiplications.
• SLTT enables online training by computing gradients instantaneously at each time step, eliminating the need to store intermediate states across time steps.
• The SLTT-K variant restricts backpropagation to only K time steps, reducing time complexity from Ω(T) to Ω(K), where T is the total sequence length.
• The framework uses time-step-wise batch normalization and replaces max pooling with average pooling to support online learning and improve training stability.
• The method is compatible with standard SNN architectures such as VGG-11, ResNet, and normalization-free ResNets (NF-ResNet), and supports both static and neuromorphic datasets.

Figure 1 : The training time and memory cost comparison between the proposed SLTT-1 method and the BPTT with SG method on ImageNet. SLTT-1 achieves similar accuracy as BPTT, while owning better training efficiency than BPTT both theoretically and experimentally. Please refer to Sections 4 and 5 for

Experimental results

Research questions

• RQ1 To what extent do temporal dependencies in backpropagation contribute to the final gradients in SNNs?
• RQ2 Can unimportant temporal routes in the computational graph be safely pruned during backpropagation without degrading model performance?
• RQ3 Can online gradient computation be achieved in SNNs without storing full sequence states, enabling constant memory usage independent of sequence length?
• RQ4 How does limiting backpropagation to only K time steps affect training efficiency and model accuracy?
• RQ5 Can the proposed method achieve state-of-the-art performance on large-scale datasets like ImageNet while significantly reducing memory and time costs?

Key findings

• On ImageNet, SLTT achieves state-of-the-art accuracy while reducing memory cost by over 70% and training time by more than 50% compared to BPTT with surrogate gradients.
• The SLTT-1 variant matches BPTT performance on ImageNet, CIFAR-10, CIFAR-100, DVS-Gesture, and DVS-CIFAR10, while drastically improving training efficiency.
• SLTT-K reduces the number of scalar multiplications from Ω(T) to Ω(K), enabling further time complexity reduction with no performance degradation.
• The method achieves competitive performance on both static and neuromorphic datasets, including SOTA results on ImageNet with a 6-time-step sequence.
• The online training framework of SLTT eliminates the need for storing intermediate states, enabling constant memory usage regardless of sequence length.
• The method is compatible with various SNN architectures, including VGG-11, NF-ResNet-34/50/101, and supports both cross-entropy and MSE-based loss functions for different datasets.

Figure 2 : Computational graph of multi-layer SNNs. Dashed arrows represent the non-differentiable spike generation functions.