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

Abstract:
Spiking neural networks (SNNs) are brain-inspired models that enable energyefficient implementation on neuromorphic hardware. However, the supervised training of SNNs remains a hard problem due to the discontinuity of the spiking neuron model. Most existing methods imitate the backpropagation framework and feedforward architectures for artificial neural networks, and use surrogate derivatives or compute gradients with respect to the spiking time to deal with the problem. These approaches either accumulate approximation errors or only propagate information limitedly through existing spikes, and usually require information propagation along time steps with large memory costs and biological implausibility. In this work, we consider feedback spiking neural networks, which are more brainlike, and propose a novel training method that does not rely on the exact reverse of the forward computation. First, we show that the average firing rates of SNNs with feedback connections would gradually evolve to an equilibrium state along time, which follows a fixed-point equation. Then by viewing the forward computation of feedback SNNs as a black-box solver for this equation, and leveraging the implicit differentiation on the equation, we can compute the gradient for parameters without considering the exact forward procedure. In this way, the forward and backward procedures are decoupled and therefore the problem of non-differentiable spiking functions is avoided. We also briefly discuss the biological plausibility of implicit differentiation, which only requires computing another equilibrium. Extensive experiments on MNIST, Fashion-MNIST, N-MNIST, CIFAR-10, and CIFAR-100 demonstrate the superior performance of our method for feedback models with fewer neurons and parameters in a small number of time steps. Our code is available at https://github.com/pkuxmq/IDE-FSNN.

heading: 1 Introduction
text: 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.

heading: 2 Related Work
text: 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.

heading: 3 Preliminaries

heading: 3.1 Spiking Neural Network Models
text: 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.

heading: 3.2 Implicit Differentiation on the Fixed-Point Equation
text: We consider a fixed-point equation a = fθ(a) parameterized by θ. Let L(a∗) denote the objective function with respect to the equilibrium state a∗, and let gθ(a) = fθ(a)− a. The implicit differentiation on the equation satisfies ( I − ∂fθ(a
∗) ∂a∗
) da∗
dθ = ∂fθ(a ∗) ∂θ [3]. Therefore, the differentiation of
L(a∗) with respect to parameters can be calculated based on implicit differentiation:
∂L(a∗) ∂θ = −∂L(a ∗) ∂a∗ ( J−1gθ |a∗ ) ∂fθ(a∗) ∂θ , (3)
where J−1gθ |a∗ is the inverse Jacobian of gθ evaluated at a ∗. For the proof please refer to [3]. To solve the inverse Jacobian, we follow [3, 4] and solve an alternative linear system ( JTgθ |a∗ ) x +(
∂L(a∗) ∂a∗
)T = 0. We can leverage Broyden’s method [3, 4], which is a second-order quasi-Newton
approach; or we can alternatively use a fixed-point update scheme x = (JTfθ |a∗)x+ ( ∂L(a∗) ∂a∗ )T since JTgθ |a∗ = J T fθ |a∗ − I , and it converges with linear convergence rate as long as ‖JTfθ |a∗‖ < 1. In this way, gradients for the parameters can be calculated only with the equilibrium state and equation.

heading: 4 Proposed IDE Method
text: In this section, we first derive the equilibrium state of FSNNs under both continuous and discrete views, and demonstrate that FSNNs can be treated as solving a fixed-point equation. Then we introduce how to train the network by the proposed IDE method based on the equation and briefly discuss the biological plausibility. Finally, we incorporate the multi-layer structure into the model for more non-linearity and stronger representation ability.

heading: 4.1 Derivation of Equilibrium States for Feedback Spiking Neural Networks

