Temporal Efficient Training of Spiking Neural Network via Gradient Re-weighting

TL;DR

This work tackles the poor generalization of directly trained spiking neural networks (SNNs) trained with surrogate gradients by introducing Temporal Efficient Training (TET), which optimizes per-time-step outputs and reweights gradient contributions to favor flatter minima. It also introduces Time Inheritance Training (TIT) to accelerate training by progressively increasing simulation length. Across CIFAR-10/100, ImageNet, and neuromorphic DVS-CIFAR10, TET/TIT achieves state-of-the-art results, notably 83.17% top-1 on DVS-CIFAR10, while reducing training time and improving temporal scalability.

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}.

Figures (8)

• Figure 1: Workflow of temporal efficient training (TET). To obtain a more generalized SNN, we modify the optimization target to adjust each moment's output distribution.
• Figure 2: Loss landscape of VGGSNN. The 2D landscape of $\mathcal{L}_\text{SDT}$ and $\mathcal{L}_\text{TET}$ from two different training methods.
• Figure 3: TET helps to jump out the local minimum point. We provide the test accuracy (A) and loss (B) change after changing the SDT to TET at epoch 200. TET efficiently improves the test performance and reduces the two kinds of loss.
• Figure 4: Time scalability robustness and network efficiency of ResNet-19 on CIFAR100. (A) The comparison of training from scratch (dots) and inheriting from a small simulation length (lines). (B) SNN network performance changes with energy consumption.
• Figure 5: STD loss landscape of ResNet-19 on CIFAR100 from different training approaches.

Theorems & Definitions (2)

• Lemma 4.1
• proof