Differentiable Spike: Rethinking Gradient-Descent for Training Spiking Neural Networks

1 Introduction . Recently , spiking neural networks ( SNNs ) have received increasing attention due to their biologyinspired neural behavior and efficient computation . SNNs deal with binary spike information and therefore enjoy the advantage of multiplication-free inference . Neuromorphic hardware such as TrueNorth [ 1 ] and Loihi [ 2 ] demonstrates that SNNs can save energy by orders of magnitude . Hybrid architecture like Tianjic [ 3 ] suggests its potential power for general intelligence when integrated with traditional artificial infrastructure . However , the bio-mimicry also causes an intrinsic disadvantage in training high-performing SNNs from scratch due to the discrete spike . As a result , obtaining a high-performing SNN has long been a critical problem limiting SNN ’ s deployment in practice . To be more specific , the appealing advantage of the multiplication-free computation in SNN is not recognized in its training [ 4 ] . Modern ML frameworks such as Pytorch [ 5 ] , JAX [ 6 ] , and Tensorflow [ 7 ] do not provide efficient and general instructions or high-level functions to accelerate the convolution with 0/1 spike activation . Moreover , the binary activation in SNN produces allor-nothing gradients [ 8 ] that are incompatible with typical gradient-based optimization methods for efficient neural network training . To circumvent this problem , many works adopt surrogate gradients [ 8 ] ( typically the derivative of a soft-relaxed function ) to replace the spike non-linear gradient . Although the heuristic choice of SG is beneficial to enable the gradient descent in SNN training , the theoretical soundness of SG has not been justified so far . What ’ s more , the gradient descent in SNN training may follow different updating rules from ANN . In fact , since the true gradient ∗Equal Contributions . Corresponding author . Primarily supported by NSF-61876032 . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . value of spike activation does not exist , the adoption of SG raises a new question of how to replace the∞-gradient adaptively through the training process for a good updating strategy of parameters . In this work , we study this specialized problem of adaptive gradient estimation in SNN training . We first revisit the gradient descent algorithm with the finite difference method to accurately depict the loss landscape of adopting a surrogate gradient for the non-differentiable spike firing function . Then we further propose the Differentiable Spike ( Dspike ) function , which can adaptively change its shape and capture the direction of finite difference gradients . With such adaptive surrogate gradients , we establish an efficient algorithm of training SNN with large-scale models . The overall workflow of our algorithm is visualized in Fig . 1 . Our contributions can be summarized as follows : • We identify the problem of adaptively estimating gradient in SNN training and propose a novel analysis framework for SNN ’ s surrogate gradient , through which we are able to not only explain why it works well but also indicate how we can improve its adaptivity to data through training . • We also introduce the novel Dspike function built from the hyperbolic tangent function . The Dspike function can learn its shape accordingly by matching it with the finite difference gradient . • Extensive experiments on both static and dynamic datasets show that our method is highly effective and efficient . For example , our model can achieve 75.4 % top-1 accuracy on the CIFAR-DVS . Our SNN also only consumes 7.8 % energy of the corresponding ANN on CIFAR-10 . 2 Related Work . ANN-SNN Conversion . Generally , the algorithms to obtain a high-performance can be divided into two routes : ( 1 ) ANN-SNN conversion and ( 2 ) direct training SNN from scratch . Conversion-based methods [ 9 , 10 ] utilize the knowledge from ANN and convert the ReLU activation to spike activation mechanism . This type of method can produce an SNN in a short time . For example , Rueckauer et al . [ 11 ] propose to traverse the training dataset to find the maximum activation . Then , the SNN can be converted by directly replacing several modules . In [ 12 ] , maximum activation is abandoned since it may out-lie the central range of activation . The authors choose to find the percentile numbers for threshold voltage to improve the robustness . TSC [ 13 ] adopts the soft-reset ( also known as reset by subtraction ) mechanism and no leakage for conversion , which proves to be effective in conversion . Deng & Gu [ 14 ] decompose the conversion error to each layer , then propose to reduce the error by shifting the bias term . [ 15 ] first propose the calibration to reduce the conversion accuracy . ANN-SNN conversion is fast and accurate . However , achieving near lossless conversion requires a considerable amount of time steps ( > 200 ) to accumulate the spikes , which may significantly add up the latency . Direct Training SNN . In contrast to the conversion method , direct training of SNNs can lower the time step and may further reduce energy consumption . Direct training leverages gradient-descent algorithms . And many works have been developed to tackle the non-differentiability of spike activation . This non-differentiable activation function is similar to the quantization or binarization neural network [ 16 , 17 , 18 ] or adder neural network [ 19 ] . [ 20 ] achieve high performance on CIFAR10 and other neuromorphic datasets . Some works use probabilistic models to approximate the gradient [ 21 ] , which is also developed in binary stochastic model [ 22 ] . Another common approach is to use rate-encoding networks for gradients calculation [ 23 , 24 ] . They quantize the spike rate as the magnitude of information transmitted . ( Wu et al . 2019 [ 25 ] improves the leaky integrate-and-fire ( LIF ) model to an iterative LIF model and develops STBP learning algorithm . Zheng et al . [ 26 ] further propose the tdBN algorithm to balance the gradient norm and smooth the loss function in SNN training , enabling large-scale model optimization . 3 Preliminary . 3.1 Leaky Integrate-and-Fire Model . We use bold italic letters for vector representations . For example , x and y denote the input and target output variables . Bold capital letters like W denote the matrices ( or tensors as clear from the text ) . Constants are denoted by small upright letters , e.g. , a . With bracketed superscript and subscript , we can denote the time dimension and the element indices , respectively . For example , u ( t ) i means the i-th membrane potential at time step t. We adopt the well-known Leaky Integrate-and-Fire ( LIF ) neuron model for a spiking neural network ( SNN ) . Given a membrane potential u ( t ) , the pre-synaptic potential is updated as u ( t+1 ) , pre = τu ( t ) + c ( t+1 ) , where c ( t+1 ) = Wx ( t+1 ) . ( 1 ) Here , τ is a constant leaky factor within ( 0 , 1 ) and c ( t+1 ) is the pre-synaptic input at time step t+ 1 . The input can be calculated from the dot-product between the weights and the spike from the previous layer . Let Vth be the firing threshold , the membrane potential will fire a spike when it exceeds the threshold , and then hard-reset to 0 , given by y ( t+1 ) = { 1 if u ( t+1 ) , pre > Vth 0 otherwise , u ( t+1 ) = u ( t+1 ) , pre · ( 1− y ( t+1 ) ) ( 2 ) After firing , the spike output y ( t+1 ) will propagate to the next layer and become the input x ( t+1 ) of the next layer , ( note that we omit the layer index for simplicity ) . Vth is set to 0.5 in this work . Softmax in SNN . In the general classification model , the final network output will be leveraged to compute the softmax and then to find the winner class . However , the network output is highly variant per input sample , and unlike ReLU activation , the output could be positive or negative . For this reason , we choose to only integrate the network output , and do not fire them across time , as did in recent practices [ 26 , 27 ] . ynet = 1 T T∑ i=1 c ( t ) net = 1 T T∑ i=1 Wx ( t ) , ( 3 ) Then , we can compute the cross-entropy loss based on the ture label and Softmax ( ynet ) . 3.2 Time-dependent Batch Normalization . Batch Normalization ( BN ) [ 28 ] is used to reduce internal covariate shift and smooth the loss landscape [ 29 ] during the optimization of an ANN . It enables the stable training of ANN even with more than 100 layers on the large-scale dataset , e.g . ResNet-152 [ 30 ] . However , BN layers are not designed to normalize spatial-temporal data at first . And simply transplanting BN in SNN training will lead to undesired results . To solve this problem , Zheng et al . [ 26 ] propose time-dependent Batch Normalization ( tdBN ) , which normalizes the data in both spatial and temporal paradigm . tdBN will gather the pre-synaptic input along the time dimension , denoted by C = [ c ( 1 ) , c ( 2 ) , . . . , c ( T ) ] , and then normalize it by C̄ = αVth ( C− E [ C ] ) √ VAR ( C ) + # normalize ( 4 ) Ĉ = γC̄ + β , # affine transformation ( 5 ) where E and VAR compute the mean and variance in channel dimension . γ and β are learnable parameters . Note that in tdBN the normalized tensor will multiply αVth , which can prevent over-fire and under-fire by choosing α wisely . After tdBN , the pre-synaptic input in each time step will be fed into Eq . ( 1 ) . tdBN stabilizes the training of SNN and achieves good results . We also adopt this training technique in our framework . See [ 26 ] for more details of tdBN . 4 Methodology . 4.1 Non-Differentiable Spike . One of the most notorious problems in SNN training is the non-differentiability of the firing function Eq . ( 2 ) . The firing function could be viewed as a variant of sign function , whose derivative is identified as the Dirac delta function . This derivative is 0 almost everywhere except for the threshold that grows to infinity . Seemingly , directly using the Dirac delta function is not an ideal choice for gradient descent-based learning . To demonstrate this problem , let L denote the task loss objective , we can compute the gradient of weights parameters using the spatial-temporal backpropagation ( STBP ) : ∂L ∂W = ∑ t ∂L ∂y ( t ) ∂y ( t ) ∂u ( t ) , pre ∂u ( t ) , pre ∂c ( t ) ∂c ( t ) ∂W . ( 6 ) In above equation , ∂y ( t ) ∂u ( t ) , pre is the gradient of firing function and is 0 almost everywhere , as aforemen- tioned . As a consequence , the gradient descent ( W←W− η ∂L∂W ) either freezes weights or updates the weights to infinity . To mitigate this problem , prior work seeks an alternative gradient ( also known as surrogate gradient ) to circumvent this problem . A popular surrogate gradient may refer to the rectangular function proposed in [ 4 ] , given by ∂y ( t ) ∂u ( t ) , pre = 1 a sign ( ∣∣∣u ( t ) , pre − Vth∣∣∣ < a 2 ) , ( 7 ) where a is a hyper-parameter and usually set to 1 . When Vth = 0.5 , a = 1 , the gradient is simplified to 1 if 0 ≤ u ( t ) , pre ≤ 1 otherwise 0 . This is known as the Straight-Through Estimator in Bengio et al . [ 22 ] and is broadly applied in binary networks [ 31 ] . Another family of surrogate gradient may refer to triangular gradients [ 32 ] , given by ∂y ( t ) ∂u ( t ) , pre = γmax ( 0 , 1− ∣∣∣∣u ( t ) , preVth − 1 ∣∣∣∣ ) . ( 8 ) This surrogate gradient is computed based on the distance between the potential and threshold . Despite various surrogate gradients achieve acceptable accuracy , there lacks a theoretical justification for its reason and effectiveness .

Summary: The authors outline a principled approach--finite difference gradients (FDG)--for choosing a surrogate gradient for spiking neural networks. FDG is demonstrated to approximate gradient descent in classic rate-based ANNs very well, approximate classic rectangular surrogate gradient function descent in SNNs for medium neighborhoods and thus theoretically justify the method, and be useful in finding an optimal neighborhood in which to approximate a gradient for SNNs.  This last step is then used to optimize a continuously varying surrogate spike gradient estimator (Dspike), modeled as a tanh approximation with the gain as a free parameter. Using the chosen Dspike function (which depends on the time step), the authors then demonstrate the appreciable accuracy and energy efficiency improvements of the resultant SNN.