原文链接:https://blog.csdn.net/qq_43622216/article/details/123591425
Going Deeper With Directly-Trained Larger Spiking Neural Networks
说明
准备将自己读的一些和SNN(脉冲神经网络)相关的一些论文记录下来,方便自己以后回看也希望能够帮到有需要的人。
删除了文中一些自认为不重要的内容(很少),其他部分尽量使用专业用语进行翻译,如果有什么出错或不恰当的地方希望各位批评指出。
相关信息
论文地址: Going Deeper With Directly-Trained Larger Spiking Neural Networks
论文于2020年12月18日上传至arXiv,发表于AAAI 2021,由清华大学和加州大学圣巴巴拉分校相关研究人员完成。
通讯作者:李国齐
李国齐老师团队在SNN领域已经发了不少具有价值的文章,为方便查看它们的关系,列表如下:
文献 | 时间 | N-MNIST | MNIST-DVS | ImageNet | CIAFR10 | CIFAR10-DVS | DVS128-Gesture | SHD Dataset |
---|---|---|---|---|---|---|---|---|
文献1 | 2018年 | 98.78% | ||||||
文献2 | 2019年 | 99.53% | 90.53% | 60.5% | ||||
文献3(本文) | 2020年 | 67.05% | 93.16% | 67.8% | 96.87% | |||
文献4 | 2021年 | 99.13% | 70.4% | 97.56% | ||||
文献5 | 2021年 | 72% | 98.61% | 91.08% | ||||
文献6 | 2022年 | 99.63% | 73.8% | 98.96% |
文献1:Wu Y, Deng L, Li G, et al. Spatio-temporal backpropagation for training high-performance spiking neural networks[J]. Frontiers in neuroscience, 2018, 12: 331.
文献2:Wu Y, Deng L, Li G, et al. Direct training for spiking neural networks: Faster, larger, better[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2019, 33(01): 1311-1318.
文献3:Zheng H, Wu Y, Deng L, et al. Going deeper with directly-trained larger spiking neural networks[J]. arXiv preprint arXiv:2011.05280, 2020.
文献4:Wu Z, Zhang H, Lin Y, et al. Liaf-net: Leaky integrate and analog fire network for lightweight and efficient spatiotemporal information processing[J]. IEEE Transactions on Neural Networks and Learning Systems, 2021.
文献5:Yao M, Gao H, Zhao G, et al. Temporal-wise Attention Spiking Neural Networks for Event Streams Classification[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2021: 10221-10230.
文献6:Ultra-low Latency Spiking Neural Networks with Spatio-Temporal Compression and Synaptic Convolutional Block (非李国齐组)
主要贡献
论文提出了基于spatio-temporal backpropagation (时空反向传播)的threshold-dependent batch normalization (tdBN) 方法,称为 “STBP-tdBN”,并使用tdBN构建了一个深度脉冲残差网络,扩展了直接训练的SNN的深度(从小于10层到50层),使用很小的timestep(时间步长)在CIFAR-10数据集上达到了93.15%的准确率,在DVS-CIFAR10数据集上达到了67.8%的准确率(SOTA)、ImageNet数据集上达到了67.05%的准确率。
启发
基于本文提出的tdBN算法以及ResNet-SNN,可以用其构建深层次的SNN网络并取得更好的性能。
Abstract
脉冲神经网络(SNN)在时空信息的bioplausible coding(生物可信编码)和event-driven(事件驱动)信号处理方面具有广阔的应用前景,非常适合应用于energy-efficient(低能耗的)神经形态硬件中。然而,其独特的工作模式使得其比传统网络更难训练。目前,主要有两种方式训练高性能的深层SNN。一是将预训练的ANN模型转换为对应的SNN版本,这种方式通常需要较长的coding window(编码窗口)进行收敛,也不能在训练时利用spatio-temporal features(时空特征)来解决时间任务(temporal tasks)。另一种方法是在时空域(spatio-temporal domain)中直接训练SNN,但由于激发函数(firing function)的二元脉冲活动(binary spike activity)和梯度消失或爆炸的问题,目前的方法还局限于浅层架构,难以利用大规模数据集(如ImageNet)。为此,我们提出了一种阈值依赖批量归一化(threshold-dependent batch normalization ,tdBN)方法,该方法基于新兴的时空反向传播(spatio-temporal backpropagation),称为“STBP-tdBN”,可以直接训练深层SNN,并在神经形态硬件上有效地实现其推理。利用提出的方法和精心设计的shortcut connection,我们将直接训练的SNN从一个浅层结构(<10层)大大地扩展到一个很深的结构(50层)。此外,我们基于“块动态等距(Block Dynamical Isometry)”理论从理论上分析了我们方法的有效性。最后,我们报告了更好的精度结果,包括在CIFAR-10上的93.15%,在DVS-CIFAR10上的67.8%,在ImageNet上的67.05%,而且时间步长(timestep)非常小。据我们所知,这是第一次在ImageNet上探索直接训练的高性能的深度SNN。我们相信这项工作将为充分发挥SNN的优势、吸引更多的研究者在这一领域做出贡献铺平道路。
Introduction
目前主要有两种方式来训练高性能的SNN。
- 第一种方式是将预训练的ANN转换为SNN,这通常需要数百个时间步长。因此,尽管这些SNN与具有类似结构的ANN具有相当的精度,但大量的时间步长会导致严重的信号延迟(signal latency)并增加计算量。
- 另一种方法是基于梯度下降方法直接训练SNN,该方法不依赖于预先训练的ANN并且能减少时间步长。
目前基于STBP算法训练的SNN仅限于浅层架构,无法在大规模数据集上达到令人满意的性能。因此,在上述算法下我们明确了两个需要解决的问题,以便直接训练更深层的SNN。
- 第一个问题是梯度消失或爆炸。由于其独特的时空信息处理机制和不可微的脉冲信号,当直接训练SNN时,梯度传播表现得非常不稳定,在大多数情况下趋于消失,这影响了深层SNN的效果。到目前为止,直接训练的SNN还没有一种有效的方法可以很好地处理这个问题。前人的工作都没能直接训练深层SNN并且他们的模型都小于10层,严重影响了他们方法的性能。
- 另一个问题是我们需要平衡每个神经元的阈值和输入,以在 SNN 中获得适当的放电率(firing rate)。当输入与阈值相比太小时,神经元发出的脉冲很少,神经元膜电位保持不变,因此神经元处理的信息无法充分表达。当输入太大时,神经元会一直放电(脉冲)并对输入的变化不敏感。对于直接训练的SNN,由于双脉冲(binary spikes)逐层传播,在训练过程中突触前输入(presynaptic inputs)的分布会发生偏移,使得输入的大小不合适。人们提出了许多方法来处理这个问题,如threshold regularization(阈值正则化)和NeuNorm。
Normalization(正则化)似乎是解决这两个问题的合适方法,能够稳定网络和梯度传播,并且能将突触前输入的分布正则化到相同的期望和方差,通过减少内部协变量移位(internal covariate shift)来平衡阈值和输入。然而,现有的正则化方法并不适合于SNN的训练。对于额外的时间维度和特殊的激活机制(activation mechanism),直接训练的SNN需要专门设计的正则化方法。
本文提出了一种直接训练深层SNN的新算法,主要工作如下:
- 我们提出阈值依赖批量归一化来解决梯度消失或爆炸问题,并调整放电率。在此基础上,利用残差网络结构,修改适合SNN的shortcut connection。
- 在此基础上,我们深入研究了直接训练的SNN(从10层扩展到50层),并在大规模非脉冲数据集(CIFAR-10, ImageNet)和神经形态数据集(DVSGesture, DVS-CIFAR10)上进行测试。
- 在CIFAR-10和ImageNet上,我们全面验证了不同的SNN架构(ResNet-18, 34, 50),并给出了使用更少的时间步长(不超过6个时间步长)与类似SNN的比较结果。据我们所知,这是第一次直接训练的带有全脉冲(full spikes)的SNN在ImageNet上显示了相当高的准确性。在神经形态数据集上,我们的模型在DVS-Gesture和DVS-CIFAR10上都取得了最好的性能(state-of-the-art),显示了SNN在处理时空信息方面的优势。
Related Work
SNN学习算法:
在过去的几年里,很多学习算法都在探索如何训练一个深层的SNN,如(1)一些将预训练的ANN转化为SNN;(2)基于梯度下降的算法。
- 第一种方法叫做"ANN-SNN conversion methods",它将ReLU函数的实值输出转换为SNN中的二进制脉冲,被认为是最流行的训练高性能深层SNN的方法。这种方法在大规模数据集上达到了有竞争力的性能,相对于ANN也没有严重的退化。然而,这种方法忽略了脉冲神经元丰富的时间动态行为(temporal dynamic behaviors)并且通常需要数百或数千个时间步长才能接近预训练的ANN的精度。
- 基于梯度下降的算法对SNN进行误差反向传播训练。利用梯度下降优化学习算法,一些SNN模型在CIFAR-10和其他神经形态数据集上达到了较高的性能。直接训练的SNN在处理空间和时间信息方面显示出巨大的潜力,并能在非常少的时间步长内达到较高的的准确性。但是,由于梯度消失和内部协变量移位,还不能直接训练出非常深的SNN,这正是我们想要克服的。
深度神经网络(DNN)中的梯度消失或爆炸:
当动态等距时(即其输入输出雅可比矩阵的每个奇异值都保持在1附近),DNN可以避免梯度消失或梯度爆炸。Chen1等人提出了“块动态等距”,作为所有复杂串并联DNN的通用统计工具,通过研究神经网络中每个块的一阶矩和二阶矩,并分析它们对梯度分布的影响,并对神经网络中权值初始化( weight initialization)、批处理归一化(batch normalization )和shortcut connection 的作用进行了理论解释,为我们算法的开发提供了理论依据。
正则化:
对于人工神经网络(ANN)来说,正则化(如batch normalization、group normalization、 layer normalization )已成为常用的方法。batch normalization(BN)通过减少内部协变量移位(internal covariate shift)来加速深度网络训练,从而提高学习速度并正则化模型。虽然它会导致很高的学习延迟(learning latency )和增加计算量,但BN使得能够训练更深的网络并避免梯度消失或爆炸。对于SNN,研究人员提出了其他的正则化技术,如data-based normalization、Spike-Norm和NeuNorm。这些正则化方法的目的是平衡输入和阈值,避免严重的信息丢失,但它们仍然忽略了梯度消失问题,对直接训练的深层SNN无效。 我们注意到BN在ANN中的效果和SNN中输入分布的重要性,因此我们修改BN以满足SNN模型的训练和推理。
Materials and Methods
迭代的LIF模型(Iterative LIF model):
迭代LIF模型最早由Wu2等人在2019年提出,利用欧拉法(Euler method)求解Leaky integrate-and-fire (LIF)模型的一阶微分方程,并将其转换为迭代表达式
u
t
=
τ
d
e
c
a
y
u
t
−
1
+
I
t
,
(1)
u^t = \tau_{decay}u^{t-1}+I^t \tag{1},
ut=τdecayut−1+It,(1)
其中
τ
d
e
c
a
y
\tau_{decay}
τdecay是描述膜电位衰减速度的常数,
u
t
u^t
ut是膜电位,
I
t
I^t
It是突触前输入,
V
t
h
V_{th}
Vth表示给定的阈值。当
u
t
>
V
t
h
u^t>V_{th}
ut>Vth时,神经元发放一个脉冲并将
u
t
u^t
ut置为0。突触前输入是由前一层的其他神经元发放的的脉冲信号之和,因此
I
t
I^t
It可以表示为
x
t
=
∑
j
w
j
o
t
(
j
)
x^t=\sum_{j}w_jo^t(j)
xt=∑jwjot(j),其中
w
j
w_j
wj表示权重,
o
t
(
j
)
o^t(j)
ot(j)表示在
t
t
t时刻来自其他神经元的二进制脉冲输出。考虑到空间结构,并设置
u
r
e
s
e
t
=
0
u_{reset}=0
ureset=0,则整个空间和时间域的迭代LIF模型可以表示为
u
t
,
n
+
1
=
τ
d
e
c
a
y
u
t
−
1
,
n
+
1
(
1
−
o
t
−
1
,
n
+
1
)
+
x
t
,
n
,
(2)
u^{t,n+1} =\tau_{decay}u^{t-1,n+1}(1-o^{t-1,n+1})+x^{t,n}, \tag{2}
ut,n+1=τdecayut−1,n+1(1−ot−1,n+1)+xt,n,(2)
o
t
,
n
+
1
=
{
1
i
f
u
t
,
n
+
1
>
V
t
h
,
0
o
t
h
e
r
w
i
s
e
(3)
o^{t,n+1} =\begin{cases}1&{\rm if}\,u^{t,n+1}>V_{th},\\ 0&{\rm otherwise} \end{cases} \tag{3}
ot,n+1={10ifut,n+1>Vth,otherwise(3)
其中
u
t
,
n
u^{t,n}
ut,n表示在第
n
n
n层神经元在
t
t
t时刻的膜电位,
o
t
,
n
o^{t,n}
ot,n是二进制脉冲,
τ
d
e
c
a
y
\tau_{decay}
τdecay是膜电势衰减常数(potential decay constant)。
迭代LIF模型可以在空间和时间维度上实现前向和后向传播,这使得它对一般的机器学习编程框架很友好。
阈值依赖批量正则化(Threshold-dependent batch normalization):
批归一化(batch normalization, BN)作为DNN的正则化组成部分,是目前常用的神经网络方法,它可以实现稳定的收敛和更深层的神经网络。但是,由于SNN额外的时间维度和特殊的激活机制,直接训练的SNN需要专门设计的归一化方法。这促使我们提出阈值依赖的批归一化(threshold-dependent batch normalization,tdBN)。
我们考虑一种脉冲卷积神经网络(SCNN)。设
o
t
o^t
ot表示时间步长(timestep)为
t
t
t时一层内所有神经元的脉冲输出,以及卷积核
W
W
W和偏差(bias)
B
B
B,我们有如下表达式
x
t
=
W
⊛
o
t
+
B
(4)
x^t = W \circledast o^t + B \tag{4}
xt=W⊛ot+B(4)
其中
x
t
∈
R
N
×
C
×
H
×
W
x^t \in R^{N \times C \times H \times W}
xt∈RN×C×H×W表示以
N
N
N为批处理轴(batch axis)且时间步长为
t
t
t时的突触前输入,
C
C
C是通道轴(channel axis),
(
H
,
W
)
(H,W)
(H,W)是空间轴。
在我们的tdBN里,高维的突触前输入将会沿着通道维度被归一化(如下图1所示)。令
x
k
t
x^t_k
xkt表示
x
t
x^t
xt的第
k
k
k个通道的特征图(feature map),
x
k
=
(
x
k
1
,
x
k
2
,
.
.
.
,
x
k
T
)
x_k=(x^1_k,x^2_k,...,x^T_k)
xk=(xk1,xk2,...,xkT)会按照如下方式进行归一化
x
^
k
=
α
V
t
h
(
x
k
−
E
[
x
k
]
)
V
a
r
[
x
k
]
+
ϵ
,
(5)
\hat{x}_k=\frac {\alpha V_{th}(x_k-E[x_k])}{\sqrt{Var[x_k]+\epsilon}}, \tag{5}
x^k=Var[xk]+ϵαVth(xk−E[xk]),(5)
y
k
=
λ
k
x
^
k
+
β
k
,
(6)
y_k = \lambda_k \hat{x}_k + \beta_k, \tag{6}
yk=λkx^k+βk,(6)
其中
V
t
h
V_{th}
Vth表示阈值,
α
\alpha
α是一个依赖于网络结构的超参数,
ϵ
\epsilon
ϵ是一个很小的常数,
λ
k
\lambda_k
λk和
β
k
\beta_k
βk是两个可以训练的参数,
E
[
x
k
]
、
V
a
r
[
x
k
]
E[x_k]、Var[x_k]
E[xk]、Var[xk]是在一个MiniBatch上统计估计的
x
k
x_k
xk的均值和方差。图1展示了
E
[
x
k
]
、
V
a
r
[
x
k
]
E[x_k]、Var[x_k]
E[xk]、Var[xk]的计算过程,可定义为
E
[
x
k
]
=
m
e
a
n
(
x
k
)
,
(7)
E[x_k] = mean(x_k), \tag{7}
E[xk]=mean(xk),(7)
V
a
r
[
x
k
]
=
m
e
a
n
(
(
x
k
−
E
[
x
k
]
)
2
)
.
(8)
Var[x_k] = mean((x_k - E[x_k])^2). \tag{8}
Var[xk]=mean((xk−E[xk])2).(8)
因此,在训练期间,
y
k
∈
R
T
×
N
×
H
×
W
y_k \in R^{T \times N \times H \times W}
yk∈RT×N×H×W正是下一层第k个通道的神经元在T时间步长内接收到的归一化突触前输入。
在推理过程中,我们遵循标准Batch Normalization的模式(schema),分别估计出代表
E
[
x
k
]
E[x_k]
E[xk]和
V
a
r
[
x
k
]
Var[x_k]
Var[xk]在整个数据集上的期望
μ
i
n
f
\mu_{inf}
μinf和
σ
i
n
f
2
\sigma^2_{inf}
σinf2,这些期望和方差可以在训练过程中通过移动平均求解(moving average solution)得到。
除此之外,在对具有tdBN的SNN进行推理时,进行批量尺度融合(batchnorm-scale-fusion)是必要的。它移除了推理期间的Batch Normalization操作,从而保持网络是全脉冲(full-spiking)的,并使其能够在神经形态平台上实现。设
W
c
,
k
W_{c,k}
Wc,k和
B
c
,
k
B_{c,k}
Bc,k表示一层里的第
c
c
c个feature map与下一层的第
k
k
k个feature map之间的卷积核和偏差,模式(schema)可以由下式定义
W
c
,
k
′
=
λ
k
α
V
t
h
W
c
,
k
σ
i
n
f
,
k
2
+
ϵ
,
(9)
W^{'}_{c,k} = \lambda_k \frac {\alpha V_{th}W_{c,k}}{\sqrt{\sigma^2_{inf,k}+\epsilon} }, \tag{9}
Wc,k′=λkσinf,k2+ϵαVthWc,k,(9)
B
c
,
k
′
=
λ
k
α
V
t
h
(
B
c
,
k
−
μ
i
n
f
,
k
)
σ
i
n
f
,
k
2
+
ϵ
+
β
k
,
(10)
B^{'}_{c,k} = \lambda_k \frac{\alpha V_{th}(B_{c,k}-\mu_{inf,k})}{\sqrt{\sigma^2_{inf,k}+\epsilon}}+ \beta_k, \tag{10}
Bc,k′=λkσinf,k2+ϵαVth(Bc,k−μinf,k)+βk,(10)
其中的
W
c
,
k
′
W^{'}_{c,k}
Wc,k′和
B
c
,
k
′
B^{'}_{c,k}
Bc,k′表示经过batchnorm-scale-fusion后的权重和偏差,因此在推理过程中,脉冲通过变换后的权重
W
c
,
k
′
W^{'}_{c,k}
Wc,k′和偏差
B
c
,
k
′
B^{'}_{c,k}
Bc,k′逐层传播,就不需要进行batch norm操作。因此,我们的tdBN只影响训练期间的计算量而不影响已经训练好的SNN的运行过程。
简而言之,我们的tdBN与标准BN有两个主要区别。首先,与ANN不同的是,SNN不仅逐层(layer by layer)地传播信息,而且从上一时刻刻传播到下一时刻。因此,tdBN应该在时间和空间两个维度上对特征输入进行归一化。其次,根据阈值
V
t
h
V_{th}
Vth归一化方差。在tdBN中,激活前的值(pre-activations )被归一化为
N
(
0
,
(
α
V
t
h
)
2
)
N(0,(\alpha V_{th})^2)
N(0,(αVth)2)而不是
N
(
0
,
1
)
N(0,1)
N(0,1)。我们用1和0初始化可训练参数
λ
\lambda
λ和
β
\beta
β,在串联神经网络里(serial neural network)超参数
α
\alpha
α都设为1。对于有
n
n
n个分支的局部并行网络结构,
α
=
1
n
\alpha= \frac{1}{\sqrt{n}}
α=n1。这使得在早期的训练过程中
V
t
h
V_{th}
Vth的激活前的值均值为0,标准差为
V
t
h
V_{th}
Vth。tdBN的代码见文末的补充材料A。
整体训练算法(Overall training algorithm):
在本节中,我们介绍STBP-tdBN的整体训练算法,用我们的tdBN从头开始训练深层SNN。
在误差反向传播中,我们将最后一层作为解码层(decoding layer),最终的输出Q表示为:
Q
=
1
T
∑
t
=
1
T
M
o
n
,
t
,
(11)
Q = \frac {1}{T} \sum_{t=1}^{T}Mo^{n,t}, \tag{11}
Q=T1t=1∑TMon,t,(11)
其中
o
n
,
t
o^{n,t}
on,t是最后一层输出层发放的脉冲,
M
M
M是解码层矩阵,
T
T
T是时间步长(timestep)。
然后我们让输出通过一层softmax层。损失函数(loss function)使用交叉熵。考虑有输出
Q
=
(
q
1
,
q
2
,
.
.
.
,
q
n
)
Q=(q_1,q_2,...,q_n)
Q=(q1,q2,...,qn)和标签向量(label vector)
Y
=
(
y
1
,
y
2
,
.
.
.
,
y
n
)
Y=(y_1,y_2,...,y_n)
Y=(y1,y2,...,yn),损失函数
L
L
L定义为:
p
i
=
e
q
i
∑
j
=
1
n
e
q
i
,
(12)
p_i = \frac{e^{q_i}}{\sum_{j=1}^ne^{q_i}}, \tag{12}
pi=∑j=1neqieqi,(12)
L
=
−
∑
i
=
1
n
y
i
l
o
g
(
P
i
)
.
(13)
L = -\sum_{i=1}^ny_ilog(P_i). \tag{13}
L=−i=1∑nyilog(Pi).(13)
利用迭代LIF模型,STBP-tdBN方法在空间和时间域上反向传播损失函数
L
L
L的梯度。利用链式法则,可以利用下式来计算
∂
L
∂
o
i
t
,
n
\frac{\partial L}{\partial o_i^{t,n}}
∂oit,n∂L和
∂
L
∂
u
i
t
,
n
\frac{\partial L}{\partial u_i^{t,n}}
∂uit,n∂L:
∂
L
∂
o
i
t
,
n
=
∑
j
=
1
l
(
n
+
1
)
∂
L
∂
u
j
t
,
n
+
1
∂
u
j
t
,
n
+
1
∂
o
i
t
,
n
+
∂
L
∂
u
i
t
+
1
,
n
∂
u
i
t
+
1
,
n
∂
o
i
t
,
n
,
(14)
\frac{\partial L}{\partial o_i^{t,n}} = \sum_{j=1}^{l(n+1)}\frac{\partial L}{\partial u_j^{t,n+1}}\frac{\partial u_j^{t,n+1}}{\partial o_i^{t,n}} + \frac{\partial L}{\partial u_i^{t+1,n}}\frac{\partial u_i^{t+1,n}}{\partial o_i^{t,n}}, \tag{14}
∂oit,n∂L=j=1∑l(n+1)∂ujt,n+1∂L∂oit,n∂ujt,n+1+∂uit+1,n∂L∂oit,n∂uit+1,n,(14)
∂
L
∂
u
i
t
,
n
=
∂
L
∂
o
i
t
,
n
∂
o
i
t
,
n
∂
u
i
t
,
n
+
∂
L
∂
u
i
t
+
1
,
n
∂
u
i
t
+
1
,
n
∂
u
i
t
,
n
,
(15)
\frac{\partial L}{\partial u_i^{t,n}} = \frac{\partial L}{\partial o_i^{t,n}}\frac{\partial o_i^{t,n}}{\partial u_i^{t,n}} + \frac{\partial L}{\partial u_i^{t+1,n}}\frac{\partial u_i^{t+1,n}}{\partial u_i^{t,n}}, \tag{15}
∂uit,n∂L=∂oit,n∂L∂uit,n∂oit,n+∂uit+1,n∂L∂uit,n∂uit+1,n,(15)
其中
o
t
,
n
o^{t,n}
ot,n和
u
t
,
n
u^{t,n}
ut,n表示第
n
n
n层神经元在
t
t
t时刻的脉冲和膜电位,由于脉冲活动不可微分,
∂
o
t
∂
u
t
\frac{\partial o^t}{\partial u^t}
∂ut∂ot实际上并不存在,为了解决这个问题,Wu3等人在2018年提出了导数曲线(derivative curve)来近似脉冲活动的导数,在本文中,我们使用矩形函数(rectangular function),它在梯度下降中是很有效的,可以由下式求得:
∂
o
t
∂
u
t
=
1
a
s
i
g
n
(
∣
u
t
−
V
t
h
∣
<
a
2
)
.
(16)
\frac {\partial o^t}{\partial u^t} = \frac {1}{a}sign(|u^t-V_{th}| < \frac {a}{2}). \tag{16}
∂ut∂ot=a1sign(∣ut−Vth∣<2a).(16)
整体训练算法的代码同样见文末的补充材料A。
Theoretical Analysis
在本节中,我们将分析tdBN对STBP-tdBN训练的SNN的影响。利用神经网络中梯度规范理论(gradient norm theory)的理论工具,我们发现我们的tdBN可以缓解训练过程中梯度消失或爆炸的问题。我们还将解释归一化过程中加入的尺度因子 α \alpha α和 V t h V_{th} Vth的作用。
梯度规范理论(Gradient norm theory):
梯度范数理论是近年来发展较好的理论,它旨在克服各种神经网络结构中的梯度消失或爆炸问题。在本文中,我们采用Chen4等人在2020年提出的“Block Dynamical Isometry”来分析tdBN在直接训练的SNN中的效果,它把网络看作为一系列的块(block):
f
(
x
)
=
f
i
,
θ
i
∘
f
i
−
1
,
θ
i
−
1
∘
.
.
.
∘
f
1
,
θ
1
(
x
)
,
(17)
f(x) = f_{i,\theta_i} \circ f_{i-1,\theta_{i-1}} \circ ... \circ f_{1,\theta_1}(x), \tag{17}
f(x)=fi,θi∘fi−1,θi−1∘...∘f1,θ1(x),(17)
其中函数
f
j
,
θ
j
f_{j,\theta_j}
fj,θj表示第
j
j
j个块,并将其输入输出雅可比矩阵定义为
∂
f
j
∂
f
j
−
1
=
J
j
\frac {\partial f_j}{\partial f_{j-1}} = J_j
∂fj−1∂fj=Jj,使用
ϕ
(
J
)
\phi(J)
ϕ(J)表示
t
r
(
J
)
tr(J)
tr(J)的期望,
φ
(
J
)
\varphi(J)
φ(J)表示
ϕ
(
J
2
)
−
ϕ
2
(
J
)
\phi(J^2) - \phi ^2(J)
ϕ(J2)−ϕ2(J),随后证明下面几个引理(lemma):
引理1. 考虑一个神经网络,它可以如式(17)那样用一系列的块表示,第 j j j个块的雅可比矩阵表示为 J j J_j Jj。如果 ∀ j , ϕ ( J j J j T ) ≈ 1 \forall j,\phi(J_jJ_j^T)\approx 1 ∀j,ϕ(JjJjT)≈1且 φ ( j j J j T ) ≈ 0 \varphi(j_jJ_j^T)\approx 0 φ(jjJjT)≈0,网络就达到了“块动态等距(Block Dynamical Isometry)”,并能避免梯度消失或梯度爆炸。
引理2. 考虑一个神经网络块,包含有0均值的数据归一化(data normalization with 0-mean)、线性转换(linear transform )和整流激活rectifier activations(General Linear Transform)。设输入输出向量的二阶矩分别为 α i n \alpha_{in} αin和 α o u t \alpha_{out} αout,有 ϕ ( J J T ) = α o u t α i n \phi(JJ^T)=\frac{\alpha_{out}}{\alpha_{in}} ϕ(JJT)=αinαout。
在梯度规范理论框架的基础上,将梯度规范与脉冲神经元的特性相结合,进一步分析tdBN算法对SNN的有效性。
LIF模型有两个独特的超参数:
τ
d
e
c
a
y
\tau_{decay}
τdecay和
V
t
h
V_{th}
Vth,其中
τ
d
e
c
a
y
\tau_{decay}
τdecay影响梯度在时间域中的传播,
V
t
h
V_{th}
Vth影响空间维度。在SNN的实验中,
τ
d
e
c
a
y
\tau_{decay}
τdecay通常设置为比较小的值(如0.25)。为了分析梯度变换,我们将模型简化,设置
τ
d
e
c
a
y
\tau_{decay}
τdecay为0,得到如下命题:
定理1: 考虑一个时间步长为T的SNN,在T时刻的第
j
j
j个雅可比矩阵记为
J
j
t
J^t_j
Jjt。当
τ
d
e
c
a
y
=
0
\tau_{decay}=0
τdecay=0时,将两个tdBN层之间每个块(block)的输入向量和输出向量的二阶矩固定为
V
t
h
2
V_{th}^2
Vth2,有
ϕ
(
J
j
t
(
J
j
t
)
T
)
≈
1
\phi(J_j^t(J_j^t)^T)\approx1
ϕ(Jjt(Jjt)T)≈1并且可以在SNN的训练过程中避免梯度消失或梯度爆炸。
证明: 定理1的证明基于引理1和引理2,详细的证明过程在补充材料B中。
膜电势衰减机制的影响(Influence of membrane decay mechanism):
下面分析
τ
d
e
c
a
y
\tau_{decay}
τdecay对梯度传播的影响。从等式(2)到等式(15),可以得到:
∂
L
∂
u
i
t
,
n
=
∂
L
∂
o
i
t
,
n
∂
o
i
t
,
n
∂
u
i
t
,
n
+
∂
L
∂
u
i
t
+
1
,
n
τ
d
e
c
a
y
(
1
−
o
i
t
,
n
)
.
(18)
\frac{\partial L}{\partial u_i^{t,n}} = \frac{\partial L}{\partial o_i^{t,n}}\frac{\partial o_i^{t,n}}{\partial u_i^{t,n}} + \frac{\partial L}{\partial u_i^{t+1,n}}\tau_{decay}(1-o_i^{t,n}). \tag{18}
∂uit,n∂L=∂oit,n∂L∂uit,n∂oit,n+∂uit+1,n∂Lτdecay(1−oit,n).(18)
如果一个神经元发放了脉冲,
(
1
−
o
i
t
,
n
)
=
0
(1-o_i^{t,n})=0
(1−oit,n)=0,此时梯度与
τ
d
e
c
a
y
\tau_{decay}
τdecay无关。另一方面,由于
τ
d
e
c
a
y
\tau_{decay}
τdecay是一个很小的常数,神经元在
t
+
1
t+1
t+1时刻的梯度对
t
t
t时刻的影响不大。
为了验证定理1和分析膜电势衰减机制的影响,我们在CIFAR-10数据集上用20层的plain spiking network评估tdBN。在图2中,我们展示了在训练过程中的前1/6个epoch中除了第一个编码层和最后一个输出层外的每一层梯度范数(gradient norm)的均值。结果表明,当
τ
d
e
c
a
y
\tau_{decay}
τdecay= 0时,梯度范数曲线表现较稳定,证实了我们的理论。但要注意的是,
τ
d
e
c
a
y
\tau_{decay}
τdecay不能设置为0,因为这样会阻碍信息沿时间维度传播从而造成严重的退化问题。因此我们在
τ
d
e
c
a
y
≠
0
\tau_{decay}\neq0
τdecay=0的条件下评估我们的方法。当
τ
d
e
c
a
y
\tau_{decay}
τdecay设置为0.25和0.5时,梯度范数随着网络的加深增长非常缓慢,也就不会影响训练过程。研究结果有力地支持了我们的结论,即在深层SNN中可以避免梯度消失或爆炸的问题。
尺度因子(Scaling factors):
众所周知,SNN模型获得有竞争力的性能的关键是设置合适的阈值以保持发放率(firing rate)和减少信息丢失。为了实现这一点,我们在tdBN的归一化实现中引入两个尺度因子用来平衡预激活(pre-activation)和阈值。在早期的训练过程中,有
α
\alpha
α和
V
t
h
V_{th}
Vth两个尺度因子,通过将可训练参数
λ
\lambda
λ和
β
\beta
β初始化为1和0来将预激活归一化为
N
(
0
,
V
t
h
2
)
N(0,V_{th}^2)
N(0,Vth2)。
首先,我们提出定理2来解释预激活和神经元膜电势之间的关系,以助于理解为什么我们的方法是有效的。
定理2 采用迭代LIF模型,假设预激活
x
t
∼
N
(
0
,
σ
i
n
2
)
x^t \sim N(0,\sigma_{in}^2)
xt∼N(0,σin2),则膜电势
u
t
∼
N
(
0
,
σ
o
u
t
2
)
u^t \sim N(0,\sigma_{out}^2)
ut∼N(0,σout2)且
σ
o
u
t
2
∝
σ
i
n
2
\sigma_{out}^2 \propto \sigma_{in}^2
σout2∝σin2。
证明 定理2的证明同样在补充材料B中。
用可视化分析方法验证定理2,在实验中设置
τ
d
e
c
a
y
=
0.25
\tau_{decay}=0.25
τdecay=0.25并展示具有不同预激活方差
σ
i
n
2
\sigma_{in}^2
σin2的膜电势分布情况,结果如图3所示。我们发现在预激活和膜电势的分布之间有高度的相似性,进一步支持了定理2。
接下来,利用LIF模型分析信息前向传播机制。
在前向过程中,当膜电势达到阈值时神经元会发放一个脉冲从而使信息前向逐层传播。根据定理2和等式(3),我们可以近似计算神经元发放脉冲的可能性 P ( u t > V t h ) P(u^t>V_{th}) P(ut>Vth)。不难看出, P ( u t > V t h ) P(u^t>V_{th}) P(ut>Vth)与膜电势的方差 σ i n 2 \sigma_{in}^2 σin2和 σ o u t 2 \sigma_{out}^2 σout2正相关。因此,使用尺度因子来调整预激活的分布,以维持深层SNN的脉冲发放率。图4显示了当我们将预激活的方差 x t ∼ N ( 0 , σ i n 2 ) x^t\sim N(0,\sigma_{in}^2) xt∼N(0,σin2)设为不同值时神经元的发放率分布情况。

由于延迟/衰退(decay)机制,即使神经元每次都接受到正的输入也可能不会触发一个脉冲(如图4a),这意味着下一层的神经元只接受到很小的非零突触前输入,使得脉冲信号在深层SNN中消失,阻止了信号的前向传播。另一种情况是一个神经元一直在发放脉冲(如图4c),这意味着一些神经元的输出对预激活的变化不敏感,会导致计算量的增加。
综上所述,为了平衡突触前输入和阈值以维持发放率,我们利用尺度因子来控制膜电势和预激活的方差,从而减轻其对阈值的依赖。因此,我们将预激活归一化为 N ( 0 , V t h 2 ) N(0,V_{th}^2) N(0,Vth2)。
Deep Spiking Residual Network
当网络变得越来越深时,ResNet是解决退化问题的最流行的架构之一。通过使用shortcut connection,He5等人于2016年在不同层之间增加了恒等映射(identity mapping),使得能够训练深层的神经网络。受到残差学习的启发,本文提出了深度脉冲残差网络(deep spiking residual network),使用tdBN代替了BN层并且修改了shortcut connection以达到更好的性能。
Basic Block:
ANN中的ResNet是由一些basic block(基本块)构建的。图5(a)显示了ResNet-ANN中传统的基本块的形式。
传统的基本块是在相对比较shallow的残差网络中实现的,它包含有两层的 3 × 3 3 \times 3 3×3卷积、BN层和ReLU激活层。在此基础上,我们提出了我们的深度脉冲残差网络的基本块。如图5(b)所示,使用LIF模型替换掉ReLU激活并使用tdBN替换掉BN层。除此之外,还在最终的相加之前修改了shortcut connection并增加了一个tdBN层。然后将最终激活层或shortcut connection中的超参数 α \alpha α设为 1 2 \frac {1}{\sqrt{2}} 21,其它tdBN层的超参数 α \alpha α设为1,这能保证每个激活的输入分布在训练开始时满足 N ( 0 , V t h 2 ) N(0,V_{th}^2) N(0,Vth2)。
网络架构(Network architecture):
使用基本块来构建深度脉冲残差网络。最初的编码层接受输入并直接使用一个stride=2的卷积核进行下采样(downsampling),然后脉冲活动通过基本块传播。和ResNet-ANN类似,当特征图(feature map)减半时我们也将channel数量增加一倍。在平均池化或者全连接层后,最后的解码层使用全连接层和softmax函数。
Experiment
本文将深度残差SNN在静态非脉冲数据集(static non-spiking datasets,如CIFAR-10和ImageNet)和神经形态数据集(neuromorphic dataset,如DVS-gesture和DVS-CIFAR10)上都进行了测试,并将结果与别的SNN模型进行比较,验证了该方法在精度和时间步长上的优势。补充材料C中对数据集的介绍、预处理、参数配置、训练细节和结果分析进行了总结。
静态数据集上的实验:
静态数据集上的实验包括CIFAR-10和ImageNet,它们都是标准的图像识别基准(benchmark)。我们用不同的时间步长、尺寸(size)和深度测试ResNet-SNN。由于脉冲的稀疏性,我们的模型与具有相同架构的ANN相比,在非常少的时间步长的情况下显著减少了计算量。与别的SNN模型相比,ResNet-SNN的速度更快、效果更好!实验结果如下表1所示,详细的分析在补充材料D中。
神经形态数据集上的实验:
与非脉冲的静态数据集相比,神经形态数据集包含更多的时间信息,这更适合于SNN来体现其优势。这里采用了两个具有挑战性的神经形态数据集:DVS-Gesture和DVS-CIFAR10。在DVS-Gesture上面达到了直接训练的SNN(directly-trained SNN)的SOTA(state-of-the-art)结果,实验结果如下表1和表2。
Conclusion
在本文中,我们提出了一种归一化(normalization)方法,使得能够直接训练高性能的深层SNN。我们结合梯度规范理论(gradient norm theory),证明该方法在训练过程中能够有效地平衡输入刺激和神经元阈值,从而促进学习收敛。在此基础上,通过进一步引入shortcut connection,我们将直接训练的SNN从一个常见的浅层结构(不到十层)大大扩展到一个非常深的结构(五十层)。最后,在大规模静态图像数据集和神经形态数据集上对模型进行了评价。与其他SNN模型相比,我们在CIFAR-10和ImageNet上实现了较高的准确率并具有显著更小的推理延迟。据我们所知,这是第一个在ImageNet上报道的一个直接训练的和非常深的SNN的工作。在神经形态数据集上,我们的模型可以有效地处理时空信息,并取得最先进的(state-of-the-art)性能。
综上所述,本工作为深层SNN的直接训练提供了一种可行的方案。它保持了基于脉冲的通信机制的高效率,使得SNN能够解决更复杂的大规模分类任务,有利于在神经形态硬件上的实现并促进SNN的实际应用。
Supplementary Material(待补充)
A Codes for Algorithms(算法代码)

B Proofs of Theorems(定理证明)
C Details of Experiments(实验细节)
D Analysis of Computation Reduction(计算量降低分析)
部分参考文献
Chen Z, Deng L, Wang B, et al. A comprehensive and modularized statistical framework for gradient norm equality in deep neural networks[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020, 44(1): 13-31. ↩︎
Wu Y, Deng L, Li G, et al. Direct training for spiking neural networks: Faster, larger, better[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2019, 33(01): 1311-1318. ↩︎
Wu Y, Deng L, Li G, et al. Spatio-temporal backpropagation for training high-performance spiking neural networks[J]. Frontiers in neuroscience, 2018, 12: 331. ↩︎
Chen Z, Deng L, Wang B, et al. A comprehensive and modularized statistical framework for gradient norm equality in deep neural networks[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020, 44(1): 13-31. ↩︎
He K, Zhang X, Ren S, et al. Deep residual learning for image recognition[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2016: 770-778. ↩︎