SPIDE: A Purely Spike-based Method for Training Feedback Spiking Neural Networks

1 INTRODUCTION . Spiking neural networks ( SNNs ) are brain-inspired models that transmit spikes between neurons for event-driven energy-efficient computation . SNNs can be implemented with less energy on neuromorphic hardware ( Akopyan et al. , 2015 ; Davies et al. , 2018 ; Pei et al. , 2019 ; Roy et al. , 2019 ) , which can remedy the defects of large energy consumption of artificial neural networks ( ANNs ) . Different from ANNs , however , directly supervised training of SNNs is a hard problem due to the complex spiking neuron model which is discontinuous . To solve this problem , converting ANNs to SNNs ( Hunsberger & Eliasmith , 2015 ; Rueckauer et al. , 2017 ; Sengupta et al. , 2019 ; Rathi et al. , 2019 ; Deng & Gu , 2021 ; Yan et al. , 2021 ) , or many other direct SNN training methods ( Wu et al. , 2018 ; Bellec et al. , 2018 ; Jin et al. , 2018 ; Shrestha & Orchard , 2018 ; Wu et al. , 2019 ; Neftci et al. , 2019 ; Zhang & Li , 2019 ; Kim et al. , 2020 ; Zheng et al. , 2021 ; Bohte et al. , 2002 ; Zhang & Li , 2020 ; Kim et al. , 2020 ; Xiao et al. , 2021 ) have been proposed . While these methods could partly tackle the problems of unsatisfactory performance or high latency , they require complex computation for gradient calculation or approximation , which can not be implemented by common spiking neurons on neuromorphic hardware . They aim at training SNNs on commonly used computational units , e.g . GPU , and deploying trained models for energy-efficient inference . However they do not consider if the training procedure could leverage the same spike-based computation for gradient calculation and training to reduce the large energy consumption during training as well . A few previous works try to train SNNs with spikes ( Guerguiev et al. , 2017 ; Neftci et al. , 2017 ; Samadi et al. , 2017 ; O ’ Connor & Welling , 2016 ; Thiele et al. , 2019b ; a ) . They either are based on direct feedback alignment ( DFA ) ( Nøkland , 2016 ) and performs poorly , or require impractical special neuron models ( Thiele et al. , 2019b ; a ) . Besides , they only focus on feedforward network structures imitated from ANNs , which ignores feedback connections that are ubiquitous in the human brain and enable neural networks to be shallower and more efficient ( Kubilius et al. , 2019 ; Xiao et al. , 2021 ) . Actually , feedback structures suit SNNs more since SNNs will naturally compute with multiple time steps , which could reuse representations and avoid uneconomical costs to unfold along time that ANNs suffer from ( Xiao et al. , 2021 ) . So training algorithms for feedback SNNs , which may also degrade to feedforward structures by taking feedback as zero , is worth more exploration . An ideal SNN training method should tackle the common problems , be suitable for flexible structures ( feedforward or feedback ) and be spike-based with high neuromorphic plausibility . The implicit differentiation on the equilibrium state ( IDE ) method ( Xiao et al. , 2021 ) , which is recently proposed to train feedback spiking neural networks ( FSNNs ) , is a promising method that may generalize to spike-based learning for requirement . They derive that the forward computation of FSNNs converges to an equilibrium state , which follows a fixed-point equation . Based on it , they propose to train FSNNs by implicit differentiation on this equation , which tackles the common difficulties for SNN training including non-differentiability and large memory costs , and has interesting local update properties . In their method , however , they leverage general root-finding methods to solve implicit differentiation , which requires complex computation on standard computation systems . In this work , we extend the IDE method to spike-based IDE ( SPIDE ) , which fulfills our requirements and has great potential for energy-efficient training of SNNs on neuromorphic hardware , by introducing ternary spiking neuron couples and proposing to solve implicit differentiation by spikes based on them . Our method is also applicable to feedforward structures by degrading the feedback connection as zero . In practice , however , it may require long time steps to stabilize the training with spikes due to approximation error for gradients . So we further dive into the approximation error from the statistical perspective , and propose to simply adjust the resting potential of SNNs to achieve an unbiased estimation of gradients and reduce the estimation variance of SNN computation . With these methods , we can train our models in a small number of time steps , which could further improve the energy efficiency as well as the latency . Our contributions include : 1 . We propose the SPIDE method that is the first to train high-performance SNNs by spikes with common neuron models . Specifically , we propose ternary spiking neuron couples and prove that implicit differentiation for gradient calculation can be solved by spikes based on this design . Our method is applicable to both feedback and feedforward structures . 2 . We theoretically analyze the approximation error of solving implicit differentiation by spikes , and propose to modify the resting potential to remove the approximation bias and reduce the estimation variance , which enables training in a small number of time steps . 3 . Experiments show the low latency and firing sparsity during training , which demonstrates the great potential for energy-efficient training of SNNs on neuromorphic hardware . The performance on MNIST , CIFAR-10 , CIFAR-100 and CIFAR10-DVS are also competitive . 2 RELATED WORK . Early works seek biologically inspired methods to train SNNs , e.g . spike-time dependent plasticity ( STDP ) ( Diehl & Cook , 2015 ) or reward-modulated STDP ( Legenstein et al. , 2008 ) . Since the rise of successful ANNs , several works try to convert trained ANNs to SNNs to obtain high performance ( Hunsberger & Eliasmith , 2015 ; Rueckauer et al. , 2017 ; Sengupta et al. , 2019 ; Rathi et al. , 2019 ; Deng & Gu , 2021 ; Yan et al. , 2021 ) . However , they suffer from extremely large time steps and their structures are limited in the scope of ANNs . Others try to directly train SNNs by imitating backpropagation throught time ( BPTT ) and use surrogate derivative for discontinuous spiking functions ( Lee et al. , 2016 ; Wu et al. , 2018 ; Bellec et al. , 2018 ; Jin et al. , 2018 ; Shrestha & Orchard , 2018 ; Wu et al. , 2019 ; Zhang & Li , 2019 ; Neftci et al. , 2019 ; Zheng et al. , 2021 ) or compute gradient with respect to spiking times ( Bohte et al. , 2002 ; Zhang & Li , 2020 ; Kim et al. , 2020 ) . However , they suffer from approximation error and large memory costs . Xiao et al . ( 2021 ) propose the IDE method to train feedback spiking neural networks , which decouples the forward and backward procedures and avoids the common SNN training problems . However , all these methods require complex computation during training rather than spike-based . A few works focusing on training SNNs with spikes either are based on feedback alignment and limited in simple datasets ( Guerguiev et al. , 2017 ; Neftci et al. , 2017 ; Samadi et al. , 2017 ; O ’ Connor & Welling , 2016 ) , or require impractical special neuron models that require consideration of accumulated spikes for spike generation ( Thiele et al. , 2019b ; a ) , which is impractical on neuromorphic hardware . And they are only applicable to feedforward architectures . Instead , we are the first to leverage spikes with common neuron models to train SNNs with feedback or feedforward structures . The comparison of different methods is illustrated in Table 1 . 3 PRELIMINARIES . We first introduce preliminaries about spiking neurons and the IDE training method . The basic thought of IDE ( Xiao et al. , 2021 ) is to identify the underlying equilibrium states of FSNN computation so that gradients can be calculated based on implicit differentiation on the equilibrium state . We will briefly introduce the conclusion of equilibrium states in Section 3.2 and the IDE method in Section 3.3 . For more descriptions about the background please refer to Appendix A . 3.1 SPIKING NEURAL NETWORK MODELS . Spiking neurons draw inspirations from the human brain to communicate with each other by spikes . Each neuron integrates information from input spike trains by maintaining a membrane potential through a differential equation , and generates an output spike once the membrane potential exceeds a threshold , following which the membrane potential is reset to the resting potential . We consider the commonly used integrate and fire ( IF ) model and simple current model , whose discretized computational form is : ui [ t+ 0.5 ] = ui [ t ] + ∑ j wijsj [ t ] + b , si [ t+ 1 ] = H ( ui [ t+ 0.5 ] − Vth ) , ui [ t+ 1 ] = ui [ t+ 0.5 ] − ( Vth − urest ) si [ t+ 1 ] , ( 1 ) where ui [ t ] is the membrane potential of neuron i at time step t , si [ t ] is the binary output spike train of neuron i , wij is the connection weight from neuron j to neuron i , b is bias , H is the Heaviside step function , Vth is the firing threshold , and urest is the resting potential . We use subtraction as the reset operation . urest is usually taken as 0 in previous work , while we will reconsider it in Section 4.3 . 3.2 EQUILIBRIUM STATES OF FEEDBACK SPIKING NEURAL NETWORKS . Xiao et al . ( 2021 ) derive that the ( weighted ) average rate of spikes during FSNN computation with common neuron models would converge to an equilibrium state following a fixed-point equation given convergent inputs . We focus on the conclusions with the discrete IF model under both singlelayer and multi-layer feedback structures . The single-layer structure has one hidden layer of neurons with feedback connections on this layer . The update equation of membrane potentials is : u [ t+ 1 ] = u [ t ] +Ws [ t ] + Fx [ t ] + b− ( Vth − urest ) s [ t+ 1 ] , ( 2 ) where u [ t ] and s [ t ] are the vectors of membrane potentials and spikes of these neurons , x [ t ] is the input at time step t , W is the feedback weight matrix , and F is the weight matrix from inputs to these neurons . The average input and average firing rate are defined as x [ t ] = 1t+1 ∑t τ=0 x [ τ ] and α [ t ] = 1t ∑t τ=1 s [ τ ] , respectively . Define σ ( x ) = min ( 1 , max ( 0 , x ) ) . The equilibrium state of the single-layer FSNN is described as ( Xiao et al. , 2021 ) : If the average inputs converge to an equilibrium point x [ t ] → x∗ , and there exists γ < 1 such that ‖W‖2 ≤ γVth , then the average firing rates of FSNN with discrete IF model will converge to an equilibrium point α [ t ] → α∗ , which satisfies the fixed-point equation α∗ = σ ( 1 Vth ( Wα∗ + Fx∗ + b ) ) . Note that they take urest = 0 in this conclusion , if we consider nonzero urest , the constraint and the fixedpoint equation should be ‖W‖2 ≤ γ ( Vth − urest ) and α∗ = σ ( 1 Vth−urest ( Wα ∗ + Fx∗ + b ) ) . The multi-layer structure incorporates more non-linearity into the equilibrium fixed-point equation , which has multiple layers with feedback connections from the last layer to the first layer . The update equations of membrane potentials are expressed as : { u1 [ t+ 1 ] = u1 [ t ] +W1sN [ t ] + F1x [ t ] + b1 − ( Vth − urest ) s1 [ t+ 1 ] , ul [ t+ 1 ] = ul [ t ] + Flsl−1 [ t+ 1 ] + bl − ( Vth − urest ) sl [ t+ 1 ] , l = 2 , · · · , N. ( 3 ) The equilibrium state of the multi-layer FSNN with urest is described as ( Xiao et al. , 2021 ) : If the average inputs converge to an equilibrium point x [ t ] → x∗ , and there exists γ < 1 such that ‖W1‖2‖FN‖2 · · · ‖F2‖2 ≤ γ ( Vth − urest ) N , then the average firing rates of multi-layer FSNN with discrete IF model will converge to equilibrium points αl [ t ] → αl∗ , which satisfy the fixedpoint equations α1∗ = f1 ( fN ◦ · · · ◦ f2 ( α1 ∗ ) , x∗ ) and αl+1∗ = fl+1 ( αl ∗ ) , where f1 ( α , x ) = σ ( 1 Vth−urest ( W 1α+ F1x+ b1 ) ) and fl ( α ) = σ ( 1 Vth−urest ( F lα+ bl ) ) .

Summary: The paper aims at porting the IDE method  into a spike-fbased and more bio-plausible version. The previous IDE used firing rates rather than spikes for computation, although reference Xiao et al in NeurIPS 2021 had already addressed implementations in spiking neural networks. The authors analyze the approximation error that results from solving implicit differentiation by spikes and report a solution based on ternary spiking neurons, that can be implemented with pairs of standard spiking neurons. They achieve in this way quite good performance for MNIST and CIFAR10.