One Timestep Is All You Need: Training Spiking Neural Networks with Ultra Low Latency

Sayeed Shafayet Chowdhury, Nitin Rathi, Kaushik Roy

Keywords: Spiking Neural Networks, One timestep Inference, Iterative Initialization and Retraining, Ultra-high energy efficiency

Abstract:

Spiking Neural Networks (SNNs) can be energy efficient alternatives to commonly used deep neural networks (DNNs). Through event-driven information processing, SNNs can considerably reduce the compute requirements of DNNs. However, high inference latency is a significant hindrance to their deployment. Computation over multiple timesteps increases latency and incurs memory access overhead of fetching membrane potentials, both of which lessen the energy benefits of SNNs. Hence, latency reduction is pivotal to obtain SNNs
with high energy efficiency. However, reducing latency can have an adverse effect on accuracy. To obtain solutions which optimize the accuracy-energy-latency trade-off, we propose an iterative training method which starts with an SNN of T (T>1) timesteps, and reduces T every iteration of training, with neuron threshold and leak as trainable parameters. This results in a continuum of SNNs, starting from an SNN trained with T timesteps, all the way up to unit latency. We use direct input encoding (analog inputs from pixels) with the first
convolutional layer of the network of leaky integrate and fire (LIF) neurons acting as spike generator. We choose T=5 as our starting point, since it is the minimum reported latency to achieve satisfactory performance on ImageNet. Training SNNs directly with 1 timestep results in convergence failure due to layerwise spike vanishing and difficulty in finding optimum thresholds. The proposed iterative
training approach overcomes this through enabling the learning of suitable layerwise thresholds with backpropagation by maintaining sufficient spiking activity, starting from T timesteps up to 1. Using the proposed training algorithm, we achieve top-1 accuracy of 93.05%, 70.15% and 67.71% on CIFAR-10, CIFAR-100 and ImageNet, respectively with VGG16, in just 1 timestep. Compared to a
5 timestep SNN, the 1 timestep SNN achieves ~5X enhancement in efficiency, with an accuracy drop of ~1%. In addition, 1 timestep SNNs perform inference with 5X reduced latency compared to state-of-the-art SNNs, and provide 25-33X higher energy efficiency compared to DNNs, while being comparable to them in performance. The proposed method also enables training reinforcement learning
agents on Cartpole and Atari pong environments which infer using 1 timestep.

One-sentence Summary: We propose an iterative initialization and retraining method for SNNs, enabling SNNs to infer in just 1 timestep while maintaining comparable or even better accuracy compared to other state-of-the-art SNNs.