Training Feedback Spiking Neural Networks by Implicit Differentiation on the Equilibrium State

1 Introduction . Spiking neural networks ( SNNs ) have gained increasing attention recently due to their inherent energyefficient computation [ 21 , 38 , 41 , 33 , 8 ] . Inspired by the neurons in the human brain , biologically plausible SNNs transmit spikes between neurons , enabling event-based computation which can be carried out on neuromorphic chips with less energy consumption [ 1 , 7 , 30 , 33 ] . Meanwhile , SNNs are computationally more powerful than artificial neural networks ( ANNs ) theoretically and are therefore regarded as the third generation of neural network models [ 24 ] . ∗Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Despite the advantages , directly supervised training of SNNs remains a hard problem , which hampers the real applications of SNNs compared with popular ANNs . The main obstacle lies in the complex spiking neuron model . While backpropagation [ 35 ] works well for ANNs , it suffers from the discontinuity of spike generation which is non-differentiable in SNN training . Most recent successful SNN training methods still imitate the backpropagation through time ( BPTT ) [ 39 ] framework by error propagation through the computational graph unfolded along time steps , and they deal with the spiking function by applying surrogate derivatives to approximate the gradients [ 41 , 5 , 14 , 38 , 42 , 27 , 16 , 47 ] , or by computing the gradients with respect to the spiking time only on the spiking neurons [ 6 , 46 , 16 ] . However , these methods either accumulate approximation error along time steps , or suffer from the “ dead neuron ” problem [ 38 ] , i.e . learning would not occur when no neuron spikes . At the same time , BPTT requires memorizing intermediate variables at all time steps and backpropagating along them , which is memory-costing and biologically implausible . So it is necessary to consider training methods other than backpropagation along computational graphs that fit SNNs better . On the other hand , most recent SNN models simply imitate the feedforward architectures of ANNs [ 41 , 38 , 42 , 46 , 47 ] , which ignores the ubiquitous feedback connections in the human brain . Feedback ( recurrent ) circuits are critical to human ’ s vision system for object recognition [ 15 ] . Meanwhile , [ 18 ] shows that shallow ANNs with recurrence achieve higher functional fidelity of human brains and similarly high performance on large-scale vision recognition tasks , compared with deep ANNs . So incorporating feedback connections enables neural networks to be shallower , more efficient , and more brain-like . As for SNNs , feedback was popular in early models like Liquid State Machine [ 25 ] , which leverages a recurrent reservoir layer with weights fixed or trained by unsupervised methods . And compared with the uneconomical cost for ANNs to incorporate feedback connections by unfolding along time , SNNs naturally compute with multiple time steps , which inherently supports feedback connections . Most recent SNN models imitate feedforward architectures because they were once lacking effective training methods and thus they borrow everything from successful ANNs . We focus on another direction , i.e . feedback SNN , which is a natural choice for visual tasks as well . In this work , we consider the training of feedback spiking neural networks ( FSNN ) , and propose a novel method based on the Implicit Differentiation on the Equilibrium state ( IDE ) . Inspired by recent advances in implicit models [ 3 , 4 ] , which treat weight-tied ANNs as solving a fixed-point equilibrium equation and propose alternative implicit models defined by the equation , we derive that when the average inputs converge to an equilibrium , the average firing rates of FSNNs would gradually evolve to an equilibrium state along time , which follows a fixed-point equation as well . Then we view the forward computation of FSNN as a black-box solver for the fixed-point equation , and borrow the idea of implicit differentiation from implicit models [ 3 , 4 ] to calculate the gradients , which only relies on the equation rather than the exact forward procedure . In this way , gradient calculation is agnostic to the spiking function in SNN , thus avoiding the common difficulties in SNN training . While implicit differentiation may seem too abstract to be computed in the brain , we briefly discuss the biological plausibility and show that it only requires computing another equilibrium along the inverse connections of neurons . Besides , we incorporate the multi-layer structure into the feedback model for better representation ability . Our contributions include : 1 . We are the first to theoretically derive the equilibrium states with a fixed-point equation for the average firing rates of FSNNs with the ( leaky ) integrate and fire model under both continuous and discrete views . According to this , the forward computation of FSNNs can be interpreted as solving a fixed-point equation . 2 . We propose a novel training method for FSNNs based on the implicit differentiation on the equilibrium state , which is decoupled from the forward computational graph and avoids SNN training problems , e.g . non-differentiability and large memory costs . We also discuss the biological plausibility and demonstrate the connection to the Hebbian learning rule . 3 . We conduct extensive experiments on MNIST , Fashion-MNIST , N-MNIST , CIFAR-10 , and CIFAR-100 , which demonstrate the superior results of our methods with fewer neurons and parameters in a small number of time steps for both static images and neuromorphic inputs . Especially , our directly trained model can outperform the state-of-the-art SNN performance on the complex CIFAR-100 dataset with only 30 time steps . 2 Related Work . Training Methods for Spiking Neural Networks . Early works apply biologically inspired method , spike-time dependent plasticity ( STDP ) [ 9 ] , to formulate a bottom-up unsupervised learning rule , or choose reward-modulated STDP [ 22 ] with limited supervision . Since the rise of successful ANNs , error backpropagation and gradient descent have inspired many methods . One direction is to convert a trained ANN to SNN [ 13 , 34 , 37 , 32 , 8 , 44 ] . However , they suffer from conversion errors and extremely large simulation time steps . The other methods are to directly calculate the gradient and train SNNs . These methods either compute the gradient with respect to spiking times [ 6 , 46 ] , or leverage a surrogate derivative for discontinuous spiking functions [ 21 , 41 , 5 , 14 , 38 , 42 , 27 , 47 ] , or combine them [ 16 ] . However , they suffer from the “ dead neuron ” problem [ 38 ] or accumulated approximation error , and typically require backpropagation along the computational graph to be unfolded by many time steps , which is memory-consuming and biologically implausible . As for SNN with feedback connection , [ 45 ] proposes the ST-RSBP method , which backpropagates errors at the spike-train level . They simply approximate the activation function of a neuron as a linear operation , and require long time steps for satisfactory results . In this work , we propose a new training method that does not rely on the exact reverse of the forward computation , which avoids problems of non-differentiability and large memory costs , and only requires short time steps for superior performance . There are also works trying methods other than BP along computational graphs to train SNNs , such as equilibrium propagation [ 29 ] . However , [ 29 ] defines a complex computation form rather than common SNN models , and can hardly achieve satisfactory results on the simple MNIST dataset . Instead , our work is based on SNN models applicable on neuromorphic hardware and demonstrates promising results on more complex datasets . Equilibrium of Neural Networks . The study on the equilibrium of neural networks originates from energy-based models , e.g . Hopfield Network [ 11 , 12 ] . They view the dynamics or iterative procedures of feedback ( recurrent ) neural networks as minimizing an energy function , which will converge to a minimum of the energy . Based on the energy , several training methods are proposed , including recurrent backpropagation [ 2 , 31 ] and more recent equilibrium propagation ( EP ) [ 36 ] . They heavily rely on the energy function and can hardly achieve competitive results as deep neural networks do . Deep equilibrium models [ 3 , 4 ] , on the other hand , are recently proposed models which treat weight-tied deep ANNs as solving a fixed-point equilibrium point , and propose alternative implicit models defined by the fixed-point equations rather than energy functions . They express the entire deep network as an equilibrium computation and solve and train it by root-finding methods and implicit differentiation , respectively , which achieves superior results . Most of these works are based on ANNs , except that [ 29 ] generalizes the EP method to neurons with binary communications . They define a complex neuron computation form and follow the methodology of energy-based EP method to approximate the gradients . Several recent works also study the equilibrium of SNNs [ 23 , 26 ] . They consider equilibrium from the perspective of solving a constrained optimization problem , but either do not propose to guide the training based on it or are limited in single-layer structure and simple problems . Differently , we are the first to derive the equilibrium state with a fixed-point equation for average firing rates of FSNNs with common SNN models , and propose to train SNNs by exact gradients through implicit differentiation , which is also scalable to multi-layer structure and deep learning problems . 3 Preliminaries . 3.1 Spiking Neural Network Models . Spiking neurons , different from real-valued ANNs , communicate with each other by binary spike trains along time . Each neuron maintains a membrane potential , which integrates input spike trains , and the neuron would generate an output spike once the membrane potential exceeds a threshold . The commonly used integrate and fire ( IF ) model and leaky integrate and fire ( LIF ) model describe the dynamics of the membrane potential as : IF : du dt = R · I ( t ) , u < Vth , LIF : τm du dt = − ( u− urest ) +R · I ( t ) , u < Vth , ( 1 ) where u is the membrane potential , I is the input current , Vth is the spiking threshold , and R and τm are resistance and time constant , respectively . Once u reaches Vth at time tf , a spike is generated and u is reset to the resting potential u = urest , which is usually taken as 0 . The spike train is expressed by the Dirac delta function : s ( t ) = ∑ tf δ ( t− tf ) . We consider the simple current model Ii ( t ) = ∑ j wijsj ( t ) + b , where wij is the weight from neuron j to neuron i , which is the same as ANN . After discretization , the general computation form for the IF and LIF model is described as : 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 ] − Vthsi [ t+ 1 ] , ( 2 ) where H ( x ) is the Heaviside step function , i.e . the non-differentiable spiking function , si [ t ] is the binary spike train of neuron i , and λ is 1 for the IF model while λ < 1 is a leaky term related to the constant τm and discretization time interval for the LIF model . The constant R , τm , and time step size are absorbed into the weights wij and bias b . We use subtraction as the reset operation .

Summary: The paper proposes a new training method for SNNs that relies on decoupling the forward propagation from the backward by viewing the SNNs as solving a fixed-point equilibrium equation and exploiting implicit differentiation results to compute the gradients. The theoretical contribution is to derive the equilibrium states for the fixed-point equations that the SNNs are solving, based on which the implicit differentiation technique can be leveraged. The practical contribution is a training algorithm that avoids the problems of non-differentiability in SNNs and achieves state-of-the-art performance on many datasets.