heading: 4.1.1 Continuous View
text: We first consider a group of spiking neurons with feedback connections. Let u(t) and s(t) denote the membrane potentials and spikes of these neurons at time t respectively, x(t) denote the inputs, W denote the feedback weight matrix, F denote the weight matrix from inputs to these neurons, and b denote a constant bias. Under the IF model, the dynamics of membrane potentials are expressed as:
du dt = Ws(t−∆td) + Fx(t) + b− Vths(t), (4)
where ∆td is a time delay of feedback connections, and Vth is the threshold. Note that W and F represent linear operations including both fully-connected and convolutional layers. Define the average firing rates as a(t) = 1t ∫ t 0 s(τ)dτ , and the average inputs as x(t) = 1t ∫ t 0 x(τ)dτ . Then through integration, we have:
a(t) = 1
Vth
( t−∆td
t Wa(t−∆td) + Fx(t) + b−
u(t)
t
) . (5)
Eq.(5) roughly follows a fixed-point update scheme except the existence of u(t). Now we dive into u(t). Since neurons will not spike when the accumulated u(t) is negative, if vi(t) =( t−∆td t Wa(t−∆td) + Fx(t) + b ) i < 0, ideally neuron i generates no spike and its accumulated negative term remains in ui(t). So ui(t) can be divided as ui(t) = u−i (t) + u + i (t), where 1 tu − i (t) = min(vi(t), 0) is the remaining negative term, and u + i (t) is the positive one typically bounded in the range between 0 and Vth. There could be random error in u+i (t) in the context of random arrival of spikes rather than the average condition (e.g. the average is 0, but a large positive input followed by a negative one will generate an unexpected spike). Despite this, we can still assume u+i (t) to be bounded by a constant when average inputs converge. By this decomposition, we have the equation with the element-wise ReLU function (ReLU(x) = max(x, 0)) and bounded u+(t):
a(t) = ReLU ( 1
Vth
( t−∆td
t Wa(t−∆td) + Fx(t) + b )) − 1 Vth u+(t) t . (6)
With Eq.(6), we can derive that the average firing rate will gradually converge to an equilibrium state. Theorem 1. If the average inputs converge to an equilibrium point x(t) → x∗, and there exists constant c and γ < 1 such that |u+i (t)| ≤ c,∀i, t and ‖W‖2 ≤ γVth, then the average firing rates of FSNN with continuous IF model in Eq.(6) will converge to an equilibrium point a(t)→ a∗, which satisfies the fixed-point equation a∗ = ReLU ( 1 Vth (Wa∗ + Fx∗ + b) ) .
The proof can be found in Appendix C. Theorem 1 rigorously shows the equilibrium state under the IF model, and we can view the forward computation of FSNN as solving this fixed-point equation. As for the LIF model, we can similarly define the weighted average firing rate â(t) = ∫ t 0 κ(t−τ)s(τ)dτ∫ t 0 κ(t−τ)dτ
and the weighted average inputs x̂(t) = ∫ t 0 κ(t−τ)x(τ)dτ∫ t 0 κ(t−τ)dτ , where κ(τ) = exp(− τ τm ) is the response kernel of the LIF model. In this setting, however, there could be random errors caused by u+(t) as its denominator ∫ t 0 κ(t− τ)dτ does not go to infinity. We consider it as an approximate solver for the equilibrium with random errors, as shown in Proposition 1. Please refer to Appendix E for details. Proposition 1. If the weighted average inputs converge to an equilibrium point x̂(t) → x∗, and there exists constant c and γ < 1 such that |u+i (t)| ≤ c,∀i, t and ‖W‖2 ≤ γVth, then the weighted average firing rates â(t) of FSNN with continuous LIF model gradually approximate an equilibrium
point a∗ with bounded random errors, which satisfies a∗ = ReLU (
1 Vth
(Wa∗ + Fx∗ + b) ) .

