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

Spiking neurons draw inspiration 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 reset potential. We consider the commonly used integrate and fire (IF) and leaky integrate and fire (LIF) neuron models, whose general discretized computational form is:

where $u_{i}[t]$ is the membrane potential of neuron $i$ at time step $t$, $s_{i}[t]$ is the binary output spike train of neuron $i$, $w_{ij}$ is the connection weight from neuron $j$ to neuron $i$, $b_{i}$ is bias, $H$ is the Heaviside step function, $V_{th}$ is the firing threshold, $u_{reset}$ is the reset potential, and $\lambda=1$ for the IF model while $\lambda<1$ is a leaky term for the LIF model. We use subtraction as the reset operation. $u_{reset}$ is usually taken as $0$ in previous work, while we will reconsider it in Section 4.3. We mainly consider the IF neuron model by default and will demonstrate that our method is applicable to the LIF model as well.

Due to the complex spiking neuron model which is discontinuous, directly supervised training of SNNs is a hard problem, since the explicit computation is non-differentiable and therefore backpropagation along the forward computational graph is problematic. The IDE training method [Xiao et al., 2021] considers another approach to calculate gradients that does not rely on the exact reverse of the forward computation, which avoids the problem of non-differentiability as well as large memory costs by BPTT methods with surrogate gradients. Specifically, the IDE training method first derives that the (weighted) average firing rate of FSNN computation with common neuron models would gradually evolve to an equilibrium state along time, which follows a fixed-point equation. Then by viewing the forward computation of FSNN as a black-box solver for this equation, and applying implicit differentiation on the equation, gradients can be calculated only based on this equation and the (weighted) average firing rate during the forward computation rather than the exact forward procedure. Therefore, the forward and backward procedures are decoupled and the non-differentiability is avoided.

We briefly introduce the conclusion of equilibrium states in Section 3.2.1 and the IDE method in Section 3.2.2.

The common spiking neuron model only generates spikes when the membrane potential exceeds a positive threshold, which limits the firing rate from representing negative information. To enable approximation of possible negative values for implicit differentiation calculation in Section 4.2, we require negative spikes, whose expression could be:

and the reset is the same as usual: $u_{i}[t+1]-(V_{th}-u_{reset})s_{i}[t+1]$. Direct realization of such ternary output, however, may be not supported by common neuromorphic computation of SNNs.

We propose to leverage two coupled common neurons to realize this computation. As illustrated in Figure 2, the two coupled neurons with the common IF model (Eq. (1)) receive opposite inputs and output opposite spikes, which aim to deal with positive information and negative information with spikes, respectively. They should share a reset operation in order to accord with Eq. (6), which can be realized by the connections between them: as we use subtraction as the reset operation, the connection whose weight equals $V_{th}-u_{reset}$ enables one neuron to reset the other equivalently. To see how this works, consider the condition that the accumulated membrane potential of neuron 1 reaches $V_{th}$, then neuron 1 would generate a spike and reset, and the output is this positive spike. At the same time, the membrane potential of neuron 2 is $-V_{th}$ and the neuron will not fire and reset, but the spike from neuron 1 will reset it to $-u_{reset}$, which accords with our desired reset for ternary output. Similarly, if the inputs are negative, neuron 2 will generate a spike which will be treated negative as output, and both neurons are reset. For the operation of taking negative, one solution is to enable the reverse operation on hardware; another is to reconnect neuron 2 with other neurons while taking the weight negative to that of neuron 1. As for most existing neuromorphic hardware, the double connections to the two coupled neurons are more feasible, and this implies that twice the synaptic operations are needed for each spike. Therefore, such kind of coupled common neurons can realize ternary output.

We note that the SpikeGrad algorithm [Thiele et al., 2020] also requires neurons for ternary output. However, they do not consider how such kind of operation can be implemented with the common neuron models that follow the basic properties of biological neurons, and moreover, they propose another modified neuron model in practice that requires consideration of accumulated spikes for spike generation, which is hardly supported in practice. Differently, our method is with the common neuron model (IF or LIF). This can be suitable for neuromorphic hardware. Also, SpikingYOLO [Kim et al., 2020b] proposes ternary neurons for the inference of SNNs converted from ANNs without considering the implementation by the common neuron models. Our ternary spiking neuron couples may also provide a way to implement their method on hardware.

Based on the coupled neurons in Section 4.1, we can solve implicit differentiation with spikes. For notation simplicity, we directly use Eq. (6) as a ternary neuron without detailing coupled neurons below. Our main focus is on solving Eq. (5) with spikes. We first consider the IF neuron model by default and will include the LIF model later. The brief outline for the derivation is: we first derive the update equation of membrane potentials, then we derive the equivalent equation of the rate of spikes with eliminating perturbation, and finally, we prove that the firing rate converges to the solution of Eq. (5).

