Sparse Spiking Gradient Descent

1 Introduction . In recent years , deep artificial neural networks ( ANNs ) have matched and occasionally surpassed human-level performance on increasingly difficult auditory and visual recognition problems [ 1 , 2 , 3 ] , natural language processing tasks [ 4 , 5 ] and games [ 6 , 7 , 8 ] . As these tasks become more challenging the neural networks required to solve them grow larger and consequently their power efficiency becomes more important [ 9 , 10 ] . At the same time , the increasing interest in deploying these models into embedded applications calls for faster , more efficient and less memory intensive networks [ 11 , 12 ] . The human brain manages to perform similar and even more complicated tasks while only consuming about 20W [ 13 ] which contrasts with the hundreds of watts required for running ANNs [ 9 ] . Unlike ANNs , biological neurons in the brain communicate through discrete events called spikes . A biological neuron integrates incoming spikes from other neurons in its membrane potential and after reaching a threshold emits a spike and resets its potential [ 14 ] . The spiking neuron , combined with the sparse connectivity of the brain , results in a highly spatio-temporally sparse activity [ 15 ] which is fundamental to achieve this level of energy efficiency . Inspired by the extraordinary performance of the brain , neuromorphic computing aims to obtain the same level of energy efficiency preferably while maintaining an accuracy standard on par with ANNs by emulating spiking neural networks ( SNNs ) on specialised hardware [ 16 , 17 , 18 , 19 ] . These efforts go in two directions : emulating the SNNs and training the SNNs . While there is a growing number of successful neuromorphic implementations for the former [ 20 , 21 , 22 , 23 , 24 , 17 , 25 , 26 ] , the latter has proven to be more challenging . Some neuromorphic chips implement local learning rules [ 26 , 27 , 28 ] and recent advances have achieved to approximate backpropagation on-chip [ 29 ] . However , a full end-to-end supervised learning via error backpropagation requires off-chip training [ 22 , 23 , 18 ] . This is usually achieved by simulating and training the entire SNN on a GPU and more 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . recently , emulating the forward pass of the SNN on neuromorphic hardware and then performing the backward pass on a GPU [ 30 , 31 ] . The great flexibility provided by GPUs allows the development of complex training pipelines without being constrained by neuromorphic hardware specifications . However , as GPUs do not leverage the event-driven nature of spiking computations this results in effective but slower and less energy efficient training . An alternative to directly training an SNN consists on training an ANN and converting it to an SNN . However , this approach has been recently shown to not be more energy efficient than using the original ANN in the first place [ 32 ] . Training SNNs has been a challenge itself even without considering a power budget due to the all-or-none nature of spikes which hinders traditional gradient descent methods [ 33 , 34 ] . Since the functional form of a spike is a unit impulse , it has a zero derivative for all membrane potentials except at the threshold where it is infinity . Recently developed methods based on surrogate gradients have been shown to be capable of solving complex tasks to near state-of-the-art accuracies [ 35 , 18 , 36 , 37 , 38 , 39 , 31 ] . The problem of non-differentiability is solved by adopting a well-behaved surrogate gradient when backpropagating the error [ 33 ] . This method , while effective for training very accurate models , is still constrained to working on dense tensors , thus , not profiting from the sparse nature of spiking , and limits the training speed and power efficiency when training SNNs . In this work we introduce a sparse backpropagation method for SNNs . Recent work on ANNs has shown that adaptive dropout [ 40 ] and adaptive sparsity [ 41 ] can be used to achieve state-of-the-art performance at a fraction of the computational cost [ 42 , 43 ] . We show that by redefining the surrogate gradient functional form , a sparse learning rule for SNNs arises naturally as a three-factor Hebbianlike learning rule . We empirically show an improved backpropagation time up to 70x faster than current implementations and up to 40 % more memory efficient in real datasets ( Fashion-MNIST ( FMNIST ) [ 44 ] , Neuromorphic-MNIST ( N-MNIST ) [ 45 ] and Spiking Heidelberg Dataset ( SHD ) [ 46 ] ) thus reducing the computational gap between the forward and backward passes . This improvement will not only impact training for neuromorphic computing applications but also for computational neuroscience research involving training on SNNs [ 39 ] . 2 Sparse Spike Backpropagation . 2.1 Forward Model . We begin by introducing the spiking networks we are going to work with ( Figs . 1 and 2 ) . For simplicity we have omitted the batch dimension in the derivations but it is later used on the complexity analysis and experiments . We have a total of L fully connected spiking layers . Each layer l consisting of N ( l ) spiking neurons which are fully connected to the next layer l+1 through synaptic weights W ( l ) ∈RN ( l ) ×N ( l+1 ) . Each neuron has an internal state variable , the membrane potential , V ( l ) j [ t ] ∈ R that updates every discrete time step t ∈ { 0 , . . . , T −1 } for some finite simulation time T ∈ N. Neurons emit spikes according to a spiking function f : R → { 0 , 1 } such that S ( l ) i [ t ] =f ( V ( l ) i [ t ] ) ( 1 ) The membrane potential varies according to a simplified Leaky Integrate and Fire ( LIF ) neuron model in which input spikes are directly integrated into the membrane [ 14 ] . After a neuron spikes , its potential is reset to Vr . We will work with a discretised version of this model ( see Appendix E for a continuous time version ) . The membrane potential of neuron j in layer l+1 evolves according to the following difference equation V ( l+1 ) j [ t+1 ] = α ( V ( l+1 ) j [ t ] − Vrest ) + Vrest + ∑ i S ( l ) i [ t ] W ( l ) ij − ( Vth − Vr ) S ( l+1 ) j [ t ] ( 2 ) The first two terms account for the leaky part of the model where we define α = exp ( −∆t/τ ) with ∆t , τ ∈ R being the time resolution of the simulation and the membrane time constant of the neurons respectively . The second term deals with the integration of spikes . The last term models the resetting mechanism by subtracting the distance from the threshold to the reset potential . For simplicity and without loss of generality , we consider Vrest=0 , Vth=1 , Vr=0 and V ( l ) i [ 0 ] =0 ∀ l , i from now on . We can unroll ( 2 ) to obtain V ( l+1 ) j [ t+1 ] = ∑ i t∑ k=0 αt−kS ( l ) i [ k ] ︸ ︷︷ ︸ Input trace W ( l ) ij ︸ ︷︷ ︸ Weighted input − t∑ k=0 αt−kS ( l+1 ) j [ k ] ︸ ︷︷ ︸ Resetting term ( 3 ) Thus , a spiking neural network can be viewed as a special type of recurrent neural network where the activation function of each neuron is the spiking function f ( · ) . The spiking function , is commonly defined as the unit step function centered at a particular threshold Vth f ( v ) = { 1 , v > Vth 0 , otherwise ( 4 ) Note that while this function is easy to compute in the forward pass , its derivative , the unit impulse function , is problematic for backpropagation which has resulted in adopting surrogate derivatives [ 33 ] to allow the gradient to flow . Interestingly , it has been shown that surrogate gradient descent performs robustly for a wide range of surrogate functional forms [ 38 ] . After the final layer , a loss function loss ( · ) computes how far the network activity is with respect to some target value for the given input . The loss function can be defined as a function of the network spikes , membrane potentials or both . It may also be evaluated at every single time step or every several steps . We deliberately leave the particular definition of this function open as it does not directly affect the sparse gradient descent method we introduce . 2.2 Backward Model . The SNN is updated by computing the gradient of the loss function with respect to the weights in each layer . This can be achieved using backpropagation through time ( BPTT ) on the unrolled network . Following ( 1 ) and ( 3 ) we can derive weight gradients to be ∇W ( l ) ij = ∑ t ε ( l+1 ) j [ t ] ︸ ︷︷ ︸ Gradient from next layer Spike derivative︷ ︸︸ ︷ dS ( l+1 ) j [ t ] dV ( l+1 ) j [ t ] ( ∑ k < t αt−k−1S ( l ) i [ k ] ) ︸ ︷︷ ︸ Input trace ( 5 ) This gradient consists of a sum over all time steps of the product of three terms . The first term , modulates the spatio-temporal credit assignment of the weights by propagating the gradient from the next layer . The second term , is the derivative of f ( · ) evaluated at a given Vj [ t ] . The last term is a filtered version of the input spikes . Notice that it is not necessary to compute this last term as this is already computed in the forward pass ( 3 ) . We have chosen not make the spike resetting term in ( 3 ) differentiable as it has been shown that doing so can result in worse testing accuracy [ 38 ] and it increases the computational cost . For a typical loss function defined on the output layer activity , the value of ε ( L−1 ) j [ t ] ( i.e . in the last layer ) is the gradient of the loss with respect to the output layer ( see Fig . 1 ) . For all other layers we have that ε ( l ) j [ t ] = ∇S ( l ) j [ t ] , that is , the gradient of the loss function with respect to the output spikes of layer l. We can obtain an expression for ∇S ( l ) j [ t ] as a function of the gradient of the next layer : ε ( l ) i [ t ] = ∇S ( l ) i [ t ] = ∑ j Wij ∑ k > t Gradient from next layer︷ ︸︸ ︷ ∇S ( l+1 ) j [ k ] Spike derivative︷ ︸︸ ︷ dS ( l+1 ) j [ k ] dV ( l+1 ) j [ k ] αk−t−1 ( 6 ) = ∑ j Wijδ ( l+1 ) j [ t ] , l = { 0 , . . . L− 2 } All hidden layers need to compute equations ( 5 ) and ( 6 ) . The first one to update their input synaptic weights and the second one to backpropagate the loss to the previous layer . The first layer ( l=1 ) only needs to compute ( 5 ) and the last layer ( l=L−1 ) will compute a different ε ( L−1 ) j [ t ] depending on the loss function . We introduce the following definitions which result in the sparse backpropagation learning rule . Definition 1 . Given a backpropagation threshold Bth∈R we say neuron j is active at time t iff |Vj [ t ] − Vth| < Bth ( 7 ) Definition 2 . The spike gradient is defined as dSj [ t ] dVj [ t ] : = { g ( Vj [ t ] ) , if Vj [ t ] is active 0 , otherwise ( 8 ) This means that neurons are only active when their potential is close to the threshold . Applying these two definitions to ( 5 ) and ( 6 ) results in the gradients only backpropagating through active neurons at each time step as shown in Fig . 2 . The consequences of this are readily available for the weight gradient in ( 5 ) as the only terms in the sum that will need to be computed are those in which the postsynaptic neuron j was active at time t. Resulting in the following gradient update : ∇W ( l ) ij [ t ] = ε ( l+1 ) j [ t ] dS ( l+1 ) j [ t ] dV ( l+1 ) j [ t ] ( ∑ k < t α t−k−1S ( l ) i [ t ] ) V ( l+1 ) j [ t ] is active , 0 , otherwise ( 9 ) ∇W ( l ) ij = ∑ t ∇W ( l ) ij [ t ] ( 10 ) For the spike gradient in ( 6 ) we have two consequences , firstly , we will only need to compute ∇S ( l ) j [ t ] for active neurons since ∇S ( l ) j [ t ] is always multiplied by dS ( l ) j [ t ] dV ( l ) j [ t ] in ( 5 ) and ( 6 ) . Secondly , we can use a recurrent relation to save time and memory when computing δj [ t ] in ( 6 ) as shown in the following proposition . Proposition 1 . We can use a recurrent relation to compute δj [ t ] given by δj [ t ] = { αnδj [ t+ n ] , if dSj [ k ] dVj [ k ] =0 , for t+1≤k≤ t+n ∇Sj [ t+1 ] dSj [ t+1 ] dVj [ t+1 ] +α nδj [ t+n ] , if dSj [ k ] dVj [ k ] =0 , for t+1 < k≤ t+n , dSj [ t+1 ] dVj [ t+1 ] ̸=0 ( 11 ) This means that we only need to compute δ ( l+1 ) j [ t ] at those time steps in which either dSj [ t+1 ] dVj [ t+1 ] ̸= 0 or V ( l ) j [ t ] is active ( see Appendix C for a visualisation of this computation ) . Thus , we end up with the following sparse expression for computing the spike gradient . ∇S ( l ) j [ t ] = { ∑ j Wijδ ( l+1 ) j [ t ] , V ( l ) j [ t ] is active 0 , otherwise ( 12 )

Summary: This paper proposes a novel sparse gradient update rule to train spiking neural networks faster and more energy-efficient on GPUs. A mathematical derivation of backpropagation through time (BPTT) is given for a network of leaky integrate-and-fire (LIF), feed-forward, fully-connected layers without skip connections. A threshold is defined that marks neurons as active when they are close to the membrane threshold and the spike gradient is only propagated back if a neuron was active at that time during the forward pass. As a result, only a fraction of the actual update steps have to be computed, because all others would be zero anyway. The authors implement a version of their algorithm in CUDA to compare their sparse gradient update rule to the dense update rule from surrogate gradients. It is found, that on the three evaluated datasets, Fashion-MNIST, N-MNIST and Spiking Heidelberg Digits, a consistent improvement in memory efficiency and training speed can be measured while achieving similar accuracies.