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

1 INTRODUCTION . The advantages of Spiking neuron networks ( SNNs ) lie in their energy-saving and fast-inference computation when embedded on neuromorphic hardware such as TrueNorth ( DeBole et al. , 2019 ) and Loihi ( Davies et al. , 2018 ) . Such advantages originate from the biology-inspired binary spike transmitted mechanism , by which the networks avoid multiplication during inference . On the other hand , this mechanism also leads to difficulty in training very deep SNNs from scratch because the non-differentiable spike transmission hinders the powerful back-propagation approaches like gradient descents . Recently , many studies on converting artificial neuron networks ( ANNs ) to SNNs have demonstrated SNNs ’ comparable power in feature representation as ANNs ( Han & Roy , 2020 ; Deng & Gu , 2020 ; Li et al. , 2021a ) . Nevertheless , it is commonly agreed that the direct training method for high-performance SNN is still crucial since it distinguishes SNNs from converted ANNs , especially on neuromorphic datasets . The output layer ’ s spike frequency or the average membrane potential increment is commonly used as inference indicators in SNNs ( Shrestha & Orchard , 2018 ; Kim et al. , 2019 ) . The current standard direct training ( SDT ) methods regard the SNN as RNN and optimize inference indicators ’ distribution ( Wu et al. , 2018 ) . They adopt surrogate gradients ( SG ) to relieve the non-differentiability ( Lee et al. , 2016 ; Wu et al. , 2018 ; Zheng et al. , 2021 ) . However , the gradient descent with SG does not match with the loss landscape in SNN and is easy to get trapped in a local minimum with low generalizability . Although using suitable optimizers and weight decay help ease this problem , the performance of deep SNNs trained from scratch still suffers a big deficit compared to that of ANNs Deng et al . ( 2020 ) . Another training issue is the memory and time consumption , which increases B Corresponding author Standard Direct Training linearly with the simulation time . Rathi & Roy ( 2020 ) initializes the target network by a converted SNN to shorten the training epochs , indicating the possibility of high-performance SNN with limited activation time . The training problem due to the non-differentiable activation function has become the main obstruction of spiking neural network development . In this work , we examine the limitation of the traditional direct training approach with SG and propose the temporal efficient training ( TET ) algorithm . Instead of directly optimizing the integrated potential , TET optimizes every moment ’ s pre-synaptic inputs . As a result , it avoids the trap into local minima with low prediction error but a high second-order moment . Furthermore , since the TET applies optimization on each time point , the network naturally has more robust time scalability . Based on this characteristic , we propose the time inheritance training ( TIT ) , which reduces the training time by initializing the SNN with a smaller simulation length . With the help of TET , the performance of SNNs has improved on both static datasets and neuromorphic datasets . Figure 1 depicts the workflow of our approach . The following summarizes our main contributions : • We analyze the problem of training SNN with SG and propose the TET method , a new loss and gradient descent regime that succeeds in obtaining more generalizable SNNs . • We analyze the feasibility of TET and picture the loss landscape under both the SDT and TET setups to demonstrate TET ’ s advantage in better generalization . • Our sufficient experiments on both static datasets and neuromorphic datasets prove the effectiveness of the TET method . Especially on DVS-CIFAR10 , we report 83.17 % top-1 accuracy for the first time , which is over 10 % better than the current state-of-the-art result . 2 RELATED WORK . In recent years , SNNs have developed rapidly and received more and more attention from the research community . However , lots of challenging problems remain to be unsolved . In general , most works on SNN training have been carried out in two strategies : ANN-to-SNN conversion and direct training from scratch . ANN-to-SNN Conversion . Conversion approaches avoid the training problem by trading high accuracy through high latency . They convert a high-performing ANN to SNN and adjust the SNN parameters w.r.t the ANN activation value layer-by-layer ( Diehl et al. , 2015 ; 2016 ) . Some special techniques have been proposed to reduce the inference latency , such as the subtraction mechanism ( Rueckauer et al. , 2016 ; Han et al. , 2020 ) , robust normalization Rueckauer et al . ( 2016 ) , spike-norm ( Sengupta et al. , 2018 ) , and channel-wise normalization ( Kim et al. , 2019 ) . Recently , Deng & Gu ( 2020 ) decompose the conversion error to each layer and reduce it by bias shift . Li et al . ( 2021a ) suggest using adaptive threshold and layer-wise calibration to obtain high-performance SNNs that require a simulation length of less than 50 . However , converted methods significantly extend the inference latency , and they are not suitable for neuromorphic data ( Deng et al. , 2020 ) . Direct training . In this area , SNNs are regarded as special RNNs and training with BPTT ( Neftci et al. , 2019 ) . On the backpropagation process , The non-differentiable activation term is replaced with a surrogate gradient ( Lee et al. , 2016 ) . Compared with ANN-to-SNN conversion , direct training achieves high accuracy with few time steps but suffers more training costs ( Deng et al. , 2020 ) . Several studies suggest that surrogate gradient ( SG ) is helpful to obtain high-performance SNNs on both static datasets and neuromorphic datasets ( Wu et al. , 2019 ; Shrestha & Orchard , 2018 ; Li et al. , 2021b ) . On the backpropagation process , SG replaces the Dirac function with various shapes of curves . Exceptionally , Wu et al . ( 2018 ) first propose the STBP method and train SNNs on the ANN programming platform , which significantly promotes direct training development . Zheng et al . ( 2021 ) further proposes the tdBN algorithm to smooth the loss function and first realize training a large-scale SNN on ImageNet . Zhang & Li ( 2020 ) proposes TSSL-BP to break down error backpropagation across two types of inter-neuron and intra-neuron dependencies and achieve low-latency and high accuracy SNNs . Recently , Yang et al . ( 2021 ) designed a neighborhood aggregation ( NA ) method to use the multiple perturbed membrane potential waveforms in the neighborhood to compute the finite difference gradients and guide the weight updates . They significantly decrease the required training iterations and improve the SNN performance . 3 PRELIMINARY . 3.1 ITERATIVE LIF MODEL . We adopt the Leaky Integrate-and-Fire ( LIF ) model and translate it to an iterative expression with the Euler method ( Wu et al. , 2019 ) . Mathematically , the membrane potential is updated as u ( t+ 1 ) = τu ( t ) + I ( t ) , ( 1 ) where τ is the constant leaky factor , u ( t ) is the membrane potential at time t , and I ( t ) denotes the pre-synaptic inputs , which is the product of synaptic weight W and spiking input x ( t ) . Given a specific threshold Vth , the neuron fires a spike and u ( t ) reset to 0 when the u ( t ) exceeds the threshold . So the firing function and hard reset mechanism can be described as a ( t+ 1 ) = Θ ( u ( t+ 1 ) − Vth ) ( 2 ) u ( t+ 1 ) = u ( t+ 1 ) · ( 1− a ( t+ 1 ) ) , ( 3 ) where Θ denotes the Heaviside step function . The output spike a ( t + 1 ) will become the post synaptic spike and propagate to the next layer . In this study , we set the starting membrane u ( 0 ) to 0 , the threshold Vth to 1 , and the leaky factor τ to 0.5 for all experiments . The last layer ’ s spike frequency is typically used as the final classification index . However , adopting the LIF model on the last layer will lose information on the membrane potential and damage the performance , especially on complex tasks ( Kim et al. , 2019 ) . Instead , we integrate the pre-synaptic inputs I ( t ) with no decay or firing ( Rathi & Roy , 2020 ; Fang et al. , 2021 ) . Finally , we set the average membrane potential as the classification index and calculate the cross-entropy loss for training . 3.2 SURROGATE GRADIENT . Following the concept of direct training , we regard the SNN as RNN and calculate the gradients through spatial-temporal backpropagation ( STBP ) ( Wu et al. , 2018 ) : ∂L ∂W = ∑ t ∂L ∂a ( t ) ∂a ( t ) ∂u ( t ) ∂u ( t ) ∂I ( t ) ∂I ( t ) ∂W , ( 4 ) where the term ∂a ( t ) ∂u ( t ) is the gradient of the non-differentiability step function involving the derivative of Dirac ’ s δ-function that is typically replaced by surrogate gradients with a derivable curve . So far , there are various shapes of surrogate gradients , such as rectangular ( Wu et al. , 2018 ; 2019 ) , triangle ( Esser et al. , 2016 ; Rathi & Roy , 2020 ) , and exponential ( Shrestha & Orchard , 2018 ) curve . In this work , we choose the surrogate gradients shaped like triangles . Mathematically , it can describe as ∂a ( t ) ∂u ( t ) = 1 γ2 max ( 0 , γ − |u ( t ) − Vth| ) , ( 5 ) where the γ denotes the constraint factor that determines the sample range to activate the gradient . 3.3 BATCH NORMALIZATION FOR SNN . Batch Normalization ( BN ) ( Ioffe & Szegedy , 2015 ) is beneficial to accelerate training and increase performance since it can smooth the loss landscape during training ( Santurkar et al. , 2018 ) . Zheng et al . ( 2021 ) modified the forward time loop form and proposed threshold-dependent Batch Normalization ( tdBN ) to normalize the pre-synaptic inputs I in both spatial and temporal paradigms so that the BN can support spatial-temporal input . We adopt this setup with the extension of the time dimension to batch dimension 1 . In the inference process , the BN layer will be merged into the pre-convolutional layer , thus the inference rule of SNN remain the same but with modified weight : Ŵ←W γ α , b̂← β + ( b− µ ) γ α , ( 6 ) where µ , α are the running mean and standard deviation on both spatial and temporal paradigm , γ , β are the affine transformation parameters , and W , b are the parameters of the pre-convolutional layer . 4 METHODOLOGY . 4.1 FORMULA OF TRAINING SNN WITH SURROGATE GRADIENTS . Standard Direct Training . We use O ( t ) to represent pre-synaptic input I ( t ) of the output layer and calculate the cross-entropy loss . The loss function of standard direct training LSDT is : LSDT = LCE ( 1 T T∑ t=1 O ( t ) , y ) , ( 7 ) where T is the total simulation time , LCE denotes the cross-entropy loss , and y represents the target label . Following the chain rule , we obtain the gradient of W with softmax S ( · ) inference function : ∂LSDT ∂W = 1 T T∑ t=1 [ S ( Omean ) − ŷ ] ∂O ( t ) ∂W , ( 8 ) where Omean denotes the average of the output O ( t ) over time , and ŷ is the one-hot coding of y. Temporal Efficient Training . In this section , we come up with a new kind of loss function LTET to realize temporal efficient training ( TET ) . It constrains the output ( pre-synaptic inputs ) at each moment to be close to the target distribution . It is described as : LTET = 1 T · T∑ t=1 LCE [ O ( t ) , y ] . ( 9 ) Recalculate the gradient of weights under the loss function LTET , and we have : ∂LTET ∂W = 1 T T∑ t=1 [ S ( O ( t ) ) − ŷ ] · ∂O ( t ) ∂W . ( 10 )

Summary: The authors propose the temporal efficient training (TET) method that features loss evaluation at every timestep. As such, the main difference between TET and SDT lies in the fact that TET evaluates loss at every timestep with the correct label vector whereas SDT evaluates loss at the last timestep. But unfortunately, this could not be done because of the use of $L_{TOTAL}=(1-\lambda)L_{TET}+\lambda L_{MSE}$, where $L_{MSE}$ is evaluated at the last moment. The proposed method is simple but its efficacy outperforms the SOTA results for various networks on various datasets.