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: This paper proposes a new training approach, the temporal efficient training (TET). This algorithm utilizes a new loss function to improve the generalizability of SNNs. Further, a new training pipeline is presented to reduce the simulation time of SNNs. This work outperforms the SOTA on the static datasets and neuromorphic datasets.