[论文解读] 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 Computing被引用 4

一句话总结

本文提出时空空间学习（SLTT），一种用于训练脉冲神经网络（SNNs）的内存和时间效率更高的反向传播方法。通过在反向传播过程中选择性忽略计算图中不重要的时间路径，SLTT 减少了标量乘法运算次数，并实现了与序列长度无关的恒定内存使用。在 ImageNet 上，SLTT 实现了当前最优的准确率，相比使用替代梯度的 BPTT，内存消耗降低超过 70%，训练速度提升 50%。

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.

研究动机与目标

• 为解决使用替代梯度的反向传播通过时间（BPTT）在脉冲神经网络（SNNs）中带来的高内存和训练时间开销问题。
• 通过识别并剪除计算图中不重要的时间路径，减少反向传播过程中的标量乘法次数。
• 在 SNN 中实现在线、增量式梯度计算，消除对存储完整序列状态的依赖，从而实现与序列长度无关的恒定内存使用。
• 开发 SLTT-K 变体，仅在 K 个选定时间步进行反向传播，进一步降低计算复杂度且不造成性能损失。
• 在基准数据集上实现当前最优的 SNN 性能，同时显著提升训练效率。

提出的方法

• 该方法将 BPTT 的梯度分解为空间分量和时间分量，以显式分析时间贡献对误差反向传播的影响。
• 提出一种剪枝策略，在反向传播过程中移除计算图中不重要的时间路径，从而减少标量乘法运算。
• SLTT 通过在每个时间步即时计算梯度，实现在线训练，无需存储跨时间步的中间状态。
• SLTT-K 变体将反向传播限制在仅 K 个时间步，将时间复杂度从 Ω(T) 降低至 Ω(K)，其中 T 为总序列长度。
• 该框架采用逐时间步的批量归一化，并将最大池化替换为平均池化，以支持在线学习并提升训练稳定性。
• 该方法兼容标准 SNN 架构，如 VGG-11、ResNet 和无归一化残差网络（NF-ResNet），并支持静态数据集和类脑数据集。

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

实验结果

研究问题

• RQ1 在 SNN 的反向传播中，时间依赖性在最终梯度中所起的作用有多大？
• RQ2 在反向传播过程中，能否安全地剪除计算图中不重要的时间路径而不降低模型性能？
• RQ3 能否在不存储完整序列状态的前提下实现在 SNN 中的在线梯度计算，从而实现与序列长度无关的恒定内存使用？
• RQ4 将反向传播限制在仅 K 个时间步，对训练效率和模型准确率有何影响？
• RQ5 所提出的方法能否在 ImageNet 等大规模数据集上实现当前最优性能，同时显著降低内存和时间开销？

主要发现

• 在 ImageNet 上，SLTT 实现了当前最优的准确率，相比使用替代梯度的 BPTT，内存消耗降低超过 70%，训练时间减少 50% 以上。
• SLTT-1 变体在 ImageNet、CIFAR-10、CIFAR-100、DVS-Gesture 和 DVS-CIFAR10 上的性能与 BPTT 相当，同时显著提升了训练效率。
• SLTT-K 将标量乘法次数从 Ω(T) 降低至 Ω(K)，在不造成性能损失的前提下进一步降低了时间复杂度。
• 该方法在静态数据集和类脑数据集上均表现出色，包括在 6 时间步序列下于 ImageNet 上实现当前最优结果。
• SLTT 的在线训练框架消除了对存储中间状态的需求，实现了与序列长度无关的恒定内存使用。
• 该方法兼容多种 SNN 架构，包括 VGG-11、NF-ResNet-34/50/101，并支持针对不同数据集的交叉熵和 MSE 基损失函数。

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