Online Training Through Time for Spiking Neural Networks
Abstract
Spiking neural networks (SNNs) are promising brain-inspired energy-efficientmodels.
Recent progress in training methods has enabled successful deep SNNson large-scale tasks with low latency.
Particularly, backpropagation through time(BPTT) with surrogate gradients (SG) is popularly used for its high performancewith extremely low latency.
However, it is at the cost of large memory consumptionfor training, theoretical unclarity for optimization, and inconsistency with biologicalonline learning.
Other works connect the spike representations of SNNs withequivalent artificial neural network formulation and train SNNs by gradients fromequivalent mappings to ensure descent directions.
But they fail to achieve lowlatency and are also inconsistent with online learning which is the learning ruleon neuromorphic hardware.
In this work, we propose online training throughtime (OTTT) for SNNs, which is derived from BPTT to enable forward-in-timelearning by tracking presynaptic activities and leveraging instantaneous loss andgradients.
Meanwhile, we theoretically analyze and prove that the gradients of
OTTT can provide a similar descent direction for optimization as gradients fromequivalent mapping between spike representations under both feedforward andrecurrent conditions.
OTTT only requires constant training memory costs agnosticto 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 Hebbianlearning, which could pave a path for online on-chip learning.
With OTTT, it is thefirst time that the two mainstream supervised SNN training methods, BPTT with
SG and spike representation-based training, are connected, and meanwhile it is in abiologically 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.
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].
SNNsare composed of brain-inspired spiking neurons that imitate biological neurons to transmit spikesbetween each other.
This allows event-based computation and enables efficient computation onneuromorphic hardware with low energy consumption [14, 15, 16].
However, the supervised training of SNNs is challenging due to the complex neuron model withdiscrete 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 trainingof deep SNNs with high performance on large-scale datasets (e.g., ImageNet) with extremely lowlatency (e.g., 4-6 time steps) [6, 10, 11, 13].
These approaches unfold the iterative expressionof 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, duringtraining, they suffer from significant memory costs proportional to the number of time steps, and theoptimization with approximated surrogate gradients is not well guaranteed theoretically.
Anotherbranch 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 equivalentmappings 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 clearerthan surrogate gradients.
However, they require a larger number of time steps compared to BPTTwith SG.
Therefore, they suffer from high latency, and more energy consumption is required if therepresentation is rate-based.
Another critical point for both methods is that they are indeed inconsistentwith 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 lowlatency, 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 commonlyused BPTT with SG method by analyzing the temporal dependency and proposing to track thepresynaptic activities in order to decouple this dependency.
With the instantaneous loss calculation, OTTT can perform forward-in-time learning without backpropagtion along the time dimension.
Thenwe 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 descentdirection for the optimization problem formulated by spike representation.
For the feedforwardnetwork 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 andgradients can be calculated by implicit differentiation.
With this formulation, the gradients correspondto an approximation of gradients calculated by implicit differentiation, which can be proved to be adescent 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 theform of three-factor Hebbian learning rule [37], which could pave a path for online learning onneuromorphic 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 spikerepresentations, 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 apath 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 withequivalent ANN-like closed-form mappings.
With the connection, SNNs can be converted from
ANNs [28, 29, 30, 31, 32, 7, 33, 8, 34], or SNNs can be optimized by gradients calculated fromequivalent mappings [2, 24, 25, 26, 9, 27].
Variants following this direction also include [12] whichconnects feedback SNNs with equilibrium states following fixed-point equations instead of closedform mappings.
These methods have a clearer descent direction for the optimization problem, butrequire a relatively large number of time steps, suffering from high latency and usually more energyconsumption with rate based representation.
Another direction is to directly calculate gradients basedon the SNN computation.
They follow the BPTT framework, and deal with the non-differentiableproblem 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 [38, 39], orcombining 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 remainsunknown why surrogate gradients work well.
[10] empirically observed that surrogate gradientshave a high similarity with numerical gradients, but it remains unclear theoretically.
And gradientsbased on spiking times suffer from the “dead neuron” problem [3], so they should be combined with
SG in practice [39, 22].
Meanwhile, methods in both directions are inconsistent with biologicalonline 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 recurrentlearning (RTRL) [40] proposes to propagate partial derivatives of hidden states over parametersthrough time to enable forward-in-time calculation of gradients.
Several recent works improve thememory costs of RTRL with approximation for more practical usage [41, 42, 43].
Another recentwork proposes a different method to online update parameters based on decoupled gradients at eachtime step [44].
However, these are all for RNNs and not tailored to SNNs.
Several online trainingmethods are proposed for SNNs [45, 46], which are derived in the spirit of RTRL and simplified for
SNNs.
However, they also leverage surrogate gradients without providing theoretical explanation foroptimization, and [45] use feedback alignment [47] while [46] is limited to single-layer recurrent SNNs, so they cannot scale to large-scale tasks.
Differently, our work explain the descent directionunder both feedforward and recurrent conditions, and is scalable to large-scale tasks including ImageNet classification.
Preliminaries
[Figure]
3.1 Spiking Neural Networks
Spiking neurons are brain-inspired models that transmit information by spike trains.
Each neuronmaintains a membrane potential u and integrates input spike trains, which will generate a spike onceu exceeds a threshold.
We consider the commonly used leaky integrate and fire (LIF) model, whichdescribes 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 potentialu =
u rest , which is usually set to be zero.
The output spike train is defined using the Dirac deltafunction: 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 thei -th neuron, w ij is the weight from neuron j to neuron i , and b i is a bias.
The discrete computationalform 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 areabsorbed into the weights and bias.
The reset operation is implemented by subtraction.
3.2 Previous SNN Training Methods
BPTT with SG.
BPTT unfolds the iterative update equation in Eq.(2) and backpropagates along thecomputational graph as shown in Fig. 1(a), 1(c). The gradients with T time steps are calculated by 1 :∂L [Equation]
[ t ] {+∂ s l +1 [ t ] ∂ u l +
 [ t ] (cid:32) ∂+}