heading: 4.1.2 Discrete View
text: In practice, we will simulate SNNs with discretization. Now consider the computation in Eq.(2). With feedback connections, the update equation of membrane potentials under the IF model is:
u[t+ 1] = u[t] + Ws[t] + Fx[t] + b− Vths[t+ 1], (7) where we treat the feedback delay in one time step for simplicity. Define the average firing rates as a[t] = 1t ∑t τ=1 s[τ ], the average inputs as x[t] = 1 t+1 ∑t τ=0 x[τ ], and u[0] = 0, s[0] = 0. By summation, we have:
a[t+ 1] = 1
Vth
( t
t+ 1 Wa[t] + Fx[t] + b− u[t+ 1] t+ 1
) . (8)
Different from the continuous view, a[t] is bounded in the range of [0, 1], since there could be at most t spikes during t time steps. Therefore, ui[t] will maintain both the negative terms and the exceeded positive ones. Similarly, ui[t] = u−i [t] +u + i [t], where 1 tu − i [t] = min (max (vi[t]− Vth, 0) , vi[t]) is the exceeded term, and u+i [t] is assumed to be bounded by a constant as previously indicated. Then:
a[t+1] = σ
( 1
Vth
( t
t+ 1 Wa[t] + Fx[t] + b )) − 1 Vth u+[t+ 1] t+ 1 ,whereσ(x) =  1, x > 1
x, 0 ≤ x ≤ 1 0, x < 0 .
(9) With Eq.(9), we can derive the equilibrium state under discrete view.
Theorem 2. If the average inputs converge to an equilibrium point x[t] → x∗, and there exists constant c and γ < 1 such that |u+i [t]| ≤ c,∀i, t and ‖W‖2 ≤ γVth, then the average firing rates of FSNN with discrete IF model in Eq.(9) will converge to an equilibrium point a[t] → a∗, which satisfies the fixed-point equation a∗ = σ ( 1 Vth (Wa∗ + Fx∗ + b) ) .
The proof can be found in Appendix C. And for the LIF model, define the weighted average firing rate â[t] = ∑t τ=1 λ
t−τ s[τ ]∑t τ=1 λ
t−τ and the weighted average inputs x̂[t] = ∑t τ=0 λ
t−τx[τ ]∑t τ=0 λ
t−τ , then we can similarly consider it as an approximation solver for the equilibrium state with random errors, as shown in Proposition 2. Please refer to Appendix E for details. Proposition 2. If the weighted average inputs converge to an equilibrium point x̂[t] → x∗, and there exists constant c and γ < 1 such that |u+i [t]| ≤ c,∀i, t and ‖W‖2 ≤ γVth, then the weighted average firing rates â[t] of FSNN with discrete LIF model gradually approximate an equilibrium
point a∗ with bounded random errors, which satisfies a∗ = σ (
1 Vth
(Wa∗ + Fx∗ + b) ) .

heading: 4.2 Training of Feedback Spiking Neural Networks
text: Based on the derivation in Section 4.1, we can view the forward computation of FSNNs as a black-box solver for the fixed-point equilibrium equations, with some errors caused by finite time steps or the LIF model. Then we assume the (weighted) average firing rates a[T ] after T time steps approximately follow the equations. We will demonstrate how to train FSNNs and its biological plausibility.

heading: 4.2.1 Loss and Gradient Computation
text: Suppose that we simulate the SNN by T time steps. Let a[T ] denote the final (weighted) average firing rates. We configure a readout layer after these spiking neurons, which performs as a fully-connected classification layer, with the number of outputs as class numbers. We assume that these neurons will not spike or reset, and do classification based on the accumulated membrane potential. Then the outputs are equivalent as a linear transformation on a[T ], i.e. o = Woa[T ]. The loss L is defined on o and labels y by commonly used loss functions L(o,y), and we leverage the cross-entropy loss.
Let fθ denote the function in fixed-point equation, e.g. fθ(a∗,x∗) = σ (
1 Vth
(Wa∗ + Fx∗ + b) )
. The gradient for parameters can be calculated based on the implicit differentiation as described in Section 3.2 by substituting a∗ with a[T ]. Then parameters can be optimized based on common gradient descent methods, e.g. SGD [35] and its variants. A pseudocode is presented in Appendix B.

