Surrogate gradient scaling for directly training spiking neural networks

Abstract

Spiking neural networks (SNNs) are considered to be biologically plausible and can yield high energy efficiency when implemented on neuromorphic hardware due to their highly sparse asynchronous binary event-driven nature. Recently, surrogate gradient (SG) approaches have enabled SNNs to be trained from scratch with backpropagation (BP) algorithms under a deep learning framework. However, a popular SG approach known as straight-through estimator (STE), which only propagates the same gradient information, does not take into account the activation differences between the membrane potentials and output spikes. To address this issue, we propose surrogate gradient scaling (SGS), which scales up or down the gradient information of the membrane potential according to the sign of the gradient of the spiking neuron output and the difference between the membrane potential and the output of the spiking neuron. This SGS approach can also be applied to unimodal functions that propagate different gradient information from the output spikes to the input membrane potential. In addition, SNNs trained directly from scratch suffer from poor generalization performance, and we introduce Lipschitz regularization (LR), which is incorporated into the loss function. It not only improves the generalization performance of SNNs but also makes them more robust to noise. Extensive experimental results on several popular benchmark datasets (CIFAR10, CIFAR100 and CIFAR10-DVS) show that our approach not only outperforms the SOTA but also has lower inference latency. Remarkably, our SNNs can lead to 34$\times$, 29$\times$, and 17$\times$ computation energy savings compared to standard Artificial neural networks (ANNs) on above three datasets.

Appendix

1.1 Detailed Network architecture

The key component residual block in ResNet-19 is shown in Fig. 9. Conv is the convolutional layer, BN denotes batch normalization, and LIF denotes the LIF neurons. I[t] is the input spike train and O[t] is the output spike train. The detailed network architecture of ResNet-19 is shown in Table 7. $C_{in}$ denote the input channels and $C_{out}$ denotes the output channels. Block indicates the residual block shown in Fig. 9. Block4 and Block7 indicate that the stride of the first convolution layer of the residual block is 2 and the stride of the second convolution layer remains 1.

Fig. 9

Residual block of Spiking ResNet-19

Table 7 The detailed architecture of ResNet-19

1.2 Energy efficiency calculation

Here we refer to [10] to give calculations on the power consumption of DNNs and SNNs. The energy consumption in DNNs is mainly focused on a large number of MAC operations, while the computation of SNNs is performed on binary events, and the energy consumption is mainly focused on AC operations. In this paper, we specify that each MAC and AC operation is performed in register transfer logic on 45nm CMOS technology. Considering 32-bit weight values, the energy consumed by the MAC operation and AC operation for a 32-bit integer is 3.2pJ and 0.1pJ, respectively. The floating point operations (FLOPS) of DNNs and SNNs can be obtained according to:

$$\begin{aligned} {F_{DNNs}} = \left\{ {\begin{array}{*{10}{l}} {{O^2}*N*{k^2}*M,Conv}\\ {{D_{in}} * {D_{out}},FC} \end{array}} \right. \end{aligned}$$

(17)

$$\begin{aligned} {F_{SNNs}} = \left\{ {\begin{array}{*{10}{l}} {{O^2} * N*{k^2} * M * {S_A},Conv}\\ {{D_{in}} * {D_{out}} * {S_A},FC} \end{array}} \right. \end{aligned}$$

(18)

As for the convolutional layer, N denotes the input channels, M denotes the output channels, the input size is $I \times I$, the kernel size is $k \times k$, and the output size is $O \times O$. For the fully connected layer, $D_{in}$ denotes the size of the input, and $D_{out}$ denotes the size of the output. $S_A$ denotes the spike rate of each layer. Then, the inference energy of SNNs and DNNs can be calculated from:

$$\begin{aligned} {E_{DNNs}}= & {} (\sum \limits _{i = 1}^N {{F_{DNNs}}} ) * {E_{MAC}}\end{aligned}$$

(19)

$$\begin{aligned} {E_{SNNs}}= & {} (\sum \limits _{i = 1}^N {{F_{SNNs}}} ) * {E_{AC}} * T \end{aligned}$$

(20)

where $E_{MAC}$ is 3.2 pJ, $E_{AC}$ is 0.1 pJ, and T is time steps. The inference energy is considering FLOPS across all N layers according to (19) and (20).