u l +1 [ t ] ∂ W l + (cid:88) τ<t τ (cid:89) i = t −(cid:18)
[-=
-]∂
{+u l +
+} [ i + 1] ∂ u l
[-=
-]{++
+}
{+[ i ] + ∂ u l +
+}
[ i + 1] ∂
s l {++
+}
[ i ] ∂ s l +1 [ i ] ∂ u l
+
[ i ] (cid:19) ∂ u l +1 [ τ ] ∂ W l {+(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 ae ( Vth − u ) /a(
+}
+ e ( Vth − u ) /a 2 ) 2 , where a 1 and a 2 are hyperparameters.
𝑠 𝑡−1 𝑙+ 𝑠 𝑡𝑙+ 𝑠
𝑡+1 𝑙+ 𝑢 𝑡−1𝑙+ 𝑢 𝑡 𝑙+ 𝑢 𝑡+1𝑙+ 𝑠 𝑡−1𝑙 𝑠 𝑡 𝑙 𝑠 𝑡+1𝑙 𝑊 𝑙 𝑊 𝑙 𝑊 𝑙
Time (a) BPTT Forward 𝑠 𝑡−1 𝑙+ 𝑠 𝑡𝑙+ 𝑠
𝑡+1 𝑙+ 𝑢 𝑡−1 𝑙+ 𝑢 𝑡 𝑙+ 𝑢 𝑡+1 𝑙+ 𝑠 𝑡−1𝑙 𝑠 𝑡𝑙 𝑠 𝑡+1𝑙
𝑊 𝑙 𝑊 𝑙 𝑊 𝑙
Time ො𝑎 𝑡−1𝑙 ො𝑎 𝑡𝑙 ො𝑎 𝑡+1𝑙 (b) OTTT Forward 𝜕𝐿 𝜕𝑠 𝑡−1 𝑙+
Time 𝜕𝐿 𝜕𝑠 𝑡𝑙+ 𝜕𝐿 𝜕𝑠 𝑡+1𝑙+ 𝜕𝐿 𝜕𝑠 𝑡−1𝑙
𝜕𝐿 𝜕𝑠 𝑡𝑙 𝜕𝐿 𝜕𝑠 𝑡+1𝑙
𝜕𝐿𝜕𝑊 𝑙 𝜕𝐿 𝜕𝑊 𝑙 𝜕𝐿𝜕𝑊 𝑙 𝜕𝐿 𝜕𝑢 𝑡−1𝑙+ 𝜕𝐿 𝜕𝑢 𝑡𝑙+ 𝜕𝐿 𝜕𝑢 𝑡+1 𝑙+ 𝑠 𝑡−1 𝑙 𝑠 𝑡𝑙 𝑠 𝑡+1 𝑙 (c) BPTT Backward 𝜕𝐿 𝜕𝑠 𝑡−1 𝑙+
Time 𝜕𝐿 𝜕𝑠 𝑡𝑙+ 𝜕𝐿 𝜕𝑠 𝑡+1𝑙+ 𝜕𝐿 𝜕𝑠 𝑡−1𝑙
𝜕𝐿 𝜕𝑠 𝑡𝑙 𝜕𝐿 𝜕𝑠 𝑡+1𝑙
𝜕𝐿𝜕𝑊 𝑙 𝜕𝐿 𝜕𝑊 𝑙 𝜕𝐿𝜕𝑊 𝑙 𝜕𝐿 𝜕𝑢 𝑡−1𝑙+ 𝜕𝐿 𝜕𝑢 𝑡𝑙+ 𝜕𝐿 𝜕𝑢 𝑡+1 𝑙+ ො𝑎 𝑡−1 𝑙 ො𝑎 𝑡𝑙 ො𝑎 𝑡+1 𝑙 (d) OTTT Backward
Spike Representation.
The (weighted) firing rate or first spiking time of spiking neurons canbe proved to follow ANN-like closed-form transformations [7, 26, 9, 12, 27].
We focus on theweighted firing rate [12, 27] which has connection with OTTT in this work.
Define weightedfiring rates and weighted average inputs
[Equation] in the discretecondition. 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) ) inthe discrete condition.
For feedforward networks, the closed-form mapping between successive layersis 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) , andgradients are calculated with such spike representation: ∂L∂ W l =
[Equation]
(cid:81)
[Equation]
[Equation]
W l .
{+For the recurrent condition, a [ t ] will converge to an equilibrium state following an implicit fixed-pointequation, e.g. a ∗ = σ (cid:16) 
+}
V th ( Wa ∗ + Fx ∗ + b ) (cid:17) for a single-layer network with input connections
F and contractive recurrent connections W , and gradients can be calculated based on implicitdifferentiation [12].
Let a = f θ ( a ) denote the fixed-point equation ( θ are parameters).
We have
[Equation]
, where J f θ | a [ T ] =
∂f θ ( a [ T ]) ∂ a [ T ] is the Jacobian of f θ at a [ T ] .
Online Training Through Time for SNNs
This section contains four sub-sections.
In Section 4.1, we introduce our proposed OTTT bydecoupling the temporal dependency of BPTT.
Then in Section 4.2, we further connect the gradientsof OTTT and spike representation-based methods, and prove that OTTT can provide a descentdirection for optimization, which is not guaranteed by BPTT with SG.
In Section 4.3, we discuss therelationship between OTTT and the three-factor Hebbian learning rule.
Implementation details arepresented 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 computationalgraph of previous time to backpropagate through time.
We will decouple such temporal dependencyto 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 inthe 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 considerthe case that we do not apply surrogate derivatives to ∂ s l +
[ i ] ∂ u l +
[ i ] in such temporal dependency.
Since thederivative 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 2 .
Then the dependency only includes ∂ u l +
[ i +1]
∂ u l +
[ i ] , which equals λ I .
Therefore, we have 1 :∂L ∂ W l =
T (cid:88) t = ∂L ∂ s l +1 [ t ] ∂ s l +1 [ t ] ∂ u l +1 [ t ]  (cid:88) τ ≤ t λ t − τ ∂ u l +1 [ τ ] ∂ W l   , ∇ W l L = T (cid:88) t = g u l +1 [ t ]  (cid:88) τ ≤ t λ t − τ s l [ τ ] 
 ⊤ , (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 trackpresynaptic 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 becalculated 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 feedbackconnections between neurons.
If there is feedback connection from layer l 2 to l 1 ( l 2 ≥ l 1 ), therewould 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) omitfeedback connections for simplicity).
We also consider not applying surrogate derivatives to ∂ s l
[ i ] ∂ u l 2 [ i ] inthe temporal dependency so that gradients are not calculated in this path.
Similar to the feedforwardcondition, we can derive that the gradients of the general weight W l i → l j from any layer l i to any layerl j can be calculated by ∇ W li → lj L [ t ] =
g u lj [ t ]ˆ a l i [ t ] ⊤ at each time step.
A theoretical explanationfor 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) T t =1 s N [ t ] , y (cid:17) , where y is the label, s N [ t ] is thespike at the last layer, and L can take cross-entropy loss.
This loss depends on all time steps and doesnot support online gradients.
We leverage the instantaneous loss and calculate the above g u l +1 [ t ] as:
[Equation]
where L is the upper bound of the typical loss 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 willapply 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 maintainingthe unfolded computational graph and only requires constant training memory costs agnostic to timesteps.
Note that instantaneous gradients of OTTT will be different from BPTT with the instantaneousloss for multi-layer or recurrent networks, as we do not consider future influence in the instantaneouscalculation: BPTT considers terms such as ∂L [ t ′ ]
∂ u N [ t ′ ] ∂ u N [ t
[Equation]
] ∂ u N [ t ] (cid:81) li = N − [Equation]
u i +
[ t ] ∂ u i [ t ] ( t ′ > t ) for u l [ t ] whilewe do not.
The equivalence of OTTT and BPTT only holds for the last layer, and we do not seekthe exact equivalence to BPTT with SG which is theoretically unclear, but will build the connectionwith spike representations and prove the descent guarantee.
Also, note that the tracked presynapticactivities are similar to the biologically plausible “eligibility traces” in the literature [45, 46], and wewill 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.
Feedforward Networks.
As introduced in Section 3.2, methods based on spike representations establish closed-form mappings between successive layers with weighted firing rate a [ t ] =
[Equation]
as a l +1 [ T ] ≈ σ (cid:16) 1 V th (cid:0) W l a l [ T ] + b l +1 (cid:1)(cid:17) , and calculate gradients by ∂L∂ W l = [Equation]
[Equation]
l
[Equation]
i
[Equation]
[Equation]
[Equation]
a i [ T ] ∂ a l +1 [ 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:( ∇ W l L sr ) sr = T (cid:88) t =   (cid:32) T λ T − t ∂L sr ∂ s N [ t ] l +1 (cid:89) i =
N −∂ a i +1 [ T ] ∂ a i [ T ] (cid:33) ⊤ ⊙ d l +1 [ T ] 
 ˆ a l
[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 theconnection 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 ) toapproximate d l +
i
[ T ] .
As introduced in Section 3.2, σ ( x ) = min(max(0 , x ) , 1) .
Thend l
+
[ T ] = sign (cid:0) | W l a l [ T ] + b l +1 − V
th | < V th (cid:1) , and it 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 approximatethe well-defined derivative of the mapping function between a [ T ] , rather than the pseudo derivativeof the non-differentiable Heaviside step function as in BPTT with SG.
The approximation is exactexcept the rare case that averagely a neuron generate spikes (i.e. the average input is positive) whilesometimes the membrane potential is less than u rest (or the reverse).
For simplicity in the followingtheoritical 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] [Equation]
(cid:80) Tt
=1 λ T − t L ( s N [ t ] , y )
so that1 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 isan 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) tt
=1 λ T [Equation]
t , y ) if L is a convexfunction.
Then let ˆ g u l +
[ 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]
L = [Equation] T (cid:80) Tt
=
ˆ g u l +
[ t ]ˆ a l [ t ] ⊤ and ( ∇ W l L sr ) sr = 1 T
[Equation]
] ⊤ .
Pleaserefer to Appendix A for details.
In the third step, we handle the remaining difference that ∇ W l L leverages instantaneous presynapticactivities ˆ a [ t ] during calculation while ( ∇ W l L sr ) sr uses the final ˆ a l [ T ] after time T .
Note that theweighted 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 Lcan 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 suchthat
[Equation]
=1 ˆ 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 problemformulated by spike representation.
Some random error can be viewed as randomness for stochasticoptimization.
Recurrent Networks.
For networks with feedback connections, we first consider the single-layercondition for simplicity (refer to Appendix A for general conditions).
We consider feedforwardconnections F from inputs to neurons and contractive recurrent connections W between neurons.
Asintroduced in Section 3.2, given inputs x [ t ] → x ∗ , a [ t ] of neurons will converge to an equilibriumstate a ∗ =
f θ ( a ∗ ) =
σ (cid:16) 1 V th ( Wa ∗ + Fx ∗ + b ) (cid:17) with bounded random error, and gradients are calculated as ( ∇ θ L sr ) sr
= (cid:0) ∂L sr ∂ θ (cid:1)
⊤ = (cid:16) ∂L sr ∂ a [ T ] (cid:0)
I − J f
θ | a [ T ] (cid:1) − 1 ∂f θ ( a [ T ]) ∂ θ (cid:17) ⊤ , where θ ∈ { W , F , b } .
We consider replacing the inverse Jacobian by an identity matrix:
[Equation]
Previous works have proved that this gradient can provide a descent direction for the optimizationproblem [35, 36], i.e. (cid:68) (cid:94) ( ∇ θ L sr ) sr , ( ∇ θ L sr )
sr (cid:69) > 0 .
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 1 λ T − t ∂L sr ∂ s [ t ] ⊤ ⊙
d [ T ] (cid:17) ˆ a [ T ] ⊤ .
Similarly, wecan 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 σ 2 max , where σ max and σ min are themaximal and minimal singular value of ∂f θ ∂ θ | a [ T ] , and the errors ϵ 1 [ t ] = a [ t ] − a [ T ] , ϵ
[ t ] = x
[ t ] −x [ T ] are small such that (cid:13)(cid:13)(cid:13)(cid:80)
[Equation]
max
[Equation]
max (cid:13)(cid:13)(cid:13)(cid:80)
[Equation] [Equation]
[Equation]
(cid:0)
[Equation]
−
[Equation] [Equation]
] (cid:1)
[Equation]
(cid:13)(cid:13)(cid:13) −(cid:13)(cid:13)(cid:13)(cid:80)
Tt =
1 λ t (1 − λ T − t ) 1 − λ T
ˆ 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 ,∇ W li → lj L [ t ] =
g u lj [ t ]ˆ a l i [ t ] ⊤ , and dive into connections between any two neurons i and j , we have:
[Equation]
[Equation]
where ˆ a i [ t ] is the tracked presynaptic activity, f ( u j [ t ]) is the surrogate derivative function which canrepresent the change rate of the postsynaptic activity as analyzed in Section 4.2, and δ j [ t ] =
g u j [ t ] isthe gradient for u j [ t ] which represents a global modulator.
This is a kind of three-factor Hebbianlearning 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 isshown biologically plausible with high-frequency bursts [48].
Note that the analysis in Section 4.2still 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 nexttime step, which we denote as OTTT O , or we can accumulate the gradients by T time steps and thenupdate parameters, which we denote as OTTT A .
For OTTT O , we assume that the online update issmall 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 threshold-dependentbatch normalization (tdBN) to achieve high performance with extremely low latency on large-scaledatasets [6, 10, 11, 13], which requires calculating the mean and variance statistics along the temporaldimension.
This technique intrinsically prevents online gradients and has to suffer from large memorycosts.
To overcome this shortcoming, we do not use tdBN, but borrow the idea from normalizationfree ResNets (NF-ResNets) [49] to replace batch normalization by scaled weight standardization(sWS) [50].
sWS standardizes weights by ˆ W i,j = γ · W i,j − µ W i, · σ W i, · √ N , and the scale γ is determined byanalyzing the signal propagation with different activation functions.
We apply sWS for VGG [51]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 [52], CIFAR100 [52], ImageNet [53],and CIFAR10-DVS [54] to demonstrate the superior performance of our proposed method onlarge-scale static and neuromorphic datasets.
We leverage the VGG network architecture (64C3128C3-AP2-256C3-256C3-AP2-512C3-512C3-AP2-512C3-512C3-GAP-FC) for experiments on
CIFAR-10, CIFAR-100 and CIFAR10-DVS, and the NF-ResNet-34 [49] network architecture forexperiments on ImageNet.
For all our SNN models, we set V th = 1 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 thetemporal 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 networkon CIFAR-10 with batch size 128 under different time steps and calculating the memory costs onthe 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 oftime steps, e.g. 6, OTTT can reduce the memory costs by 2 ∼ 3 × .
This advantage may also allowtraining acceleration of SNNs by larger the
[Table] [Table]
5.2 Comparison of Performance
We conduct experiments on both large-scale static and neuromorphic datasets.
We first verify theeffectiveness of the surrogate derivative ∂ s l +
[ t ] ∂ u l +
[ t ] = sign (cid:0) |
u l +1 [ t ] − V th | < V th (cid:1) in Section 4.2.
Forthe VGG (sWS) network on CIFAR-10, OTTT A and OTTT O achieves 91.46% and 92.47% testaccuracy respectively.
We empirically observe that applying the sigmoid-like surrogate derivativeachieves a higher generalization performance, e.g. OTTT A and OTTT O achieves 93.58% and93.73% test accuracy respectively under the same random seed, and a possible reason may bethat this introduces some noise for the approximation to regularize the training and improve thegeneralization.
Therefore, in the following performance evaluation, we take sigmoid-like surrogatederivative for OTTT.
As shown in Table 1, both OTTT A and OTTT O achieve satisfactory performanceon all datasets, and compared with BPTT under the same training settings, OTTT achieves higherperformance.
The proposed OTTT also achieves promising performance on all datasets comparedwith other representative conversion and direct training methods.
5.3 Effectiveness for Recurrence
As introduced in Sections 4.1 and 4.2, the proposed OTTT is also valid for networks with feedbackconnections.
Previous work has shown that adding feedback connections can improve the performanceof SNNs, especially on the CIFAR-100 datasets [12, 56].
Therefore, we conduct experiments on
CIFAR-100 with the VGG-F network architecture which simply adds a feedback connection fromthe last feature layer to the first feature layer following [12], and this weight is zero-intialized.
Asshown in Table 2, the training of VGG-F is valid and VGG-F achieves a higher performance than
VGG due to the introduction of feedback connections.
The architectures with feedback connectionscan be further improved with neural architecture search [56].
[Table] [Table]
5.4 Effectiveness for Training with Batch Size 1
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 theeffectiveness for training with batch size 1.
The VGG network on CIFAR-10 is studied, and batchsize 1 is compared with the default batch size 128 under the same random seed.
Models are onlytrained for 20 epochs due to the relatively long training time with batch size 1.
As shown in Table 3,training with one sample per iteration is still valid, indicating the potential to conduct full onlinetraining with the proposed OTTT.
5.5 Influence of Inference Time Steps
We study the influence of inference time steps on ImageNet as shown in Fig.
It illustrates that themodel 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. 4.
Itdemonstrates that models trained by OTTT have higher firing rates in first layers while lower firingrates in later layers compared with BPTT.
Overall the firing rate is around 0.19 and with 6 time stepseach neuron averagely generate 1.1 spikes, indicating the low energy consumption.
Considering thateach neuron has more synaptic operations in later layers than first layers (because the channel size isincreasing with layers), the synaptic operations of models trained by OTTT and BPTT are about thesame ( 1 . 98 × 10 8 vs 1 .
93 × 10 8 ).
More results please refer to Appendix D.
[Figure] [Figure]
Conclusion
In this work, we propose a new training method, online training through time (OTTT), for spikingneural networks.
We first derive OTTT from BPTT with SG by decoupling the temporal dependencywith the tracked pre-synaptic activities, which only requires constant training memory agnostic totime steps and avoids the large training memory costs of BPTT.
Then we theoretically analyze andconnect the gradients of OTTT and gradients of methods based on spike representations, and provethe descent guarantee of OTTT for the optimization problem under both feedforward and recurrentnetwork conditions.
Additionally, we show that OTTT is in the form of three-factor Hebbian learningrule, which is the first to connect BPTT with SG, spike representation-based methods, and biologicallearning rules.
Extensive experiments demonstrate the superior performance of our methods onlarge-scale static and neuromorphic datasets in a small number of time steps.