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 authors propose a relatively novel way of training spiking networks by considering the equilibrium state of their firing rates in response to input, and training using implicit differentiation, as has recently been explored in deep equilibrium models and implicit layers. They consider leaky and perfect integrate-and-fire neurons, as well as single layers with recurrent feedback, and hierarchical layers with feedback from the last layer to the first. The authors prove that the dynamics do indeed converge to an equilibrium state, and demonstrate the performance of the model on several state-of-the-art datasets.