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

TL;DR

This paper introduces a novel temporal efficient training (TET) method for spiking neural networks that improves training convergence, generalizability, and scalability, achieving state-of-the-art results on multiple datasets including a significant accuracy boost on DVS-CIFAR10.

Contribution

The paper proposes the TET approach that enhances surrogate gradient training of SNNs by improving convergence and generalizability, and introduces temporal inheritable training for acceleration.

Findings

TET outperforms existing methods on CIFAR-10/100 and ImageNet.

Achieved 83% top-1 accuracy on DVS-CIFAR10, over 10% improvement.

TET improves temporal scalability and training efficiency of SNNs.

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