SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Department of Physiology, Anatomy and Genetics
Oxford, United Kingdom
luke.taylor@hertford.ox.ac.uk
Andrew J King
andrew.king@dpag.ox.ac.uk
Nicol S Harper
nicol.harper@dpag.ox.ac.uk
The adaptive leaky integrate-and-fire (ALIF) model is fundamental within computational neuroscience and has been instrumental in studying our brains in silico. Due to the sequential nature of simulating these neural models, a commonly faced issue is the speed-accuracy trade-off: either accurately simulate a neuron using a small discretisation time-step (DT), which is slow, or more quickly simulate a neuron using a larger DT and incur a loss in simulation accuracy. Here we provide a solution to this dilemma, by algorithmically reinterpreting the ALIF model, reducing the sequential simulation complexity and permitting a more efficient parallelisation on GPUs. We computationally validate our implementation to obtain over a  $50 \times$  training speedup using small DTs on synthetic benchmarks. We also obtained a comparable performance to the standard ALIF implementation on different supervised classification tasks - yet in a fraction of the training time. Lastly, we showcase how our model makes it possible to quickly and accurately fit real electrophysiological recordings of cortical neurons, where very fine sub-millisecond DTs are crucial for capturing exact spike timing.
The surge of progress in artificial neural networks (ANNs) over the last decade has advanced our understanding of the potential computational principles underlying the processing of sensory information [1-9]. Although these networks architecturally bear a resemblance to the brain [10], they tend to omit a key physiological constraint: the spike. With their increased biological realism, spiking neural networks (SNNs) have shown great promise in bridging the gap between experimental data and computational models. SNNs can be fitted to real neural data [11-15], or used to simulate the
brain, offering a new level of understanding of the complex workings of the nervous system [16-19]. They also have engineering applications in energy-efficient machine learning [20].
An established class of spiking models in computational neuroscience is the leaky integrate-and-fire (LIF) neuron, with origins dating back to 1907 [21]. Just like in real neurons, input current (resulting from presynaptic input) charges the membrane potential of the model neurons, which then output binary signals (i.e. spikes) as a form of communication (Figure 1a). The adaptive leaky integrate-and-fire (ALIF) model [22] is a modern extension of the LIF. It more closely mimics the biology, capturing a key property of neurons, which is their adaptive firing threshold (i.e. spikes become less frequent in response to a steady input current [23]). ALIF neurons have been shown to accurately fit real neural recordings [11, 22, 24, 25] and to outperform the simpler LIF neurons on various machine learning benchmarks [26-29].
Despite these modelling advances, a major shortcoming of LIF and ALIF neurons is their slow inference and training times. Unlike real neurons, modelling neuron dynamics involves sequential computation over discretised time. This leads to a problematic trade-off between speed and accuracy when simulating SNNs. A small DT enables accurate modelling of dynamics, but is slow to stimulate and train on computer systems such as GPUs. A large DT obtains less accurate dynamics, but at the benefit of being faster to simulate and train [30] (Figure 1b). This raises the important question of capturing the best of both worlds: is there a way to accelerate the inference and training of spiking LIF and ALIF neurons without sacrificing simulation accuracy?
In this work, we address the speed-accuracy trade-off when simulating and training LIF and ALIF SNNs, and present a solution that is both fast and accurate. We take advantage of a fundamental property of neurons that is sometimes not modelled - the absolute refractory period (ARP). This is a short period of time following spike initiation, during which a neuron cannot fire again. As a result, a neuron can spike at most once within such a period (Figure 1c). We leverage this observation to develop a novel algorithmic reformulation of the ALIF model which reduces the sequential simulation complexity of the standard ALIF algorithm. Specifically, we outline how ALIF recurrent networks can be simulated with a constant sequential complexity  $O(1)$  over the ARP simulation length  $T_{R}$ , and how this approach can be extended to longer simulation lengths to obtain identical dynamics to the standard ALIF network algorithm. Faster simulation and training are theoretically obtained for growing length  $T_{R}$ , by either employing physiologically plausible ARPs of  $\sim 2\mathrm{ms}$  and decreasing the DT to a very fine value  $\sim 0.1\mathrm{ms}$  (for realistic neural modelling), or setting the DT to a coarser value and adopting larger non-physiological ARPs (for machine learning tasks). Our main contributions are the following:
- We develop a novel algorithmic reformulation of the ALIF model with an  $O(T / T_R)$  rather than  $O(T)$  - sequential complexity, for simulation length  $T$  and ARP length  $T_R$ .<sup>1</sup>
- We find that our model achieves substantial inference (up to  $40 \times$ ) and training speedup (up to  $53 \times$ ) over the standard ALIF SNN for increasing ARP time steps, and find this to hold over different numbers of simulation lengths, batch sizes, number of neurons and layers.
- We demonstrate the feasibility of our model to be trained using surrogate gradient descent, with our accelerated ALIF SNN achieving comparable accuracies to the standard ALIF SNN on temporal spiking classification datasets - yet in a fraction of the training time.
- Finally, we showcase how our ALIF implementation makes it possible to quickly fit real electrophysiological recordings of cortical neurons, where very fine sub-millisecond discretisation steps are important for accurately capturing exact spike timing.
To avoid ambiguity, simulation length  $T$  and ARP length  $T_{R}$  are number of time steps (dimensionless) and not unit time.
Figure 1: Problem overview. a. Schematic of an ALIF neuron: input current  $I$  charges membrane potential  $V$  and outputs spikes  $S$  if firing threshold is reached (with the neuron's internal state evolving over time). b. An example of the simulation trade-off problem when simulating a single ALIF neuron with fixed synaptic weights receiving Poisson spike input. The simulation error and the speed grow for increasing discretisation time (DT). c. Observation for our solution: a neuron emits at most a single spike during a simulation span  $T_{R}$  equal in length to the neuron's absolute refractory period (ARP).
2 Background and related work
Standard ALIF model We introduce the recurrent SNN of ALIF neurons with fixed ARP and latency of recurrent transmission, defined by the following set of equations [26, 31].
$$
S _ {i} ^ {(l)} [ t ] = f \left(V _ {i} ^ {(l)} [ t ]\right) = \mathbb {1} _ {V _ {i} ^ {(l)} [ t ] > \theta_ {i} ^ {(l)} [ t ]} \quad (\text {O u t p u t s p i k e}) \tag {1}
$$
$$
V _ {i} ^ {(l)} [ t ] = \left(\beta_ {i} ^ {(l)} V _ {i} ^ {(l)} [ t - 1 ] + \left(1 - \beta_ {i} ^ {(l)}\right) I _ {i} ^ {(l)} [ t ]\right) \left(1 - S _ {i} ^ {(l)} [ t - 1 ]\right) \quad (\text {M e m b r a n e p o t e n t i a l}) \tag {2}
$$
$$
\tilde {I} _ {i} ^ {(l)} [ t ] = \left(b _ {i} ^ {(l)} + \underbrace {\sum_ {j = 1} ^ {N ^ {(l - 1)}} W _ {i j} ^ {(l)} S _ {j} ^ {(l - 1)} [ t ]} _ {\text {F e e d f o r w a r d c u r r e n t}} + \underbrace {\sum_ {j = 1} ^ {N ^ {(l)}} W _ {i j} ^ {\text {r e c}} ^ {(l)} S _ {j} ^ {(l)} [ t - D ]} _ {\text {R e c u r r e n t c u r r e n t}}\right) \quad (\text {I n p u t c u r r e n t}) \tag {3}
$$
$$
I _ {i} ^ {(l)} [ t ] = \left\{ \begin{array}{l l} \tilde {I} _ {i} ^ {(l)} [ t ] & \text {i f} C _ {i} ^ {(l)} \geq T _ {R} \\ 0 & \text {o t h e r w i s e} \end{array} \right. \quad (\text {A b s o l u t e r e f r a c t o r y p e r i o d}) \tag {4}
$$
$$
\theta_ {i} ^ {(l)} [ t ] = 1 + d _ {i} ^ {(l)} a _ {i} ^ {(l)} [ t ] \quad (\text {A d a p t a t i o n}) \tag {5}
$$
$$
a _ {i} ^ {(l)} [ t ] = p _ {i} ^ {(l)} a _ {i} ^ {(l)} [ t - 1 ] + S _ {i} ^ {(l)} [ t - 1 ]
$$
At time  $t$ , neuron  $i$  within layer  $l$  (consisting of  $N^{(l)}$  neurons) receives input current  $I_{i}^{(l)}[t]$  and outputs a binary value  $S_{i}^{(l)}[t] \in \{1,0\}$  (i.e. spike or no spike) if a neuron's membrane potential  $V_{i}^{(l)}[t]$  reaches firing threshold  $\theta_{i}^{(l)}$  (Equation 1). The evolution of the membrane potential is described by the normalised discretised LIF equation (Equation 2), in which the membrane potential dissipates by a learnable factor  $0 \leq \beta_{i}^{(l)} \leq 1$  at every time step and resets to zero if a spike occurred at the previous time step. The input current is comprised of a constant bias source  $b_{i}^{(l)}$  and from incoming spikes reflecting feedforward  $W^{(l)} \in \mathbb{R}^{N^{(l)} \times N^{(l-1)}}$  and recurrent connectivity  $W^{\mathrm{rec}}(l) \in \mathbb{R}^{N^{(l)} \times N^{(l)}}$  (Equation 3).
Here, we assume the recurrent transmission latencies  $D$  to be of fixed length and equal in length to the ARP, that is  $T_{R} = D$ . With a simple modification however, our methods can work for longer  $D$ , or different  $D$  on each connection, so long as  $D \geq T_{R}$ . The ARP is enforced by only allowing input current to enter the neuron if the number of time steps  $C_i^{(l)}$  following the last spike equals or exceeds the ARP length (Equation 4). Lastly, adaptation is implemented by raising the firing threshold  $\theta_i^{(l)}[t]$
Here notation  $\mathbb{1}_{\mathrm{condition}}$  denotes the indicator function, which is equal to one if the condition is true and zero otherwise.
following each spike  $S_{i}^{(l)}[t - 1]$  (Equation 5), which decays exponentially to baseline  $\theta_{i}^{(l)} = 1$  in the absence of any spikes (using learnt decay factor  $0 \leq p_{i}^{(l)} \leq 1$  and adaptation scalar  $d_{i}^{(l)}$ ).
The ARP in biological neurons is typically about  $\sim 1 - 2\mathrm{ms}$  [32-34]. Our method takes advantage of the fact that the monosynaptic connection latency between neurons in local circuits is typically also often around  $\sim 1 - 2\mathrm{ms}$  [35]. LIF and ALIF neuronal models are typically run with a DT of about  $\sim 0.1\mathrm{ms}$  in the computational neuroscience literature and  $\sim 1\mathrm{ms}$  in machine learning literature [30]. Furthermore, the firing rate of neurons is often less than the reciprocal of the ARP, suggesting that a higher  $T_{R}$  than the ARP may still provide a reasonable approximation of neural behaviour for some purposes.
SNN training The main problem with training SNNs is the non-differentiable nature of their activation function in Equation 1 (whose derivative is undefined at  $V_{i}^{(l)}[t] = \theta_{i}^{(l)}[t]$  and zero otherwise). This precludes the direct use of the backprop algorithm [36], which has underpinned the successful training of ANNs. A popular solution is to replace the undefined spike derivative with a well-behaved function, referred to as a surrogate gradient [37-40], and training the network with backprop-through-time [41]. This method also supports the training of neural parameters other than synaptic connectivity, such as membrane time constants [42, 28] and adaptive firing thresholds [26-28], both of which we include in our work, as they have been shown to improve performance (and increase biological realism). It is worth mentioning that other SNN training methods exist such as mapping trained ANNs to SNNs by transforming ANN neuron activations to SNN neuron firing rates [43-46]. These methods are, however, of less interest to computational neuroscientists as they discard all temporal spike precision.
Related work Recent work has provided new theoretical insights for increasing the simulation accuracy when employing a larger  $\mathrm{DT} \geq 1\mathrm{ms}$  [30]. We are not aware of any work that accelerates the simulation and training times on GPUs when employing a smaller  $\mathrm{DT} \leq 1\mathrm{ms}$  (although see the NEST simulation library for simulating SNNs on CPUs [47-49]). There are, however, different methods for speeding up the backward pass (i.e. how gradients are computed within the SNN), which consequentially speeds up training times. Rather than being viewed as competing approaches, these methods could further augment the speed of our solution. In sparse gradient descent, gradients are only passed through neurons whose membrane potential is close to firing threshold, which can significantly accelerate the backward pass when neurons are mostly silent [50]. Inspired by work on training non-spiking networks [51, 52], other methods completely bypass the backward pass and adjust weights online [53, 27, 29]. Another approach is to propagate gradient information through spikes [54-60], which - similar to the idea of sparse gradient descent - is fast when neurons are mostly silent. This method, however, enforces neurons to spike at most once and can suffer from training instabilities (although see [61, 62]).
3 Theoretical results
We outline a novel reformulation of the ALIF model, which theoretically reduces the sequential simulation complexity from  $O(T)$  to  $O(T / T_R)$  (for simulation length  $T$  and ARP length  $T_R$ ). Using the observation that a neuron can spike at most once over simulation length  $T_R$  (equal to the ARP; Figure 1c), we propose simulating network dynamics sequentially in blocks of length  $T_R$  - as opposed to simulating every time step individually (Figure 2a) - as we show that these blocks can be simulated with a constant sequential complexity  $O(1)$ . Our reformulated ALIF model is mathematically the same as the standard ALIF model, but substantially faster to simulate and train. All the proofs for the propositions can be found in the Supplementary material.
3.1 Block: simulating single-spike ALIF dynamics with a constant sequential complexity
A SNN exhibits a sequential dependence due to the spike reset mechanism, where a neuron's membrane potential  $V_{i}[t]$  can be reset based on its previous output  $S_{i}[t - 1]$ . Consequently, the simulation of a SNN necessitates a sequential approach. This restriction can, however, be alleviated if we assume a neuron to spike at most once over a particular simulation length (such as the ARP). The following steps - grouped together into a module referred to as a Block (Figure 2b) - compute
b
Figure 2: Solution overview. a. Our proposed solution: instead of simulating network dynamics one time step after another (top), we sequentially simulate blocks of time equal in length to the neuron ARP (bottom), in which a neuron can spike at most once. b. A schematic of a Block: our proposed solution for emulating ALIF dynamics with a constant sequential complexity  $O(1)$  over a short duration in which a neuron spikes at most once.
ALIF dynamics (assuming at most one spike) without any sequential operations.
$$
\tilde {V} _ {i} [ t ] = \left(I _ {i} * \tilde {\beta} _ {i}\right) [ t ] \quad (\text {N o - r e s e t m e m b r a n e p o t e n t i a l}) \tag {6}
$$
$$
\tilde {S} _ {i} [ t ] = f \left(\tilde {V} _ {i} [ t ]\right) \quad (\text {F a u l t y o u p t s p i k e s}) \tag {7}
$$
$$
z _ {i} [ t ] = \phi (\tilde {S} _ {i}) [ t ] \quad \text {(L a t e n t t i m i n g o f s p i k e s)} \tag {8}
$$
$$
S _ {i} [ t ] = \mathbb {1} _ {z _ {i} [ t ] = 1} \quad (\text {C o r r e c t o u t p u t s p i k e s}) \tag {9}
$$
1. Calculate membrane potentials without reset The first step in the Block (Equation 6) converts input current  $I_{i}[t]$  to membrane potentials  $\tilde{V}_{i}[t]$  without spike reset (i.e. excluding the reset mechanism in Equation 2). This transformation is achieved using a convolution (Proposition 1), thus avoiding any sequential operations.
Proposition 1. Membrane potentials without spike reset are computed as a convolution  $\tilde{V}_i[t] = (I_i * \tilde{\beta}_i)[t]$  between input current  $I_i[t]$  and kernel  $\tilde{\beta}_i[t] = (1 - \beta_i)\beta_i^t$  with the initial membrane potential encoded as  $I_i[0] = \frac{V_i[0]}{1 - \beta_i}$ .
2. Faulty output spikes No-reset membrane potentials  $\tilde{V}_i[t]$  are mapped to erroneous output spikes  $\tilde{S}_i[t] = f(\tilde{V}_i[t])$  (Equation 7) using spike function  $f$  (Equation 1). This output can contain more than one spike, but only the first spike complies with the standard model dynamics, due to the omission of the reset mechanism and ARP constraint. Thus, to ensure that the Block only emits a single spike, all spikes succeeding the first spike occurrence are removed using the following steps.
3. Latent timing of spikes Erroneous spike output  $\tilde{S}_i[t]$  is mapped to a latent representation  $z_i[t] = \phi(\tilde{S}_i)[t]$  (Equation 8), encoding the timing of spikes. Function  $\phi(\cdot)$  (taking vector  $\tilde{S}_i$  as input; Proposition 2) is constructed to map all erroneous spikes  $\tilde{S}_i[t]$ , besides the first spike occurrence, to a value other than one (i.e.  $z_i[t] \neq 1$  for all  $t$  except for the smallest  $t$  satisfying  $\tilde{S}_i[t] = 1$  if such  $t$  exists).
Proposition 2. Function  $\phi (\tilde{S}_i)[t] = \sum_{k = 1}^t\tilde{S}_i[k](t - k + 1)$  acting on  $\tilde{S}_i\in \{0,1\} ^T$  contains at most one element equal to one  $\phi (\tilde{S}_i)[t] = 1$  for the smallest  $t$  satisfying  $\tilde{S}_i[t] = 1$  (if such  $t$  exists).
4. Correct output spikes Lastly, the correct spike output  $S_{i}[t] = \mathbb{1}_{z_{i}[t]} = 1$  is obtained by setting every value in  $z_{i}[t]$ , besides the value one (i.e. the first spike), to zero (Equation 9).
3.2 Blocks: simulating ALIF SNN dynamics with a  $O(T / T_R)$  sequential complexity
The standard ALIF SNN model can be reformulated using a chaining of Blocks, which reduces the sequential simulation complexity (as each Block is simulated in  $O(1)$ ). For a given ARP of length  $T_{R}$ , we observed that a neuron spikes at most once over simulation length  $T_{R}$  (Figure 1c). Thus, a simulation length  $T$  can be simulated using  $N = \frac{T}{T_{R}}$  Blocks, each of length  $T_{R}$ . Next, we outline how to simulate ALIF dynamics across Blocks to emulate the dynamics of the standard ALIF SNN. We introduce new notation to index Block  $1 \leq n \leq N$  using subscript  $n$  (e.g. input current to neuron  $i$  simulated in Block  $n$  is expressed as  $I_{i,n}[t]$ ) with time steps indexed between  $[1, T_{R}]$  (as opposed to  $[1, T]$ ) within a Block.
Input current and the ARP of a Block The input current in the standard model (Equation 3 and 4) is modified for the Block model (Proposition 3). The feedforward and recurrent current to Block  $n + 1$  are derived from the presynaptic and postsynaptic spikes from Block  $n + 1$  and Block  $n$  respectively. In addition, the ARP is enforced by applying a mask derived from the latent timing of spikes  $z_{i,n}[t]$  (Equation 8).
Proposition 3. The input current  $I_{i,n+1}[t]$  of neuron  $i$  simulated in Block  $n + 1$  (of length  $T_R$ ) is defined as follows, and enforces an absolute refractory period of length  $T_R$  and a monosynaptic transmission latency of  $D = T_R$ .
$$
I _ {i, n + 1} [ t ] = \left(b _ {i} + \underbrace {\sum_ {j = 1} ^ {N ^ {i n}} W _ {i j} S _ {j , n + 1} [ t ]} _ {\text {F e e d f o r w a r d c u r r e n t}} + \underbrace {\sum_ {j = 1} ^ {N ^ {o u t}} W _ {i j} ^ {r e c} S _ {j , n} [ t ]} _ {\text {R e c u r r e n t c u r r e n t}}\right) \underbrace {\mathbb {1} _ {z _ {i , n} [ t ] \geq \max _ {t} S _ {i , n} [ t ]}} _ {\text {A R P m a s k}}
$$
**Evolving membrane potentials between Blocks** Two cases are distinguished to correctly emulate the evolution of the membrane potentials between Blocks (Proposition 4): 1) if neuron  $i$  did not spike in Block  $n$  (i.e.  $\max_t S_{i,n}[t] = 0$ ), then its initial membrane potential  $V_{i,n+1}[0]$  in Block  $n+1$  is set to its final membrane potential  $V_{i,n}[T_R]$  in Block  $n$ . Otherwise 2), the initial membrane potential is set to zero to emulate a spike reset (and no state needs to be transferred between Blocks as the neuron is in a refractory state).
Proposition 4. The initial membrane potential  $V_{i,n+1}[0]$  of neuron  $i$  simulated in Block  $n+1$  (of length  $T_R$ ) is equal to the last membrane potential in Block  $n$  if no spike occurred and zero otherwise.
$$
V _ {i, n + 1} [ 0 ] = \left\{ \begin{array}{l l} V _ {i, n} [ T _ {R} ] & i f \max  _ {t} S _ {i, n} [ t ] = 0 \\ 0 & o t h e r w i s e \end{array} \right.
$$
**Evolving adaptive firing thresholds between Blocks** The adaptive firing threshold  $\theta_{i,n+1}[t]$  of neuron  $i$  in Block  $n+1$  is derived from the initial adaptive parameter  $a_{i,n+1}[0]$  (Proposition 5). Two cases are distinguished for deriving this parameter: 1) if the neuron did not spike during the previous Block  $n$ , this parameter is set to its last value,  $a_{i,n}[T_R]$ , in the prior Block; otherwise 2), the effect of the spike needs to be taken into account, for which the initial adaptive parameter is expressed as  $p_i^m (a_s + p_i^{-1})$ . Here,  $m$  is the number of time steps remaining in Block  $n$  following the spike and  $a_s$  is the adaptive parameter value at the time of spiking.
Proposition 5. The adaptive firing threshold  $\theta_{i,n+1}[t]$  of neuron  $i$  simulated in Block  $n+1$  (of length  $T_R$ ) is constructed from the initial adaptive parameter  $a_{i,n+1}[0]$ , which is equal to its last value in the previous Block if no spike occurred, and otherwise equal to an expression which accounts for the effect of the spike on the adaptive threshold.
$$
\theta_ {i, n + 1} [ t ] = 1 + d _ {i} p _ {i} ^ {t} a _ {i, n + 1} [ 0 ]
$$
$$
a _ {i, n + 1} [ 0 ] = \left\{ \begin{array}{l l} a _ {i, n} [ T _ {R} ] & \text {i f \max  _ {t} S _ {i , n} [ t ] = 0} \\ p _ {i} ^ {m} (a _ {s} + p _ {i} ^ {- 1}) & \text {o t h e r w i s e} \end{array} \right.
$$
$$
m = \sum_ {k} ^ {T _ {R}} \mathbb {1} _ {z _ {i, n} [ k ] > 1}, \quad a _ {s} = \sum_ {k} ^ {T _ {R}} a _ {i, n} [ k ] S _ {i, n} [ k ]
$$
<sup>3</sup>With appropriate zero padding to the simulation length if it is not divisible by  $T_{R}$ .
3.3 Theoretical speedup for simulating ALIF SNNs using Blocks
Method	Computational Complexity	Sequential Operations
Standard	O(NinNoutT)	O(T)
Blocks	O(NinNoutTR2N)	O(T/TR)
Table 1: Computational and sequential complexity of simulating a layer of  $N^{\mathrm{out}}$  neurons with  $N^{\mathrm{in}}$  input neurons for  $T$  time steps using the standard method and our method (with  $N$  Blocks each of length  $T_R$ ).
Simulating an ALIF SNN using our Blocks, rather than the conventional approach, requires fewer sequential operations (Table 1). Although the computational complexity of our Blocks approach is larger than the standard approach, the number of sequential operations is less. If we assume that the sequential steps in both methods are executed in an equal amount of time (as all the non-sequential operations can be run in parallel on a GPU), we obtain a theoretical speedup equal to the length of the ARP  $\frac{T}{N} = T_R$ .
We evaluated the training speedup of our accelerated ALIF SNN relative to the standard ALIF SNN and explored the capacity of our model to be fitted using different spiking classification datasets and real electrophysiological recordings of cortical neurons. Implementation details can be found in the Supplementary material and the code at https://github.com/webstorms/Blocks.
4.1 Training speedup scales with an increasing ARP
a
Figure 3: Training speedup of our model. a. Training speedup of our accelerated ALIF model compared to the standard ALIF model for different simulation lengths  $T$ , ARP time steps and batch sizes. b. Training speedup over different number of layers and hidden units (with an ARP of  $T_R = 40$  time steps,  $T = 1024$  time steps and batch size = 64; bars plot the mean and standard error over ten runs). Assuming  $\mathrm{DT} = 0.1\mathrm{ms}$ , then 10 time steps = 1ms.