Direct Training Needs Regularisation: Anytime Optimal Inference Spiking Neural Network

TL;DR

The paper tackles the challenge of training Spiking Neural Networks (SNNs) for Anytime Optimal Inference (AOI), where reliable predictions should be maintained across varying timesteps. It introduces Spatial-Temporal Factor (STF) and Spatial-Temporal Regulariser (STR) to regulate the balance between spike activity and membrane potential at each timestep, enabling AOI through direct training, and combines this with Temporal Efficient Training (TET) loss and a softmax-based cutoff. Across frame-based and event-based datasets, STR reduces prediction uncertainty, lowers latency, and maintains or improves accuracy; with cutoff, the approach achieves 2.14 to 2.89× faster inference with only 0.50% to 0.64% accuracy loss on event-based data. The work offers a practical pathway to fast, reliable, and energy-efficient SNNs suitable for real-time neuromorphic applications.

Abstract

Spiking Neural Network (SNN) is acknowledged as the next generation of Artificial Neural Network (ANN) and hold great promise in effectively processing spatial-temporal information. However, the choice of timestep becomes crucial as it significantly impacts the accuracy of the neural network training. Specifically, a smaller timestep indicates better performance in efficient computing, resulting in reduced latency and operations. While, using a small timestep may lead to low accuracy due to insufficient information presentation with few spikes. This observation motivates us to develop an SNN that is more reliable for adaptive timestep by introducing a novel regularisation technique, namely Spatial-Temporal Regulariser (STR). Our approach regulates the ratio between the strength of spikes and membrane potential at each timestep. This effectively balances spatial and temporal performance during training, ultimately resulting in an Anytime Optimal Inference (AOI) SNN. Through extensive experiments on frame-based and event-based datasets, our method, in combination with cutoff based on softmax output, achieves state-of-the-art performance in terms of both latency and accuracy. Notably, with STR and cutoff, SNN achieves 2.14 to 2.89 faster in inference compared to the pre-configured timestep with near-zero accuracy drop of 0.50% to 0.64% over the event-based datasets. Code available: https://github.com/Dengyu-Wu/AOI-SNN-Regularisation

Figures (7)

• Figure 1: (a) Forward propagation in SNN. The input events $\boldsymbol{X}(t)$ stimulate neurons to generate spikes over the time. The output $f(\boldsymbol{X}(t))$ can respond when a sufficient number of events are received within a specific time window. $\boldsymbol{\theta}(t)$ and $\tau\boldsymbol{\Delta}(t)$ represents the spatial and temporal information at $t$, respectively. (b) The state update of one LIF neuron at input layer during forward propagation process. The weight $W_i^l$ influences the current contributing to the membrane potential $V_i^l(t)$. The threshold $V^l_{thr}$ determines the threshold level of $V_i^l(t)$ required to generate the spikes.
• Figure 2: Comparison of accuracy with respect to timestep using different loss functions on Cifar10-DVS.
• Figure 3: Visualisation of STF $\xi^l(t)$ at 8-th layer in two models, before and after regularisation, on Cifar10-DVS. Regularisation leads to a reduction in variance of $\xi^l(t)$ from 0.0025 to 0.0022, indicating enhanced stability across timestep. Additionally, the mean value of $\xi^l(t)$ rises from 0.2736 to 0.3336, reflecting an enhancement in the representation of spatial information $\lVert\theta^l(t)\rVert_2$.
• Figure 4: Comparision of average STF $\bar{\xi}^l$ (over timestep) on Cifar10-DVS in distinguishing correct vs. wrong predictions. (a) illustrates the $\bar{\xi}^l$ using the baseline method, while (b) showcases the $\bar{\xi}^l$ after applying the regularisation technique. In both cases, the left side shows the distribution of STF over different layers, and the right side displays the average STF values for correct and wrong predictions, alongside the gap value between them.
• Figure 5: Comparison of uncertainty with respect to timestep on six datasets. The assessment of uncertainty is conducted on various models with distinct settings of $\alpha$, such as {0.05, 0.1} for Cifar10/100 and {0.3, 0.5} for event-based inputs. Due to the expensive training for ImageNet, we solely evaluate the setting of {0.05}.