heading: 4.2.2 Biological Plausibility of Implicit Differentiation
text: While implicit differentiation may seem too abstract for information propagation compared with backpropagation, we will briefly discuss the biological possibility for this calculation and its connection to the Hebbian learning rule [10]. Consider the equilibrium state a∗ following a∗ = fθ(a∗,x∗) = ReLU ( 1 Vth (Wa∗ + Fx∗ + b) ) . Let m = f ′θ(a ∗,x∗) = H ( 1 Vth (Wa∗ + Fx∗ + b) ) =
H(a∗),M = Diag(m),W̃ = MW, where H(x) = { 1, x > 0
0, x ≤ 0 . As indicated in Section 3.2,
the gradient can be calculated as ∇θL = ( ∂L ∂θ )T = ( ∂fθ(a ∗,x∗) ∂θ )T ( I − 1VthW̃ T )−1 ( ∂L ∂a∗ )T , and
we can leverage a fixed-point update scheme to solve for β∗ = ( I − 1VthW̃ T )−1 ( ∂L ∂a∗ )T by iterating βk+1 = 1VthW̃ Tβk + ∂L∂a∗ , β
k → β∗. This can be viewed as computing another equilibrium for these neurons: in this stage, neurons receive the inputs ∂L∂a∗ and they use the inverse directions of connections with a mask (M can be viewed as a mask matrix based on the firing condition in the first stage, which may be realized by some inhibition mechanisms) to compute for the equilibrium β∗. If W is symmetric and Vth = 1, it is similar to the energy-based method equilibrium propagation [36] to use the same weight as the forward computation for a second equilibrium
computation. Plugging β∗ into the gradient and calculating ( ∂fθ(a
∗,x∗) ∂θ
)T explicitly, we have:
∇WL = 1VthMβ ∗a∗T ,∇FL = 1VthMβ ∗x∗T . It is interesting to find that the change of weight
from neuron j to neuron i is proportional to the the equilibrium state of neuron j in the first stage and that of neuron i in the second stage, and is related to whether neuron i fires in the first stage, because ∇Wi,jL = 1Vthmiβ ∗ i a ∗ j ,∇Fi,jL = 1Vthmiβ ∗ i x ∗ j . It is to some extent similar to the locally updated Hebbian learning rule ∆wi,j ∝ xixj meaning that neurons wire together if they fire together [10], except that we take average firing rates and some temporal information (two stages) into account. Therefore, implicit differentiation calculation is only related to two equilibrium states by the neurons with a mask possibly realized by inhibition mechanisms, and may correspond to modified locally updated rules, which is more biologically plausible than BPTT with surrogate derivatives for SNNs. Please note that ‘biological plausibility’ here is a brief discussion in the context of the above properties, while there may be other aspects of biological implausibility as well.

heading: 4.3 Incorporating Multi-layer Structure into The Feedback Model
text: The multi-layer structure is commonly adopted in ANNs due to its stronger non-linearity and representation ability. To enhance the non-linearity of the fixed-point equilibrium equation, we propose to incorporate multi-layer structure into FSNNs as well. We configure N subgroups of neurons as different layers, where the inputs have connections to the first layer, the (l − 1)-th layer has connections to the l-th layer, and the last layer has feedback connections to the first layer. Let ul(t) and sl(t) denote the l-th layer, x(t) denote the inputs, W1 denote the feedback connection from the last layer to the first layer, and Fl denote the weight from the (l − 1)-th layer (or input) to the l-th layer. The discrete update equations of membrane potentials are expressed as:{
u1[t+ 1] = λu1[t] + W1sN [t] + F1x[t] + b1 − Vths1[t+ 1], ul+1[t+ 1] = λul+1[t] + Fl+1sl[t+ 1] + bl+1 − Vthsl+1[t+ 1], l = 1, 2, · · · , N − 1. (10)
An illustration figure is presented in Appendix A. With similar definitions of average firing rates al[t] for different layers, and uli[t] = u l i − [t] + uli + [t], we have the equilibrium state as the following. Theorem 3. If the average inputs converge to an equilibrium point x[t] → x∗, and there exists constant c and γ < 1 such that |uli + [t]| ≤ c,∀i, l, t and ‖W1‖2‖FN‖2 · · · ‖F2‖2 ≤ γV Nth , then the average firing rates of multi-layer FSNN with discrete IF model will converge to equilibrium points al[t]→ al∗, which satisfy the fixed-point equations a1∗ = f1 ( fN ◦ · · · ◦ f2(a1 ∗ ),x∗ ) and al+1 ∗ = fl+1(a l∗), where f1(a,x) = σ ( 1 Vth (W1a + F1x + b1) ) and fl(a) = σ ( 1 Vth (Fla + bl) ) .
The proof can be found in Appendix D. There is also a similar proposition for the LIF model, please refer to Appendix E for details. We will do classification based on the (weighted) average firing rate of the last layer and calculate the implicit differentiation for the equation on aN [T ]. The loss, solution for implicit differentiation, and optimization methods are the same as those in Section 4.2.1.

heading: 5 Experiments
text: In this section, we conduct extensive experiments to demonstrate the superior performance of our proposed method. Please refer to Appendix F for implementation details, including restrictions on the spectral norm and batch normalization, as well as training parameters. Since few previous methods leverage feedback architectures, we compare the results of our IDE method for FSNNs with most feedforward SNNs and report network structures2 as well as the number of neurons and parameters during computation (calculated according to the papers or released codes) for comparison.

heading: 5.1 MNIST and Fashion-MNIST
text: We first evaluate our method on simple static image datasets including MNIST [19] and FashionMNIST [43], and compare the results with other directly trained SNNs [21, 41, 38, 14, 45, 46] or similar ANNs. Inputs are the same images with binary or real values at all time steps, which can be
2The notations are: ‘64C5’ means a convolution with 64 output channels and kernel size 5, ‘s’ after ‘64C5’ means convolution with stride 2 while ‘u’ after that means a transposed convolution to upscale 2×, ‘P2’ means average pooling with size 2, ‘400’ means fully-connected to 400 neurons, and ‘F’ means feedback layers.
regarded as input currents [46]. We leverage single-layer FSNNs, and adopt convolutional layers for MNIST while using fully-connected layers for Fashion-MNIST following [45]. As shown in Table 1, our models achieve comparable or better results with fewer neurons and parameters in a relatively small number of time steps, compared with other direct SNN training methods on feedforward or feedback architectures. Especially, our model achieves superior results on Fashion-MNIST with the similar structure in only 5 time steps. The LIF model performs slightly better than the IF model, probably because it leverages temporal information by encoding weighted average firing rates.
Table 1: Performance on MNIST and Fashion-MNIST. Results are based on 5 runs of experiments.
MNIST
Method Network structure Time steps Mean±Std Best Neurons Params BP [21] 20C5-P2-50C5-P2-200 >200 / 99.31% 33K 518K
STBP [41] 15C5-P2-40C5-P2-300 30 / 99.42% 26K 607K SLAYER [38] 12C5-P2-64C5-P2 300 99.36%±0.05% 99.41% 28K 51K HM2BP [14] 15C5-P2-40C5-P2-300 400 99.42%±0.11% 99.49% 26K 607K ST-RSBP [45] 15C5-P2-40C5-P2-300 400 99.57%±0.04% 99.62% 26K 607K TSSL-BP [46] 15C5-P2-40C5-P2-300 5 99.50%±0.02% 99.53% 26K 607K IDE-IF (ours) 64C5s (F64C5) 30 99.49%±0.04% 99.55% 13K 229K
IDE-LIF (ours) 64C5s (F64C5) 30 99.53%±0.04% 99.59% 13K 229K Fashion-MNIST
Method Network structure Time steps Mean±Std Best Neurons Params ANN [45] 512-512 / / 89.01% 1.8K 670K
HM2BP [45] 400-400 400 / 88.99% 1.6K 478K TSSL-BP [46] 400-400 5 89.75%±0.03% 89.80% 1.6K 478K ST-RSBP [45] 400 (F400) 400 90.00%±0.14% 90.13% 1.2K 478K IDE-IF (ours) 400 (F400) 5 90.04%±0.09% 90.14% 1.2K 478K
IDE-LIF (ours) 400 (F400) 5 90.07%±0.10% 90.25% 1.2K 478K

heading: 5.2 N-MNIST
text: We also evaluate our method on the neuromorphic dataset N-MNIST [28], whose inputs are spikes collected by dynamic vision sensors. We follow the same data pre-possessing as [46] and take 30 time steps, and we can view the (weighted) average inputs gradually converge to that at the last time step. Table 2 demonstrates the comparison results of our models and other directly trained models [14, 38, 46, 42]. It shows that our method can achieve satisfactory performance on neuromorphic data as well. Especially, only 30 time steps are required by our method for satisfactory performance.

heading: 5.3 CIFAR-10 and CIFAR-100
text: Then we evaluate our method on more complex CIFAR-10 and CIFAR-100 datasets [17]. We leverage multi-layer FSNNs with structures modified from AlexNet and CIFARNet proposed in [42], as indicated in the footnote of Table 3. We compare our model with SNNs converted from ANNs [37, 8, 32, 44] and directly trained SNNs [42, 46, 20, 40]. For CIFAR-100, no result of
directly trained SNN is reported, so we only compare with the converted ones. Table 3 demonstrates the superior results of our directly trained models with fewer neurons and parameters in a small number of time steps. Especially, our model can outperform the state-of-the-art SNN performance on CIFAR-100 with only 30 time steps, and achieves 1.59% accuracy improvement when 100 time steps are adopted. Please refer to Appendix G for more comparison results between IF and LIF models.

heading: 5.4 Convergence to the Equilibrium
text: To verify the convergence of FSNNs to equilibrium states, we plot the difference norm on the fixedpoint equation at each time step, i.e. ‖fθ(a[t])− a[t]‖, where x = fθ(x) is the fixed-point equation and a[t] is the (weighted) average firing rate at time step t. Figure 1 demonstrates the convergence of different models. It is almost the same among different samples. Since the numerical precision of firing rates is only 1t , there would be a certain convergence error due to the finite time steps. And for the LIF model, there could be random errors compared with the IF model, as indicated in Section 4. The results conform to the theorems as the difference norm gradually decreases, i.e. firing rates converge to the equilibrium following the equation. It also indicates that exact precision is not necessary for satisfactory performance. Networks with fewer neurons converge faster, so a smaller number of time steps is needed, explaining why only 5 time steps are enough in the Fashion-MNIST experiment. For results on more datasets and different time steps, please refer to Appendix G.

heading: 5.5 Training Memory Costs
text: As described in the Introduction, an important advantage of our method is that we can avoid the large memory costs, from which the methods that backpropagate along the computational graph would suffer. To quantify this ease of training, we compare the GPU memory costs of our method and the
representative STBP method [41, 42]. The architecture and training settings are the same, and the results are shown in Table 4. It well illustrates the smaller memory costs of our method, which is also agnostic to time steps. Meanwhile, our method could achieve higher performance.

heading: 5.6 Firing Sparsity
text: As for the efficient neuromorphic computation, the firing rate is an important statistic since the energy consumption is proportional to the number of spikes. We calculate the average firing rate of trained models, and compare the IF and LIF model trained by our IDE method, as well as the LIF model trained by STBP method [41, 42] with the same structure. The results in Table 5 demonstrate the firing sparsity of our model, as the average firing rate is only around or less than 0.7%. And it shows that the LIF model has a slightly sparser response compared with the IF model. We note that TSSL-BP [46] also reported the statistics about firing rate. According to their results, their trained model on CIFAR-10 has a roughly total 9.86% firing rate within 5 time steps. So it is interesting to find that our models have fewer spikes than theirs, even if we have more time steps (30 vs. 5), not to mention that our models have fewer neurons. The results also show that the model trained by our method has sparser spikes compared with STBP, demonstrating the superiority of our method.
Table 5: The average firing rate of trained models. The model is trained on CIFAR-10 with AlexNet-F structure and 30 time steps.
Layer IDE-IF IDE-LIF STBP-LIF*
Layer 1 0.0172 0.0166 0.0190 Layer 2 0.0041 0.0039 0.0082 Layer 3 0.0025 0.0024 0.0113 Layer 4 0.0008 0.0008 0.0055 Layer 5 0.0200 0.0177 0.0108
Total 0.0070 0.0066 0.0102 * Our implementation

heading: 6 Conclusion
text: In this work, we propose a novel training method for feedback spiking neural networks based on implicit differentiation on the equilibrium state. We first derive the equilibrium states of (weighted) average firing rates for the IF and the LIF models of FSNNs under both continuous and discrete views. Then we propose to optimize parameters of FSNNs only based on the implicit differentiation on the underlying fixed-point equation. This enables the backward procedure to be decoupled from the forward computational graph and therefore avoids the common training problems for SNNs, such as non-differentiability and large memory costs. Meanwhile, we briefly discuss the biological plausibility for the calculation of implicit differentiation, which only requires computing another equilibrium and is related to the locally updated Hebbian learning rule. Extensive experiments demonstrate the superior results of our method and models with fewer neurons and parameters in a small number of time steps, and the spikes are sparser in our trained models as well.

heading: Acknowledgement
text: Z. Lin was supported by the NSF China (No.s 61625301 and 61731018), NSFC Tianyuan Fund for Mathematics (No. 12026606) and Project 2020BD006 supported by PKU-Baidu Fund. Yisen Wang is partially supported by the National Natural Science Foundation of China under Grant 62006153, and Project 2020BD006 supported by PKU-Baidu Fund.

Review summaries:
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.

This paper extends prior work [3,4] termed Implicit Differentiation on the Equilibrium state (IDE) to spiking neural networks. The authors derive a learning rule for the weights that is a product of two local equilibrium points in 2 phases. They perform experiments to compare their method with other SOTA SNNs, outperforming some.

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.