[论文解读] Temporal Efficient Training of Spiking Neural Network via Gradient Re-weighting

Shikuang Deng, Yuhang Li+2 more|arXiv (Cornell University)|Feb 24, 2022

Advanced Memory and Neural Computing被引用 102

一句话总结

论文引入 Temporal Efficient Training (TET) 与梯度重权用于尖峰神经网络（SNNs）以提升泛化和时序可扩展性，在神经形态和静态数据集上取得最先进的结果，包括 DVS-CIFAR10 的 top-1 83.17%。

ABSTRACT

Recently, brain-inspired spiking neuron networks (SNNs) have attracted widespread research interest because of their event-driven and energy-efficient characteristics. Still, it is difficult to efficiently train deep SNNs due to the non-differentiability of its activation function, which disables the typically used gradient descent approaches for traditional artificial neural networks (ANNs). Although the adoption of surrogate gradient (SG) formally allows for the back-propagation of losses, the discrete spiking mechanism actually differentiates the loss landscape of SNNs from that of ANNs, failing the surrogate gradient methods to achieve comparable accuracy as for ANNs. In this paper, we first analyze why the current direct training approach with surrogate gradient results in SNNs with poor generalizability. Then we introduce the temporal efficient training (TET) approach to compensate for the loss of momentum in the gradient descent with SG so that the training process can converge into flatter minima with better generalizability. Meanwhile, we demonstrate that TET improves the temporal scalability of SNN and induces a temporal inheritable training for acceleration. Our method consistently outperforms the SOTA on all reported mainstream datasets, including CIFAR-10/100 and ImageNet. Remarkably on DVS-CIFAR10, we obtained 83$\%$ top-1 accuracy, over 10$\%$ improvement compared to existing state of the art. Codes are available at \url{https://github.com/Gus-Lab/temporal_efficient_training}.

研究动机与目标

• 动机：在直接训练中，使用替代梯度（SG）训练的SNNs泛化能力有限。
• 提出 Temporal Efficient Training (TET)，以使梯度信号与每个时间步的输出对齐并促进更平坦的极小值。
• 证明 TET 提升时序可扩展性并使时间继承训练（Time Inheritance Training, TIT）成为可能，从而降低训练时间。
• 展示在静态与 neuromorphic 数据集上的最先进性能（例如 CIFAR-10/100、ImageNet、DVS-CIFAR10）。

提出的方法

• 采用带Spike/Membrane 动态的迭代 Leaky Integrate-and-Fire (LIF) 模型用于SNNs。
• 用 LE_TET 的损失替代标准的直接训练损失，该损失通过在每个时间 t 处对输出进行交叉熵优化每一个时间步的输出。
• 结合正则化项 L_MSE 以约束每时间步输出并稳定训练。
• 表明 SDT 与 TET 产生不同的损失景观，TET 趋向于更平坦的极小值。
• 引入 Time Inheritance Training (TIT) 在较短的仿真长度上启动训练并逐步扩展 T。
• 提供一个损失公式，通过 L_TET 与 Jensen/AM-GM 不等式上界传统 SDT 损失（L_SDT）。
• 保持推断规则不变；仅修改训练损失。

实验结果

研究问题

• RQ1 替代梯度基直接训练（SDT）是否会将 SNNs 陷入比 TET 更尖锐的局部极小值？
• RQ2 每时间步训练损失（L_TET）是否能使极小值更平坦、提高 SNN 的泛化能力？
• RQ3 Time Inheritance Training (TIT) 是否在保持或提高准确率的同时减少训练时间？
• RQ4 与 state-of-the-art 方法相比，TET 在静态数据集（CIFAR-10/100、ImageNet）和神经形态数据（DVS-CIFAR10）上的表现如何？

主要发现

• TET 在多个网络和数据集上显著提升了相对于使用 SG 的 SDT 的准确性（如 CIFAR-100 上的 ResNet-19）。
• 在局部极小值处，TET 下的损失景观比 SDT 下更平坦，表明更好的泛化。
• 在训练中段从 SDT 切换到 TET 可以帮助摆脱尖锐的局部极小值并提升测试性能。
• Time Inheritance Training (TIT) 将训练时间大幅减少约一半，同时保持或提升准确率，尤其随着仿真长度的增加。
• 在 DVS-CIFAR10 上使用 VGGSNN 且不进行增强时，TET 达到 83.17% 的 top-1 准确率，比先前 SOTA 高出超过 10 个点；在使用增强时，83.17% 也被报道。