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 work introduces a new method to train spiking neural networks using far fewer gradients during the backwards back-propagation pass.  By only retaining the potential of so-called "active" neurons (those within a threshold epsilon of spiking) and learning through a surrogate function, these networks of spiking neurons can be sparsely trained (requiring only around 1% of the gradients).  The authors test these claims on three datasets (a time-encoded version of Fashion-MNIST, N-MNIST, and Spiking Heidelberg Digits), and achieve speedups of up to 70x and memory saving of 40% while maintaining approximately the same accuracy.  Further contributions interpret the learning rule from a Hebbian perspective and establish the savings in computational complexity.