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: This paper studies the differentiable training problem of spiking neural network (SNN), which is a promising alternative to the widely used ANN. The authors first identify the problem of the current surrogate gradient of SNN and propose a new surrogate gradient estimator. I really like the author's research method of thinking about problems from experiments, but the result of the experiment is just a tuning experiment of some parameters, and I did not see the author's insight into the existing SNN training problems. Furthermore, the Spike method proposed by the author lacks clear motivation. The proposed Dspike method is just a differentiable piecewise function, which does not have much methodological contribution. The author exaggerated the experimental settings and did not present the experimental results of neuromorphic images. In summary, I think this article does not meet the standards of the NeurIPS conference.  The author’s reply addressed most of my concerns. I agree that AdderNet is another line of work with different purpose. I can see that the authors are very thoughtful in the model design, however it is not well-described in the main paper. I have raised my rating by 1.