We first consider the single-layer condition. Let $\bm{\alpha}[T_{F}]$ denote the average firing rate of these neurons after the forward computation with time steps $T_{F}$ as an approximate equilibrium state (we treat the forward procedure as the first stage), $\mathbf{g}=\left(\frac{\partial\mathcal{L}}{\partial\bm{\alpha}[T_{F}]}\right)^{\top}$ denote the gradient of the loss function on this approximate equilibrium state, and $\mathbf{m}=\sigma^{\prime}(\bm{\alpha}[T_{F}]),\mathbf{M}=\text{Diag}(\mathbf{m})$ denote a mask indicator based on the firing condition in the first stage, where $\sigma^{\prime}(x)=\left\{\begin{aligned} &1,\,0<x<1\\ &0,\,\text{else}\\ \end{aligned}\right.$. We will have another $T_{B}$ time steps in the second backward stage to calculate implicit differentiation. We set the input to these neurons as $\mathbf{g}$ at all time steps, which can be viewed as input currents [Zhang & Li, 2020, Xiao et al., 2021]. Then along the inverse connections of neurons and with a mask on neurons or weights and an output rescale, the computation of FSNN with ternary neurons is:

where $V_{th},u_{reset}$ and $V_{th}^{b},u_{reset}^{b}$ are the threshold and reset potential during the first and second stage, respectively. Define the ‘average firing rate’ at this second stage as $\bm{\beta}[t]=\frac{1}{t}\sum_{\tau=1}^{t}\mathbf{s}[\tau]$, and $\mathbf{u}[0]=\mathbf{0},\mathbf{s}[0]=\mathbf{0}$, then through summation, we have (let $V_{u}=V_{th}-u_{reset},V_{u}^{b}=V_{th}^{b}-u_{reset}^{b}$):

Since there would be at most $t$ spikes during $t$ time steps, $\bm{\beta}$ should be bounded in the range of $[-1,1]$. The membrane potential $\mathbf{u}_{i}[t]$ will maintain the exceeded terms, i.e. define $\mathbf{v}_{i}[t]=\left(\frac{t}{t+1}\frac{1}{V_{u}}(\mathbf{M}\mathbf{W})^{\top}\bm{\beta}[t]+\mathbf{g}\right)_{i}$, we can divide $\mathbf{u}_{i}[t]$ as $\mathbf{u}_{i}^{E}[t]+\mathbf{u}_{i}^{B}[t]$, where $\mathbf{u}_{i}^{E}[t]=\max\left(\mathbf{v}_{i}[t]-V_{th}^{b},0\right)+\min\left(\mathbf{v}_{i}[t]+V_{th}^{b},0\right)$ is the exceeded term while $\mathbf{u}_{i}^{B}[t]$ is a bounded term [Xiao et al., 2021] which is typically bounded in the range of $[-V_{th}^{b},V_{th}^{b}]$. Then, Eq. (8) turns into:

where $\phi(x)=\min(1,\max(-1,x))$. Note that if $\mathbf{g}$ and $(\mathbf{M}\mathbf{W})^{\top}$ are in an appropriate range, there would be no exceeded term and $\phi$ will not take effect. Indeed we will rescale the loss to control the range of $\mathbf{g}$, as will be indicated in Section 4.3. With this consideration, we can prove that $\bm{\beta}[t]$ converges to the solution of Eq. (5).

The proof and discussion of assumptions are in A. With Theorem 1, we can solve Eq. (5) by simulating this second stage of SNN computation to obtain the ‘firing rate’ $\bm{\beta}[T_{B}]$ as the approximate solution. Plugging this solution to Eq. (4), the gradients can be calculated by: $\nabla_{\mathbf{W}}\mathcal{L}=\frac{1}{V_{th}-u_{reset}}\mathbf{M}\bm{\beta}[T_{B}]\bm{\alpha}[T_{F}]^{\top},\nabla_{\mathbf{F}}\mathcal{L}=\frac{1}{V_{th}-u_{reset}}\mathbf{M}\bm{\beta}[T_{B}]\mathbf{\overline{x}}[T_{F}]^{\top},$ and
$\nabla_{\mathbf{b}}\mathcal{L}=\frac{1}{V_{th}-u_{reset}}\mathbf{M}\bm{\beta}[T_{B}]$.

Note that in practice, even if the data distribution is not in a proper range, we can still view $\phi$ as a kind of clipping for improperly large numbers, which is similar to empirical techniques like “gradient clipping” to stabilize the training.

Then we consider the extension to the multi-layer condition. Let $\bm{\alpha}^{l}[T_{F}]$ denote the average firing rate of neurons in layer $l$ after the forward computation, $\mathbf{g}=\left(\frac{\partial\mathcal{L}}{\partial\bm{\alpha}^{N}[T_{F}]}\right)^{\top}$ denote the gradient of the loss function on the approximate equilibrium state of the last layer, and $\mathbf{m}^{l}=\sigma^{\prime}(\bm{\alpha}^{l}[T_{F}]),\mathbf{M}^{l}=\text{Diag}(\mathbf{m}^{l})$ denote the mask indicators. Similarly, we will have another $T_{B}$ time steps in the second stage and set the input to the last layer as $\mathbf{g}$ at all time steps. The computation of FSNN with ternary neurons is calculated as:

The ‘average firing rates’ $\bm{\beta}^{l}[t]$ are similarly defined for each layer, and the equivalent form can be similarly derived as:

The convergence of the ‘firing rate’ at the last layer to the solution of Eq. (5) can be similarly derived as Theorem 1. However, we need to calculate gradients for each parameter as Eq. (4), which is more complex than the single-layer condition. Actually, we can derive that the ‘firing rates’ at each layer converge to equilibrium points, based on which the gradients can be easily calculated with information from the adjacent layers. Theorem 2 gives a formal description.

The functions $f_{l}$ are defined in Section 3.2.1. For the proof please refer to B. With Theorem 2, by putting the solutions into Eq. (4), the gradients can be calculated by: $\nabla_{\mathbf{F}^{l}}\mathcal{L}=\frac{1}{V_{th}-u_{reset}}\mathbf{M}^{l}\bm{\beta}^{l}[T_{B}]\bm{\alpha}^{l-1}[T_{F}]^{\top}$ $(l=2,\cdots,N)$, $\nabla_{\mathbf{F}^{1}}\mathcal{L}=\frac{1}{V_{th}-u_{reset}}\mathbf{M}^{1}\bm{\beta}^{1}[T_{B}]\mathbf{\overline{x}}[T_{F}]^{\top}$, $\nabla_{\mathbf{W}^{1}}\mathcal{L}=\frac{1}{V_{th}-u_{reset}}\mathbf{M}^{l}\bm{\beta}^{l}[T_{B}]\bm{\alpha}^{N}[T_{F}]^{\top}$, and $\nabla_{\mathbf{b}^{l}}\mathcal{L}=\frac{1}{V_{th}-u_{reset}}\mathbf{M}^{l}\bm{\beta}^{l}[T_{B}]$ $(l=2,\cdots,N)$.

Note that the gradient calculation shares an interesting local property, i.e. it is proportional to the firing rates of the two neurons connected by it: $\nabla_{\mathbf{F}^{l}_{i,j}}\mathcal{L}=\frac{1}{V_{th}-u_{reset}}\mathbf{m}^{l}_{i}\bm{\beta}^{l}_{i}\bm{\alpha}^{l-1}_{j}$. During calculation, since we will have the firing rate of the first stage before the second stage, this calculation can also be carried out by event-based calculation triggered by the spikes in the second stage. So the weight update could be event-driven as well.

Also, note that the theorems still hold if we degrade our feedback models to feedforward ones by setting feedback connections as zero. In this setting, the dynamics and equilibriums degrade to direct functional mappings, and the implicit differentiation degrades to the explicit gradient. We can still approximate gradients with this computation.

In the following, we take $V_{th}^{b}-u_{reset}^{b}=1$ by default to fulfill the assumption of theorems (it may take other values if we correspondingly rescale the outputs and we set $1$ for simplicity). Other techniques like dropout can also be included. Please refer to C for details.

Section 4.2 derives that we can solve implicit differentiation with spikes, as the average firing rate will gradually converge to the solution. In practice, however, we will simulate SNNs for finite time steps, and a smaller number of time steps is better for lower energy consumption. This will introduce approximation error which may hamper training. In this subsection, we theoretically study the approximation error and propose to adjust the reset potential to reduce it. Inspired by the theoretical analysis on quantized gradients [Chen et al., 2020], we will analyze the error from the statistical perspective.

For the ‘average firing rates’ $\bm{\beta}^{l}[t]$ in Eq. (8) and the multilayer counterparts, the approximation error $e$ to the equilibrium states consists of three independent parts $e_{e}$, $e_{r}$ and $e_{i}$: the first is ${\mathbf{u}^{l}}^{E}[t+1]$ that is the exceeded term due to the limitation of spike number, the second is ${\mathbf{u}^{l}}^{B}[t+1]$ which can be viewed as a bounded random variable, and the third is the convergence error of the iterative update scheme without ${\mathbf{u}}^{l}[t+1]$, i.e. let $\mathbf{b}^{l}[t]$ denote the iterative sequences for solving ${\bm{\beta}^{l}}^{*}$ as $\mathbf{b}^{l}[t+1]=\frac{t}{t+1}\frac{1}{V_{th}-u_{reset}}(\mathbf{M}^{l+1}\mathbf{F}^{l+1})^{\top}\mathbf{b}^{l}[t]$, the convergence error is $\lVert\mathbf{b}^{l}[t]-{\bm{\beta}^{l}}^{*}\rVert$. The second part $e_{r}$ can be again decomposed into two independent components $e_{r}=e_{q}+e_{s}$: $e_{q}$ is the quantization effect due to the precision of firing rates ($\frac{1}{T}$ for $T$ time steps) if we first assume the same average inputs at all time steps, and $e_{s}$ is due to the random arrival of spikes rather than the average condition, as there might be unexpected output spikes, e.g. the average input is $0$ and the expected output should be $0$, but two large positive inputs followed by one larger negative input at the last time would generate two positive spikes while only one negative spike. So the error is divided into: $e=e_{e}+e_{q}+e_{s}+e_{i}$. Since the iterative formulation is certain for $e_{i}$, we focus on $e_{e}$, $e_{q}$ and $e_{s}$.

Firstly, the error $e_{q}$ due to the quantization effect is influenced by the input scale and time steps $T_{B}$. To enable proper input scale and smaller time steps, we rescale the loss function by a factor $s_{l}$, since the magnitude of gradients with cross-entropy loss is relatively small. We scale the loss to an appropriate range so that information can be propagated by SNNs in smaller time steps, and most signals are in the range of $\phi$ as analyzed in Section 4.2. The base learning rate is scaled by $\frac{1}{s_{l}}$ correspondingly. This is also adopted by Thiele et al. [2020].

Then given the scale and number of time steps, $e_{q}$, $e_{e}$ and $e_{s}$ can be treated as random variables from statistical perspective, and we view $\bm{\beta}^{l}[t]$ as stochastic estimators for the equilibrium states with $e_{i}$. For the stochastic optimization algorithms, the expectation and variance of the gradients are important for convergence and convergence rate [Bottou, 2010], i.e. we hope an unbiased estimation of gradients and smaller estimation variance. As for $e_{e}$ and $e_{s}$, they depends on the input data and the expectations are $\mathbb{E}[e_{e}]=0,\mathbb{E}[e_{s}]=0$ (the positive and negative parts have the same probability). While for $e_{q}$, it will depend on our hyperparameters $V_{th}^{b}$ and $u_{reset}^{b}$. Since the remaining terms in ${\mathbf{u}^{l}}^{B}[t+1]$ caused by the quantization effect is in the range of $[u_{reset}^{b},V_{th}^{b}]$ for positive terms while $[-V_{th}^{b},-u_{reset}^{b}]$ for negative ones, given $V_{th}^{b}-u_{reset}^{b}$ and considering the uniform distribution, only when $u_{reset}^{b}=-V_{th}^{b}$, $\mathbb{E}[e_{q}]=0$ for both positive and negative terms. Therefore, we should adjust the reset potential from commonly used 0 to $-V_{th}^{b}$ for unbiased estimation, as described in Proposition 3.

Also, taking $u_{reset}^{b}=-V_{th}^{b}$ achieves the smallest estimation variance for the quantization effect $e_{q}$, considering the uniform distribution on $[u_{reset}^{b},V_{th}^{b}]\cup[-V_{th}^{b},-u_{reset}^{b}]$. Since the effects of $e_{e}$, $e_{s}$ and $e_{i}$ are independent of $e_{q}$ and their variance is certain given inputs, it leads to Proposition 4.

With this analysis, we will take $V_{th}^{b}=0.5,u_{reset}^{b}=-0.5$ in the following. For $V_{th}$ and $u_{reset}$ during the first forward stage, we will also take $u_{reset}=-V_{th}$.

The original IDE method [Xiao et al., 2021] leverages other training techniques including modified batch normalization (BN) and restriction on weight spectral norm. Since the batch statistical information might be hard to obtain for calculation of biological systems and neuromorphic hardware, and training with the BN operation is hard to be implemented by spike-based calculation, we drop BN in our SPIDE method. The restriction on the weight norm, however, is necessary for the convergence of feedback models, as indicated in theorems. We will adjust it for a more friendly calculation, please refer to C for details.