Online Training Through Time for Spiking Neural Networks
Abstract
Spiking neural networks (SNNs) are promising brain-inspired energy-efficient models.
Recent progress in training methods has enabled successful deep SNNs on large-scale tasks with low latency.
Particularly, backpropagation through time (BPTT) with surrogate gradients (SG) is popularly used to enable models to achieve high performance in a very small number of time steps.
However, it is at the cost of large memory consumption for training, lack of theoretical clarity for optimization, and inconsistency with the online property of biological learning rules and rules on neuromorphic hardware.
Other works connect the spike representations of SNNs with equivalent artificial neural network formulation and train SNNs by gradients from equivalent mappings to ensure descent directions.
But they fail to achieve low latency and are also not online.
In this work, we propose online training through time (OTTT) for SNNs, which is derived from BPTT to enable forward-in-time learning by tracking presynaptic activities and leveraging instantaneous loss and gradients.
Meanwhile, we theoretically analyze and prove that the gradients of
OTTT can provide a similar descent direction for optimization as gradients from equivalent mapping between spike representations under both feedforward and recurrent conditions.
OTTT only requires constant training memory costs agnostic to time steps, avoiding the significant memory costs of BPTT for GPU training.
Furthermore, the update rule of OTTT is in the form of three-factor Hebbian learning, which could pave a path for online on-chip learning.
With OTTT, it is the first time that the two mainstream supervised SNN training methods, BPTT with
SG and spike representation-based training, are connected, and meanwhile it is in a biologically plausible form.
Experiments on CIFAR-10, CIFAR-100, ImageNet, and CIFAR10-DVS demonstrate the superior performance of our method on largescale static and neuromorphic datasets in a small number of time steps.
Our code is available at https://github.com/pkuxmq/OTTT-SNN .
Introduction
Spiking neural networks (SNNs) are regarded as the third generation of neural network models [1] and have gained increasing attention in recent years [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13].
SNNs are composed of brain-inspired spiking neurons that imitate biological neurons to transmit spikes between each other.
This allows event-based computation and enables efficient computation on neuromorphic hardware with low energy consumption [14, 15, 16].
However, the supervised training of SNNs is challenging due to the non-differentiable neuron model with discrete spike-generation procedures.
Several kinds of methods are proposed to tackle the problem, and recent progress has empirically obtained successful results.
Backpropagation through time (BPTT) with surrogate gradients (SG) is one of the mainstream methods which enables the training of deep SNNs with high performance on large-scale datasets (e.g., ImageNet) with extremely low latency (e.g., 4-6 time steps) [6, 10, 11, 13].
These approaches unfold the iterative expression of spiking neurons, backpropagate the errors through time [17], and use surrogate derivatives to approximate the gradient of the spiking function [3, 4, 18, 19, 20, 21, 22, 23].
As a result, during training, they suffer from significant memory costs proportional to the number of time steps, and the optimization with approximated surrogate gradients is not well guaranteed theoretically.
Another branch of works builds the closed-form formulation for the spike representation of neurons, e.g. the (weighted) firing rate or spiking time, which is similar to conventional artificial neural networks (ANNs).
Then SNNs can be either optimized by calculating the gradients from the equivalent mappings between spike representations [2, 24, 25, 26, 9, 27], or converted from a trained equivalent
ANN counterpart [28, 29, 30, 31, 32, 7, 33, 8, 34].
The optimization of these methods is clearer than surrogate gradients.
However, they require a larger number of time steps compared to BPTT with SG.
Therefore, they suffer from high latency, and more energy consumption is required if the representation is rate-based.
Another critical point for both methods is that they are indeed inconsistent with biological online learning, which is also the learning rule on neuromorphic hardware [15].
In this work, we develop a novel approach for training SNNs to achieve high performance with low latency, and maintain the online learning property to pave a path for learning on neuromorphic chips.
We call our method online training through time (OTTT).
We first derive OTTT from the commonly used BPTT with SG method by analyzing the temporal dependency and proposing to track the presynaptic activities in order to decouple this dependency.
With the instantaneous loss calculation, OTTT can perform forward-in-time learning, i.e. calculations are done online in time without computing backward through the time.
Then we theoretically analyze the gradients of OTTT and gradients of spike representation-based methods.
We show that they have similar expressions and prove that they can provide the similar descent direction for the optimization problem formulated by spike representation.
For the feedforward network condition, gradients are easily calculated and analyzed.
For the recurrent network condition, we follow the framework in [12] that weighted firing rates will converge to an equilibrium state and gradients can be calculated by implicit differentiation.
With this formulation, the gradients correspond to an approximation of gradients calculated by implicit differentiation, which can be proved to be a descent direction for the optimization problem as well [35, 36].
In this way, a connection between OTTT and spike representation-based methods is bridged.
Finally, we show that OTTT is in the form of three-factor Hebbian learning rule [37], which could pave a path for online learning on neuromorphic chips.
Our contributions include:
We propose online training through time (OTTT) for SNNs, which enables forward-in-time learning and only requires constant training memory agnostic to time steps, avoiding the large training memory costs of backpropagation through time (BPTT).
We theoretically analyze and connect the gradients of OTTT and gradients based on spike representations, and prove the descent guarantee for optimization under both feedforward and recurrent conditions.
We show that OTTT is in the form of three-factor Hebbian learning rule, which could pave a path for on-chip online learning.
With OTTT, it is the first time that a connection between BPTT with SG, spike representation-based methods, and biological three-factor Hebbian learning is bridged.
We conduct extensive experiments on CIFAR-10, CIFAR-100, ImageNet, and CIFAR10DVS, which
demonstrate the superior results of our methods on large-scale static and neuromorphic datasets in a small number of time steps.
Related Work
SNN Training Methods.
As for supervised training of SNNs, there are two main research directions.
One direction is to build a connection between spike representations (e.g. firing rates) of SNNs with equivalent ANN-like closed-form mappings.
With the connection, SNNs can be converted from
ANNs [28, 29, 30, 31, 32, 7, 33, 8, 34, 38], or SNNs can be optimized by gradients calculated from equivalent mappings [2, 24, 25, 26, 9, 27].
Variants following this direction also include [12] which connects feedback SNNs with equilibrium states following fixed-point equations instead of closedform mappings.
These methods have a clearer descent direction for the optimization problem, but require a relatively large number of time steps, suffering from high latency and usually more energy consumption with rate based representation.
Another direction is to directly calculate gradients based on the SNN computation.
They follow the BPTT framework, and deal with the non-differentiable problem of spiking functions by applying surrogate gradients [3, 4, 18, 19, 20, 21, 6, 23, 10, 11, 13], or computing gradients with respect to spiking times based on the linear assumption [39, 40], or combining both [22].
BPTT with SG can achieve extremely low latency.
However, it requires large training memory to maintain the computational graph unfolded along time, and it remains unknown why surrogate gradients work well.
[10] empirically observed that surrogate gradients have a high similarity with numerical gradients, but it remains unclear theoretically.
And gradients based on spiking times suffer from the “dead neuron” problem [3], so they should be combined with
SG in practice [40, 22].
Meanwhile, methods in both directions are inconsistent with biological online learning, i.e. forward-in-time learning, to pave a path for learning on neuromorphic hardware.
Differently, our proposed method avoids the above problems and maintain the online property.
Online Training of Neural Networks.
In the research of recurrent neural networks (RNNs), there are several alternatives for BPTT to enable online learning.
Particularly, real time recurrent learning (RTRL)
[41] proposes to propagate partial derivatives of hidden states over parameters through time to enable forward-in-time calculation of gradients.
Several recent works improve the memory costs of RTRL with approximation for more practical usage [42, 43, 44].
Another work proposes to online update parameters based on decoupled gradients with regularization at each time step [45].
However, these are all for RNNs and not tailored to SNNs.
Several online training methods are proposed for SNNs [46, 47, 48], which are derived in the spirit of RTRL and simplified for SNNs.
[49] leverages local losses and ignores temporal dependencies for online local training of SNNs, and [50] directly apply the method in [45] to train SNNs.
However, these methods also leverage surrogate gradients without providing theoretical explanation for optimization.
Meanwhile, [46, 49] use feedback alignment [51], [47] is limited to single-layer recurrent SNNs, and [48] requires much larger memory costs for eligibility traces, so they cannot scale to large-scale tasks.
[50] requires a specially designed neuron model and more computation for parameter regularization, and also does not consider large tasks.
Differently, our work explain the descent direction under both feedforward and recurrent conditions with convergent inputs, and is efficient and scalable to large-scale tasks including ImageNet classification.
Preliminaries 3.1 Spiking Neural Networks
Spiking neurons are brain-inspired models that transmit information by spike trains.
Each neuron maintains a membrane potential u and integrates input spike trains, which will generate a spike once u exceeds a threshold.
We consider the commonly used leaky integrate and fire (LIF) model, which describes the dynamics of the membrane potential as:
[Equation]
where I is the input current, V th is the threshold, and R and τ m are resistance and time constant, respectively.
A spike is generated when u reaches V th at time t f , and u is reset to the resting potential u = u rest , which is usually set to be zero.
The output spike train is defined using the Dirac delta function: s ( t )
= (cid:80) t f δ ( t − t f ) .
A spiking neural network is composed of connected spiking neurons with connection coefficients.
We consider a simple current model
I i ( t ) =
(cid:80)
j w ij s j ( t )
+ b i , where the subscript i represents the i -th neuron, w ij is the weight from neuron j to neuron i , and b i is a bias.
The discrete computational form is:
[Equation]
where H ( x ) is the Heaviside step function, s i [ t ] is the spike train of neuron i at discrete time step t , and λ < 1 is a leaky term (typically taken as 1 − 1 τ m ).
The constant R , τ m , and time step size are absorbed into the weights and bias.
The reset operation is implemented by subtraction.
3.2 Previous SNN Training Methods
Spike Representation.
The (weighted) firing rate or first spiking time of spiking neurons can be proved to follow ANN-like closed-form transformations [7, 26, 9, 12, 27].
We focus on the weighted firing rate [12, 27] which has connection with OTTT in this work.
Define weighted firing rates and weighted average inputs a [
[Equation]
=
[Equation] [Equation]
in the discrete condition.
Given convergent weighted average inputs x [ t ] → x ∗ , it can be proved that a [ t ] → a ∗ = σ (cid:16) 1 V th x ∗ (cid:17) with bounded random error, where σ is a clamp function ( σ ( x ) = min(max(0 , x ) , 1) ) in the discrete condition while a ReLU function in the continuous condition.
For feedforward networks, the closedform mapping between successive layers is established based on weighted firing rate after time T : a l +1 [ T ] ≈ σ (cid:16) 1 V th (cid:0) W l a l [ T ] + b l +1 (cid:1)(cid:17) , and gradients are calculated with such spike representation: ∂L∂
[Equation]
=
[Equation] [Equation]
[Equation]
W l .
For the recurrent condition, a [ t ] will converge to an equilibrium state following an implicit fixed-point equation, e.g. a ∗ = σ (cid:16)
[Equation] [Equation]
[Equation]
for a single-layer network with input connections F and contractive recurrent connections W , and gradients can be calculated based on implicit differentiation [12].
Let a = f θ ( a ) denote the fixed-point equation ( θ are parameters).
We have ∂L ∂ θ =
[Equation]
a [ T ] (cid:0) I − J f
θ | a [ T ] (cid:1) − 1 ∂f θ ( a [ T ]) ∂ θ , where
J f θ | a [ T ] =
∂f
[Equation]
[ T ])
[Equation]
a [ T ] is the Jacobian of f θ at a [ T ] .
BPTT with SG.
BPTT unfolds the iterative update equation in Eq.(2) and backpropagates along the computational graph as shown in Fig. 1(a), 1(c).
The gradients with T time steps are calculated by 2 : ∂L
[Equation]
l
[Equation]
l
[Equation] [Equation]
l
[Equation]
[ i ] (cid:19)
[Equation] [Equation]
[Equation]
(cid:33) , (3) where W l is the weight from layer l to l + 1 and L is the loss.
The non-differentiable terms ∂ s l [ t ] ∂ u l [ t ] will be replaced by surrogate derivatives, e.g. derivatives of rectangular or sigmoid functions [4]: ∂s∂u = 1 a 1 sign (cid:0)
| u − V th | < a 1 2 (cid:1) or ∂s∂u = 1 a
[Equation]
1+ e ( Vth
[Equation]
) /a 2 ) 2 , where a 1 and a 2 are hyperparameters.
Online Training Through Time for SNNs
This section contains four sub-sections.
In Section 4.1, we introduce our proposed OTTT by decoupling the temporal dependency of BPTT.
Then in Section 4.2, we further connect the gradients of OTTT and spike representation-based methods, and prove that OTTT can provide a descent direction for optimization, which is not guaranteed by BPTT with SG.
In Section 4.3, we discuss the relationship between OTTT and the three-factor Hebbian learning rule.
Implementation details are presented in Section 4.4.
4.1 Derivation of Online Training Through Time
Decouple temporal dependency.
As shown in Fig.
1(c), BPTT has to maintain the computational graph of previous time to backpropagate through time.
We will decouple such temporal dependency to enable online gradient calculation, as illustrated in Fig. 1(d).
We first focus on the feedforward network condition.
In this setting, all temporal dependencies lie in the dynamics of each spiking neuron, i.e. ∂ u l +
[ i +1]
∂ u l +
[ i ] and ∂ u l +
[ i +1] ∂
s l +
[ i ] ∂ s l +1 [ i ] ∂ u l
+
[ i ] in Eq.(3).
We consider the case that we do not apply surrogate derivatives to ∂ s l +
[ i ] ∂ u l +
[ i ] in such temporal dependency.
Since the derivative of the Heaviside step function is 0 almost everywhere, we have ∂ u l
+
[ i +1] ∂
s l +
[ i ] ∂ s l +1 [ i ] ∂ u l
+
[ i ] ≈ 0 3 .
𝑠 𝑡−1 𝑙+ 𝑠 𝑡𝑙+ 𝑠
𝑡+1 𝑙+ 𝑢 𝑡−1𝑙+ 𝑢 𝑡 𝑙+ 𝑢 𝑡+
[Figure]
𝑙
[Figure]
𝑠 𝑡−1𝑙 𝑠 𝑡 𝑙 𝑠
[Figure]
𝑊 𝑙 𝑊 𝑙 𝑊 𝑙
Time
[Figure]
Then the dependency only includes ∂ u l
[Figure] [Figure]
1] ∂ u
[Figure]
+
[ i ] , which equals λ I .
Therefore, we have 2 : ∂L
[Equation]

