[论文解读] One Timestep is All You Need: Training Spiking Neural Networks with Ultra Low Latency

Sayeed Shafayet Chowdhury, Nitin Rathi+1 more|arXiv (Cornell University)|Oct 1, 2021

Advanced Memory and Neural Computing被引用 20

一句话总结

本文提出 IIR-SNN，一种训练方法，使脉冲神经网络（SNNs）能够实现单时间步推理——将延迟降低高达 2500 倍——同时保持高精度。通过从更高时间步长模型初始化的 SNN 进行迭代微调，该方法实现了单位延迟推理，其能效比传统人工神经网络（ANNs）高出 25–33 倍。

ABSTRACT

Spiking Neural Networks (SNNs) are energy efficient alternatives to commonly used deep neural networks (DNNs). Through event-driven information processing, SNNs can reduce the expensive compute requirements of DNNs considerably, while achieving comparable performance. However, high inference latency is a significant hindrance to the edge deployment of deep SNNs. Computation over multiple timesteps not only increases latency as well as overall energy budget due to higher number of operations, but also incurs memory access overhead of fetching membrane potentials, both of which lessen the energy benefits of SNNs. To overcome this bottleneck and leverage the full potential of SNNs, we propose an Iterative Initialization and Retraining method for SNNs (IIR-SNN) to perform single shot inference in the temporal axis. The method starts with an SNN trained with T timesteps (T>1). Then at each stage of latency reduction, the network trained at previous stage with higher timestep is utilized as initialization for subsequent training with lower timestep. This acts as a compression method, as the network is gradually shrunk in the temporal domain. In this paper, we use direct input encoding and choose T=5, since as per literature, it is the minimum required latency to achieve satisfactory performance on ImageNet. The proposed scheme allows us to obtain SNNs with up to unit latency, requiring a single forward pass during inference. We achieve top-1 accuracy of 93.05%, 70.15% and 67.71% on CIFAR-10, CIFAR-100 and ImageNet, respectively using VGG16, with just 1 timestep. In addition, IIR-SNNs perform inference with 5-2500X reduced latency compared to other state-of-the-art SNNs, maintaining comparable or even better accuracy. Furthermore, in comparison with standard DNNs, the proposed IIR-SNNs provide25-33X higher energy efficiency, while being comparable to them in classification performance.

研究动机与目标

• 解决脉冲神经网络（SNNs）中高推理延迟的问题，该问题限制了其在边缘设备上的部署。
• 克服 SNN 中低延迟与高精度之间的权衡，特别是在静态图像分类和强化学习任务中。
• 通过实现单时间步推理，消除存储和获取膜电位所导致的内存访问开销。
• 在保持具有竞争力的分类性能的同时，实现与标准深度神经网络（DNNs）相当或更优的能效比。
• 通过最小化延迟，使 SNN 在机器人和自主系统等实时应用中实现实际部署。

提出的方法

• 使用 5 个时间步和直接输入编码，通过替代梯度反向传播训练初始 SNN。
• 将已在较高时间步（如 T=5）训练好的 SNN 作为后续在更低时间步（如 T=4, 3, ..., 1）训练的初始化。
• 采用逐步减少时间步的迭代微调方法，通过渐进式微调保持性能。
• 将神经元阈值视为可学习参数，以提升在不同延迟阶段的训练稳定性和性能。
• 利用混合训练流程：先训练一个 ANN，将其转换为 SNN，然后在 SNN 领域中通过反向传播进行微调。
• 使用直接输入编码以减少收敛所需的时间步数，从而实现高效的低延迟训练。

实验结果

研究问题

• RQ1 SNN 是否能够仅通过一个推理时间步就实现高精度，且性能与多时间步 SNN 相当或更优？
• RQ2 从更高时间步长模型进行迭代微调，能在多大程度上降低延迟而不造成显著的精度下降？
• RQ3 与最先进 SNN 及等效 DNN 相比，单时间步 SNN 推理在能效和延迟方面表现如何？
• RQ4 IIR-SNN 是否能有效应用于强化学习任务（如 DQN 在 CartPole 和 Atari Pong 上）并实现极低延迟？
• RQ5 减少脉冲数量并消除膜电位内存访问对整体能效的影响是什么？

主要发现

• 在仅使用一个推理时间步的情况下，IIR-SNN 在 CIFAR-10 上实现 93.05% 的 top-1 精度，在 CIFAR-100 上实现 70.15%，在 ImageNet 上实现 67.71%，均基于 VGG16 模型。
• 与最先进 SNN 相比，该方法将推理延迟降低了 5–2500 倍，同时保持或提升了精度。
• IIR-SNN 的能效比等效 DNN 高出 25–33 倍，在强化学习中相比 ANN-DQN 提升 7.55 倍。
• 在 Atari Pong 上，T1 SNN 达到 17.4 ± 3.2 的奖励，而 T5 SNN 达到 19.4 ± 1.3，与 ANN-DQN 的性能（19.7 ± 1.1）相当，但延迟显著降低。
• T1 SNN-DQN 的平均脉冲频率仅为 0.08，相比 ANN-DQN 实现了 7.55 倍的能效提升，表明其具有极低计算量。
• 该方法消除了多时间步 SNN 中存储和检索膜电位所导致的内存访问开销，从而提升了计算和能效效率。