[Equation]
t
[Equation]
g u l
[Equation] [Equation]
]  (cid:88)
[Equation]
≤ t λ
[Equation]
[ τ ] 
 ⊤ , (4) where g u l +1 [ t ] = (cid:16) ∂L
∂ s l +1 [ t ] ∂ s l +
[ t ] ∂ u l +
[ t ] (cid:17) ⊤ is the gradient for u l +1 [ t ] .
Based on Eq.(4), we can track presynaptic activities ˆ a l [ t ] = (cid:80)
τ ≤ t λ t − τ s l [ τ ] for each neuron during the forward procedure by ˆ a l [ t + 1] =
λ ˆ a l [ t ] + s l
[ t + 1] , so that when given g u l +1 [ t ] , gradients at each time step can be calculated independently by ∇ W l L [ t ] =
g u l +
[ t ]ˆ a l [ t ] ⊤ without backpropagation through ∂ u l +
[ i +1]
∂ u l +
[ i ] .
As for the recurrent network condition, there are additional temporal dependencies due to the feedback connections between neurons.
If there is feedback connection from layer l 2 to l 1 ( l 2 ≥ l 1 ), there would be terms such as ∂ u l 1 [ i +1]
∂ s l 2 [ i ] ∂ s l 2 [ i ] ∂
u l 2 [ i ] ∂ u l 2 [ i ] ∂ u l 1 [ i ] in the calculation of gradients (note that Eq.
(3) omit feedback connections for simplicity).
We also consider not applying surrogate derivatives to ∂ s l
[ i ] ∂ u l 2 [ i ] in the temporal dependency so that gradients are not calculated in this path.
Similar to the feedforward condition, we can derive that the gradients of the general weight W l i → l j from any layer l i to any layer l j can be calculated by ∇ W li → lj L [ t ] =
g u lj [ t ]ˆ a l i [ t ] ⊤ at each time step.
A theoretical explanation for optimization will be presented in Section 4.2.
Instantaneous Loss and Gradient.
Calculating online gradients, e.g. the above g u l +
[ t ] for ∇ W l L [ t ] , requires instantaneous computation of the loss at each time step.
Previous typical loss for
SNNs is based on the firing rate, e.g. L fr =
L (cid:16) 1 T (cid:80)
Tt =
1 s N [ t ] , y (cid:17) , where y is the label, s N [ t ] is the spike at the last layer, and L can take cross-entropy loss.
This loss depends on all time steps and does not support online gradients.
We leverage the instantaneous loss and calculate the above g u l +1 [ t ] as:
[Equation]
The total loss L := (cid:80) Tt
=1 L [ t ] is the upper bound of L fr when L is a convex function such as cross-entropy.
Then gradients can be calculated independently at each time step, as shown in Fig. 1(d).
We apply surrogate derivatives for ∂ s l [ t ] ∂ u l [ t ] in this calculation, which will be explained in Section 4.2.
Since gradients are calculated instantaneously at each time step, OTTT does not require maintaining the unfolded computational graph and only requires constant training memory costs agnostic to time steps.
Note that instantaneous gradients of OTTT will be different from BPTT with the instantaneous loss for multi-layer or recurrent networks, as we do not consider future influence in the instantaneous calculation: BPTT considers terms such as ∂L [ t ′ ]
∂ u N [ t ′ ] ∂ u N [ t ′ ]
∂ u N [ t ] (cid:81)
l i = N − 1 ∂
u
i +
[ t ] ∂ u i [ t ] ( t ′ > t ) for u l [ t ] while we do not.
The equivalence of OTTT and BPTT only holds for the last layer, and we do not seek the exact equivalence to BPTT with SG which is theoretically unclear, but will build the connection with spike representations and prove the descent guarantee.
Also, note that the tracked presynaptic activities are similar to the biologically plausible “eligibility traces” in the literature [46, 47, 52], and we will provide a more solid theoretical grounding for optimization in Section 4.2.
4.2 Connection with Spike Representation-Based Methods for Descent Directions
In this section, we connect gradients of OTTT and spike representation-based methods, and prove that
OTTT can provide a descent direction for optimization under feedforward and recurrent conditions with convergent inputs.
Feedforward Networks.
As introduced in Section 3.2, with convergent inputs, methods based on spike representations establish closed-form mappings between successive layers with weighted firing rate
[Equation]
as
[Equation] [Equation]
+1 (cid:1)(cid:17) , and calculate gradients by ∂L∂
[Equation]
=
[Equation]
N [ T ] (cid:81) l +1 i
= N − 1 ∂ a
i +1 [ T ] ∂ a i
[ T ] ∂ a l
+
[ T ] ∂ W l .
Note that a [ t ] is similar to the tracked presynaptic activities ˆ a [ t ] = (cid:80) tτ
=1 λ t − τ s [ τ ] in OTTT.
We can obtain:
[Equation]
where L sr is the loss based on spike representation, d l +1 [ T ] = σ ′ (cid:16) 1 V th (cid:0) W l a l [ T ] + b l +1 (cid:1)(cid:17) , and ‘ ⊙ ’ is element-wise product.
The detailed derivation can be found in Appendix A.
It can be easily seen that Eq.
(6) has a similar form as gradients in Eq.
(4), (5), and we build the connection between them in three steps.
For sake of clarify, in the following, we denote gradients of
OTTT and spike representation by ∇ W l L and ( ∇ W l L sr ) sr , respectively.
In the first step, we can use appropriate surrogate derivatives of ∂ s l +1 i
[ t ] ∂ u l +1 i
[ t ] ( t = 1 , · · · , T ) to approximate d l +1 i
[ T ] .
As introduced in Section 3.2, σ is a clamp function ( σ ( x ) =
min(max(0 , x ) , 1) ) in the discrete condition while a ReLU function in the continuous condition.
Then d l +1 [ T ] almost equals sign (cid:0) | W l a l [ T ] + b l
+1 − V
th | < V th (cid:1) if we slightly relax the clamp bound caused by the discretization, and this can be approximated by diag (cid:16) ∂ s l +1 [ t ] ∂ u l +1 [ t ] (cid:17) =
sign (cid:0) | u l +1 [ t ] − V th | < V th (cid:1) at each time step.
Note that surrogate derivatives here approximate the well-defined derivative of the mapping function between a [ T ] , rather than the pseudo derivative of the non-differentiable Heaviside step function as in BPTT with SG.
The approximation is exact except the rare case that averagely a neuron generate spikes (i.e. the average input is positive) while sometimes the membrane potential is less than u rest (or the reverse).
For simplicity in the following theoritical analysis, we assume the equivalence between surrogate derivatives and d [ T ] .
Assumption
[Equation]
In the second step, we take L sr in Eq.
(6) as L srup =
[Equation]
t L ( s N [ t ] , y ) so that 1 T 1 λ T − t ∂L sr ∂ s N [ t ] in Eq.
(6) aligns with ∂L [ t ] ∂ s N [ t ] in Eq.
(5) except a constant term.
Note that L srup is an upper bound of the common loss L ′
sr
=
L ( a N [ T ] , y )
=
L ( (cid:80) Tt =
1 λ T − t s N [ t ] (cid:80) t t =1 λ T − t , y )
if L is a convex function.
Then let ˆ g u l +1 [ t ] = (cid:16)
∂ L ( s N [ t ] , y ) ∂
s N [ t ] (cid:81)
l +1 i
= N − 1 ∂
s i +
[ t ] ∂ s i [ t ] ∂ s l +
[ t ] ∂ u l +
[ t ] (cid:17) ⊤ , with Assumption 1 we have: ∇
[Equation]
[ t ]
[Equation]
l [ t ] ⊤ and ( ∇ W l L sr ) sr
=
[Equation]
Please refer to Appendix A for details.
In the third step, we handle the remaining difference that ∇ W l L leverages instantaneous presynaptic activities ˆ a [ t ] during calculation while ( ∇ W l L sr ) sr uses the final ˆ a l [ T ] after time T .
Note that the weighted firing rate gradually converges a [ t ] → a ∗ with bounded random error.
Suppose the errors ϵ l [ t ] = a l [ t ] − a l [ T ] are small ( l = 0 represents inputs, i.e. a 0 [ t ] = x
[ t ] ), then we have that −∇ W l L can provide a descent direction, as shown in Theorem 1.
Theorem 1.
If Assumption 1 holds, V th = 1 , and the errors ϵ l [ t ] = a l [ t ] − a l [ T ] are small such that (cid:13)(cid:13)(cid:13) (cid:80)
Tt
=
ˆ g u l +
[ t ] ϵ l [ t ] ⊤ (cid:13)(cid:13)(cid:13) < (cid:13)(cid:13)(cid:13) (cid:80)
Tt
=
ˆ g u l +
[ t ] a l [ T ] ⊤ (cid:13)(cid:13)(cid:13) − (cid:13)(cid:13)(cid:13) (cid:80)
Tt =
1 λ t (1 − λ T − t ) 1 − λ T
ˆ g u l +1 [ t ] a l [ t ] ⊤ (cid:13)(cid:13)(cid:13) when ( ∇ W l L sr )
sr ̸ = 0 , then we have ⟨∇ W l L, ( ∇ W l L sr ) sr ⟩ > 0 .
For the proof and discussion of the assumption please refer to Appendix A. With this conclusion, we can explain the descent direction of gradient descent by OTTT for the optimization problem formulated by spike representation.
Some random error can be viewed as randomness for stochastic optimization.
Recurrent Networks.
For networks with feedback connections, we first consider the single-layer condition for simplicity (see Appendix A for general conditions).
We consider feedforward connections F from inputs to neurons and contractive recurrent connections W between neurons.
As introduced in Section 3.2, given convergent inputs x [ t ] → x ∗ , a [ t ] of neurons will converge to an equilibrium state a ∗ =
f θ ( a ∗ ) =
σ (cid:16) 1 V th ( Wa ∗ + Fx ∗ + b ) (cid:17) with bounded random error, and gradients are calculated as ( ∇
[Equation]
where θ ∈ { W , F , b } .
We consider replacing the inverse Jacobian by an identity matrix:
[Equation]
(cid:17) ⊤ .
Previous works have proved that this gradient can provide a descent direction for the optimization problem [35, 36], i.e. (cid:68)
[Equation]
It has a similar form as the OTTT gradient: ∇ W L = (cid:80)
Tt
=1 g u [ t ]ˆ a [ t ] ⊤ and (cid:94) ( ∇ W L sr )
sr = (cid:80)
Tt =1 (cid:16) 1 T
[Equation]
s [ t ] ⊤ ⊙
d [ T ] (cid:17) ˆ a [ T ] ⊤ .
Similarly, we can prove the descent guarantee for OTTT as shown in Theorem 2 .
For details refer to Appendix A.
Theorem 2.
If Assumption 1 holds, V th = 1 , (cid:13)(cid:13)
J f θ | a [ T ] (cid:13)(cid:13) ≤ η < σ 2 min
[Equation]
2 max , where σ max and σ min are the maximal and minimal singular value of ∂f θ ∂ θ | a [ T ] , and the errors ϵ 1 [ t ] = a [ t ] − a [ T ] , ϵ 0 [ t ] = x [ t ] − x [ T ] are small such that (cid:13)(cid:13)(cid:13)(cid:80)
[Equation]
(cid:13)(cid:13)(cid:13)(cid:80)
[Equation]
ˆ g u [ t ] a l [ t ] ⊤ (cid:13)(cid:13)(cid:13) (where l = 0 , 1 , a 1 [ t ] and a
[ t ] represent a [ t ] and x [ t ] , respectively) when ( ∇ θ L sr )
sr ̸ = 0 , then we have ⟨∇ θ L, ( ∇ θ L sr ) sr ⟩ > 0 , where θ are parameters in the network.
4.3 Connection with Three-factor Hebbian Learning Rule
By explicitly writing the instantaneous gradients of OTTT for the general weight from layer l i to l j , ∇
[Equation]
→
[Equation]
[ t ] = g u lj [ t ]ˆ a l i [ t ] ⊤ , and dive into connections between any two neurons i and j , we have:
[Equation]
(7) where ˆ a i [ t ] is the tracked presynaptic activity, f ( u j [ t ]) is the surrogate derivative function which can represent the change rate of the postsynaptic activity as analyzed in Section 4.2, and δ j [ t ] =
g s j [ t ] is the gradient for neuron output s j [ t ] which represents a global modulator.
This is a kind of three-factor Hebbian learning rule [37] and the weight can be updated locally with a global signal.
The error signal δ j [ t ] can be propagated in an error feedback path simultaneously with feedforward propagation, which is shown biologically plausible with high-frequency bursts [53].
Note that the analysis in
Section 4.2 still holds if we consider the delay of the propagation of the error signal, i.e. the update is based on ˆ a i [ t + ∆ t ] f ( u j [ t + ∆ t ])
δ j [ t ] .
4.4 Implementation Details
As introduced in Section 4.1, we will calculate instantaneous gradients ∇ W li → lj L [ t ] =
g u lj [ t ]ˆ a l i [ t ] ⊤ at each time step.
We can choose to immediately update parameters before the calculation of the next time step, which we denote as OTTT O , or we can accumulate the gradients by T time steps and then update parameters, which we denote as OTTT A .
For OTTT O , we assume that the online update is small and has negligible affects for the following calculation.
Pseudo-codes are in Appendix B.
An important issue in practice is that previous BPTT with SG works leverage batch normalization (BN) along the temporal dimension to achieve high performance with extremely low latency on large-scale datasets [6, 10, 11, 13], which requires calculating the mean and variance statistics for all time steps during the forward procedure.
This technique intrinsically prevents online gradients and has to suffer from large memory costs.
To overcome this shortcoming, we do not use BN, but borrow the idea from normalization-free ResNets (NF-ResNets) [54, 55] to replace batch normalization by scaled weight standardization (sWS) [56].
sWS standardizes weights by ˆ W i,j = γ · W i,j − µ W i, · σ W i, · √ N , and the scale γ is determined by analyzing the signal propagation with different activation functions.
We apply sWS for VGG [57] and NF-ResNet architectures in our experiments.
For details please refer to
Appendix C.
Experiments
In this section, we conduct extensive experiments on CIFAR-10 [58], CIFAR100 [58], ImageNet [59],
CIFAR10-DVS [60], and DVS128-Gesture [61] to demonstrate the superior performance of our proposed method on large-scale static and neuromorphic datasets.
We leverage the VGG network architecture (64C3-128C3-AP2-256C3-256C3-AP2-512C3-512C3-AP2-512C3-512C3-GAP-FC) for experiments on CIFAR-10, CIFAR-100, CIFAR10-DVS, and DVS128-Gesture, and the NF-ResNet34 [54] network architecture for experiments on ImageNet.
For all our SNN models, we set V
th =and λ = 0 .
Please refer to Appendix C for training details.
5.1 Comparison of Training Memory Costs
[Figure]
A major advantage of OTTT over BPTT is that OTTT does not require backpropagation along the temporal dimension and therefore only requires constant training memory costs agnostic to time steps, which avoids the large memory costs of BPTT.
We verify this by training the VGG network on CIFAR-10 with batch size 128 under different time steps and calculating the memory costs on the GPU.
As shown in Fig.
2, the training memory of BPTT grows linearly with time steps, while
OTTT maintains the constant memory (both OTTT A and OTTT O ).
Even with a small number of time steps, e.g. 6, OTTT can reduce the memory costs by 2 ∼ 3 × .
This advantage may also allow training acceleration of SNNs by larger batch sizes with the same computational resources.
5.2 Comparison of Performance
We conduct experiments on both large-scale static and neuromorphic datasets.
We first verify the effectiveness of the surrogate derivative ∂ s l +1 [ t ] ∂ u l +
[ t ] = sign (cid:0) |
u l +1 [ t ] − V th | < V th (cid:1) in Section 4.2.
For the VGG (sWS) network on CIFAR-10, OTTT A and OTTT O achieves 91.46% and 92.47% test accuracy respectively.
We empirically observe that applying the sigmoid-like surrogate derivative achieves a higher generalization performance, e.g. OTTT A and OTTT O achieves 93.58% and 93.73% test accuracy respectively under the same random seed, and a possible reason may be that this introduces some noise for the approximation to regularize the training and improve the generalization.
Therefore, in the following performance evaluation, we take sigmoid-like surrogate derivative for OTTT.
As shown in Table 1, both OTTT A and OTTT O achieve satisfactory performance on all datasets, and compared with BPTT under the same training settings, OTTT achieves higher performance.
The proposed OTTT also achieves promising performance on all datasets compared with other representative conversion and direct training methods.
Besides, it shows that the performance gap between our SNN model and ANN is around 0.7% and 2.08% on CIFAR-10 and CIFAR-100, respectively.
Usually, SNNs with a very small number of time steps do not reach the performance of equivalent ANNs due to the information propagation with discrete spikes rather than floating-point numbers.
The results of our model with 6 time steps are competitive.
We also evaluate our method on the DVS128-Gesture dataset, which is more time-varying with different hand gestures recorded by a DVS camera.
As shown in Table 2, our method can achieve the same high performance as BPTT.
While our theoretical analysis mainly focus on convergent inputs (e.g. static images or neuromorphic inputs converted from images like CIFAR10-DVS), the results show that our method can also work well for time-varying inputs.
[Table] [Table]
5.3 Effectiveness for Recurrence
As introduced in Section
[Table]
the proposed OTTT is also valid for networks with feedback connections.
Previous works have shown that adding feedback connections can improve the performance of SNNs without much additional costs, especially on the CIFAR-100 datasets [12, 63].
Therefore, we conduct experiments on CIFAR-100 with the VGG-F network architecture which simply adds a feedback connection from the last feature layer to the first feature layer following [12], and this weight is zero-intialized.
As shown in Table 3, the training of VGG-F is valid and VGG-F achieves a higher performance than VGG due to the introduction of feedback connections.
Results in Appendix D show that the improvement of OTTT from feedback connections is more significant than that of BPTT.
The architectures with feedback connections can be further improved with neural architecture search [63].
5.4 Effectiveness for Training with Batch Size
To further study the online training, i.e. not only online in time but also one sample per training, which is consistent with biological learning and learning on neuromorphic hardware, we verify the effectiveness for training with batch size 1.
The VGG network on CIFAR-10 is studied, and batch size 1 is compared with the default batch size 128 under the same random seed.
Models are only trained for 20 epochs due to the relatively long training time with batch size 1.
As shown in Table 4, training with one sample per iteration is still valid, indicating the potential to conduct full online training with the proposed OTTT.
[Table]
[Table]
5.5 Influence of Inference Time Steps
We study the influence of inference time steps on ImageNet as shown in Fig.
It illustrates that the model trained with time step 6 can achieve higher performance with more inference time steps.
5.6 Firing Rate Statistics
We study the firing rate statistics of the models trained by OTTT and BPTT, as shown in Fig.
It demonstrates that models trained by OTTT have higher firing rates in first layers while lower firing rates in later layers compared with BPTT.
Overall the firing rate is around 0.19 and with 6 time steps each neuron averagely generate 1.1 spikes, indicating the low energy consumption.
Considering that each neuron has more synaptic operations in later layers than first layers (because the channel size is increasing with layers), the synaptic operations of models trained by OTTT and BPTT are about the same ( 1 . 98 × 10 8 vs 1 .
93 × 10 8 ).
More results please refer to Appendix D.
[Table]
[Figure]
[Table]
[Figure] [Figure]
[Table]
Conclusion
In this work, we propose a new training method, online training through time (OTTT), for spiking neural networks.
We first derive OTTT from BPTT with SG by decoupling the temporal dependency with the tracked pre-synaptic activities, which only requires constant training memory agnostic to time steps and avoids the large training memory costs of BPTT.
Then we theoretically analyze and connect the gradients of OTTT and gradients of methods based on spike representations, and prove the descent guarantee of OTTT for the optimization problem under both feedforward and recurrent network conditions.
Additionally, we show that OTTT is in the form of three-factor Hebbian learning rule, which is the first to connect BPTT with SG, spike representation-based methods, and biological learning rules.
Extensive experiments demonstrate the superior performance of our methods on large-scale static and neuromorphic datasets in a small number of time steps.
Acknowledgement
Z. Lin was supported by the major key project of PCL (No. PCL2021A12), the NSF China (No.s 62276004 and 61731018), and Project 2020BD006 supported by PKU-Baidu Fund.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contribu- tions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] See Appendix E. (c) Did you discuss any potential negative societal impacts of your work?
[Yes] See Appendix E. (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] See Section 4.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix A. 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] See Supplementary Materials.
(b) Did you specify all the training details (e g., data splits, hyperparameters, how they were chosen)?
[Yes] See Appendix C. (c) Did you report error bars (e g., with respect to the random seed after running experiments multiple times)?
[Yes] See Section 5.
(d) Did you include the total amount of compute and the type of resources used (e g., type of GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix C. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] See Section 5.
(b) Did you mention the license of the assets?
[Yes] See Appendix C. (c) Did you include any new assets either in the supplemental material or as a URL?
[N/A] We do not use new assets.
(d) Did you discuss whether and how consent was obtained from people whose data you’re us- ing/curating?
[No] We use the existing common datasets.
(e) Did you discuss whether the data you are using/curating contains personally identifiable informa- tion or offensive content?
[No] We use the existing common datasets.
[-+
-]
[-+
-]
[-1 [t] ∂ u l+
-]
If you
[-=
-]
[-ˆ g u l +
-]
[-With Assumption 1, we have ∂ s l +
-]
[-As for other parameters such as biases b l , the gradients of OTTT do not involve pre-synaptic activities, so under Assumptionthey are exactly the same as gradients based on spike representation except a constant scaling factor 1 (cid:80) T − 1 τ =
-]
[-J f θ | a [ T ] (cid:13)(cid:13) ≤ η < σ 2 min σ 2 max , where σ max and σ min are the maximal and minimal singular value of ∂f θ ∂ θ | a [ T ] , and the errors ϵ 1 [ t ] = a [ t ] − a [ T ] , ϵ 
-]
[-(where l = 0 , 1 , a 1 [ t ] and a 
-]
[-=
 ˆ g u [ t ] . For gradients based on spike representation, ( ∇ θ L sr ) sr = (cid:16)-]crowdsourcing or conducted research with human subjects...
[-=
-]
(a) Did you include the full text of instructions given to participants and [-with Assumption 1 which indicates ∂ s l +
 i [ t ] ∂ u l +1 i [ t ] =-]screenshots, if applicable?
[-= 1 T (cid:80) T t =
-]
[-T 1 (cid:80) T − 1 τ =
-]
[-=
-]
[-(cid:16)(cid:80) T − 1 τ =
-]
[-T − 1 τ =
-]
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[-Var ( x l ) + Var ( αf l ( x l /β l )) , so dividing β l ensures that the residual branch keeps the identity variance 1 (combined with sWS), and this also indicates to calculate β l by β 2 l +
-]
[-DVS-CIFAR
-]
[N/A]