脉冲神经网络(SNN)论文阅读(一)-Going Deeper With Directly-Trained Larger Spiking Neural Networks

内容

阐明

准备将自己读的一些和SNN(脉冲神经网络)相关的一些论文记录下来,方便自己以后回看也希望能够帮到有需要的人。
文中一些我认为不重要(极少)的内容被删去,其他部分尽量用专业术语翻译。如有错误或不当之处,希望大家批评指正。

相关信息

论文地址:Going Deeper With Directly-Trained Larger Spiking Neural Networks
论文于2020年12月18日上传至arXiv,由清华大学和加州大学圣巴巴拉分校相关研究人员完成。
Corresponding author: Li Guoqi
李国齐老师团队在SNN领域已经发了不少具有价值的文章,为方便查看它们的关系,列表如下:

文献时间N-MNISTMNIST-DVSImageNetCIAFR10CIFAR10-DVSDVS128-GestureSHD Dataset
文献12019年99.53%90.53%60.5%
文献2(本文)2020年67.05%93.16%67.8%96.87%
文献32021年99.13%70.4%97.56%
文献42021年72%98.61%91.08%
文献52022年99.63%73.8%98.96%

文献1: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.
文献2:Zheng H, Wu Y, Deng L, et al。Going deeper with directly-trained larger spiking neural networks[J]。arXiv preprint arXiv:2011.05280, 2020.
文献3: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.
文献4: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.
文献5: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%的准确率(STOA)、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。
  1. 第一个问题是梯度消失或爆炸。由于其独特的时空信息处理机制和不可微的脉冲信号,当直接训练SNN时,梯度传播表现得非常不稳定,在大多数情况下趋于消失,这影响了深层SNN的效果。到目前为止,直接训练的SNN还没有一种有效的方法可以很好地处理这个问题。前人的工作都没能直接训练深层SNN并且他们的模型都小于10层,严重影响了他们方法的性能。
  2. 另一个问题是我们需要平衡每个神经元的阈值和输入,以在 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%5Et%20%3D%20%5Ctau_%7Bdecay%7Du%5E%7Bt-1%7D%2BI%5Et%20%5Ctag%7B1%7D%2C
其中%5Ctau_%7Bdecay%7D是描述膜电位衰减速度的常数,u%5Et是膜电位,I%5Et是突触前输入,V_%7Bth%7D表示给定的阈值。当u%5Et%3EV_%7Bth%7D时,神经元发放一个脉冲并将u%5Et置为0。突触前输入是由前一层的其他神经元发放的的脉冲信号之和,因此I%5Et可以表示为x%5Et%3D%5Csum_%7Bj%7Dw_jo%5Et%28j%29,其中w_j表示权重,o%5Et%28j%29表示在t时刻来自其他神经元的二进制脉冲输出。考虑到空间结构,并设置u_%7Breset%7D%3D0,则整个空间和时间域的迭代LIF模型可以表示为
u%5E%7Bt%2Cn%2B1%7D%20%3D%5Ctau_%7Bdecay%7Du%5E%7Bt-1%2Cn%2B1%7D%281-o%5E%7Bt-1%2Cn%2B1%7D%29%2Bx%5E%7Bt%2Cn%7D%2C%20%5Ctag%7B2%7D
o%5E%7Bt%2Cn%2B1%7D%20%3D%5Cbegin%7Bcases%7D1%26%7B%5Crm%20if%7D%5C%2Cu%5E%7Bt%2Cn%2B1%7D%3EV_%7Bth%7D%2C%5C%5C%200%26%7B%5Crm%20otherwise%7D%20%5Cend%7Bcases%7D%20%5Ctag%7B3%7D
其中u%5E%7Bt%2Cn%7D表示在第n层神经元在t时刻的膜电位,o%5E%7Bt%2Cn%7D是二进制脉冲,%5Ctau_%7Bdecay%7D是膜电势衰减常数(potential decay constant)。
迭代LIF模型可以在空间和时间维度上实现前向和后向传播,这使得它对一般的机器学习编程框架很友好。

阈值依赖批量正则化(Threshold-dependent batch normalization):

批归一化(batch normalization, BN)作为DNN的正则化组成部分,是目前常用的神经网络方法,它可以实现稳定的收敛和更深层的神经网络。但是,由于SNN额外的时间维度和特殊的激活机制,直接训练的SNN需要专门设计的归一化方法。这促使我们提出阈值依赖的批处理规范化(threshold-dependent batch normalization,tdBN)。
我们考虑一种脉冲卷积神经网络(SCNN)。设o%5Et表示时间步长(timestep)为t时一层内所有神经元的脉冲输出,以及卷积核W和偏差(bias)B,我们有如下表达式
x%5Et%20%3D%20W%20%5Ccircledast%20o%5Et%20%2B%20B%20%5Ctag%7B4%7D
其中x%5Et%20%5Cin%20R%5E%7BN%20%5Ctimes%20C%20%5Ctimes%20H%20%5Ctimes%20W%7D表示以N为批处理轴(batch axis)且时间步长为t时的突触前输入,C是通道轴(channel axis),%28H%2CW%29是空间轴。
在我们的tdBN里,高维的突触前输入将会沿着通道维度被归一化(如下图1所示)。令x%5Et_k表示x%5Et的第k个通道的特征图(feature map),x_k%3D%28x%5E1_k%2Cx%5E2_k%2C...%2Cx%5ET_k%29会按照如下方式进行归一化
%5Chat%7Bx%7D_k%3D%5Cfrac%20%7B%5Calpha%20V_%7Bth%7D%28x_k-E%5Bx_k%5D%29%7D%7B%5Csqrt%7BVar%5Bx_k%5D%2B%5Cepsilon%7D%7D%2C%20%5Ctag%7B5%7D
y_k%20%3D%20%5Clambda_k%20%5Chat%7Bx%7D_k%20%2B%20%5Cbeta_k%2C%20%5Ctag%7B6%7D
其中V_%7Bth%7D表示阈值,%5Calpha是一个依赖于网络结构的超参数,%5Cepsilon是一个很小的常数,%5Clambda_k%5Cbeta_k是两个可以训练的参数,E%5Bx_k%5D%E3%80%81Var%5Bx_k%5D是在一个MiniBatch上统计估计的x_k的均值和方差。图1展示了E%5Bx_k%5D%E3%80%81Var%5Bx_k%5D的计算过程,可定义为
E%5Bx_k%5D%20%3D%20mean%28x_k%29%2C%20%5Ctag%7B7%7D
Var%5Bx_k%5D%20%3D%20mean%28%28x_k%20-%20E%5Bx_k%5D%29%5E2%29.%20%5Ctag%7B8%7D
因此,在训练期间,y_k%20%5Cin%20R%5E%7BT%20%5Ctimes%20N%20%5Ctimes%20H%20%5Ctimes%20W%7D正是下一层第k个通道的神经元在T时间步长内接收到的归一化突触前输入。
图1

图1 在tdBN中计算E[x]和Var[x],每一个立方体表示在t时间步长时的一个feature map tensor,其中N是batch axis,C是channel axis,(H,W)是空间轴。在C维向量E[x]和Var[x]中的每一个元素由对应通道的黄色张量估计。

在推理过程中,我们遵循标准Batch Normalization的模式(schema),分别估计出代表E%5Bx_k%5DVar%5Bx_k%5D在整个数据集上的期望%5Cmu_%7Binf%7D%5Csigma%5E2_%7Binf%7D,并在训练过程中通过移动平均求解(moving average solution)得到。
除此之外,在对具有tdBN的SNN进行推理时,进行批量尺度融合(batchnorm-scale-fusion)是必要的。它移除了推理期间的Batch Normalization操作,从而保持网络是全脉冲(full-spiking)的,并使其能够在神经形态平台上实现。设W_%7Bc%2Ck%7DB_%7Bc%2Ck%7D表示一层里的第c个feature map与下一层的第k个feature map之间的卷积核和偏差,模式(schema)可以由下式定义
W%5E%7B%27%7D_%7Bc%2Ck%7D%20%3D%20%5Clambda_k%20%5Cfrac%20%7B%5Calpha%20V_%7Bth%7DW_%7Bc%2Ck%7D%7D%7B%5Csqrt%7B%5Csigma%5E2_%7Binf%2Ck%7D%2B%5Cepsilon%7D%20%7D%2C%20%5Ctag%7B9%7D
B%5E%7B%27%7D_%7Bc%2Ck%7D%20%3D%20%5Clambda_k%20%5Cfrac%7B%5Calpha%20V_%7Bth%7D%28B_%7Bc%2Ck%7D-%5Cmu_%7Binf%2Ck%7D%29%7D%7B%5Csqrt%7B%5Csigma%5E2_%7Binf%2Ck%7D%2B%5Cepsilon%7D%7D%2B%20%5Cbeta_k%2C%20%5Ctag%7B10%7D
其中的W%5E%7B%27%7D_%7Bc%2Ck%7DB%5E%7B%27%7D_%7Bc%2Ck%7D表示经过batchnorm-scale-fusion后的权重和偏差,因此在推理过程中,脉冲通过变换后的权重W%5E%7B%27%7D_%7Bc%2Ck%7D和偏差B%5E%7B%27%7D_%7Bc%2Ck%7D逐层传播,就不需要进行batchnorm操作。因此,我们的tdBN只影响训练期间的计算量而不影响已经训练好的SNN的运行过程。
简而言之,我们的tdBN与标准BN有两个主要区别。首先,与ANN不同的是,SNN不仅逐层(layer by layer)地传播信息,而且从上一时刻刻传播到下一时刻。因此,tdBN应该在时间和空间两个维度上对特征输入进行归一化。其次,根据阈值V_%7Bth%7D归一化方差。在tdBN中,预激活(pre-activations )被归一化为N%280%2C%28%5Calpha%20V_%7Bth%7D%29%5E2%29而不是N%280%2C1%29。我们用1和0初始化可训练参数%5Clambda%5Cbeta,在串联神经网络里(serial neural network)超参数%5Calpha都设为1。对于有n个分支的局部并行网络结构,%5Calpha%3D%20%5Cfrac%7B1%7D%7B%5Csqrt%7Bn%7D%7D。这使得在早期的训练过程中V_%7Bth%7D的预激活均值为0,标准差为V_%7Bth%7D。tdBN的代码见文末的补充材料A。

整体训练算法(Overall training algorithm):

在本节中,我们介绍STBP-tdBN的整体训练算法,用我们的tdBN从头开始训练深层SNN。
在误差反向传播中,我们将最后一层作为解码层(decoding layer),最终的输出Q表示为:
Q%20%3D%20%5Cfrac%20%7B1%7D%7BT%7D%20%5Csum_%7Bt%3D1%7D%5E%7BT%7DMo%5E%7Bn%2Ct%7D%2C%20%5Ctag%7B11%7D
其中o%5E%7Bn%2Ct%7D是最后一层输出层发放的脉冲,M是解码层矩阵,T是时间步长(timestep)。
然后我们让输出通过一层softmax层。损失函数(loss function)使用交叉熵。考虑有输出Q%3D%28q_1%2Cq_2%2C...%2Cq_n%29和标签向量(label vector)Y%3D%28y_1%2Cy_2%2C...%2Cy_n%29,损失函数L定义为:
p_i%20%3D%20%5Cfrac%7Be%5E%7Bq_i%7D%7D%7B%5Csum_%7Bj%3D1%7D%5Ene%5E%7Bq_i%7D%7D%2C%20%5Ctag%7B12%7D
L%20%3D%20-%5Csum_%7Bi%3D1%7D%5Eny_ilog%28P_i%29.%20%5Ctag%7B13%7D
利用迭代LIF模型,STBP-tdBN方法在空间和时间域上反向传播损失函数L的梯度。利用链式法则,可以利用下式来计算%5Cfrac%7B%5Cpartial%20L%7D%7B%5Cpartial%20o_i%5E%7Bt%2Cn%7D%7D%5Cfrac%7B%5Cpartial%20L%7D%7B%5Cpartial%20u_i%5E%7Bt%2Cn%7D%7D
%5Cfrac%7B%5Cpartial%20L%7D%7B%5Cpartial%20o_i%5E%7Bt%2Cn%7D%7D%20%3D%20%5Csum_%7Bj%3D1%7D%5E%7Bl%28n%2B1%29%7D%5Cfrac%7B%5Cpartial%20L%7D%7B%5Cpartial%20u_j%5E%7Bt%2Cn%2B1%7D%7D%5Cfrac%7B%5Cpartial%20u_j%5E%7Bt%2Cn%2B1%7D%7D%7B%5Cpartial%20o_i%5E%7Bt%2Cn%7D%7D%20%2B%20%5Cfrac%7B%5Cpartial%20L%7D%7B%5Cpartial%20u_i%5E%7Bt%2B1%2Cn%7D%7D%5Cfrac%7B%5Cpartial%20u_i%5E%7Bt%2B1%2Cn%7D%7D%7B%5Cpartial%20o_i%5E%7Bt%2Cn%7D%7D%2C%20%5Ctag%7B14%7D
%5Cfrac%7B%5Cpartial%20L%7D%7B%5Cpartial%20u_i%5E%7Bt%2Cn%7D%7D%20%3D%20%5Cfrac%7B%5Cpartial%20L%7D%7B%5Cpartial%20o_i%5E%7Bt%2Cn%7D%7D%5Cfrac%7B%5Cpartial%20o_i%5E%7Bt%2Cn%7D%7D%7B%5Cpartial%20u_i%5E%7Bt%2Cn%7D%7D%20%2B%20%5Cfrac%7B%5Cpartial%20L%7D%7B%5Cpartial%20u_i%5E%7Bt%2B1%2Cn%7D%7D%5Cfrac%7B%5Cpartial%20u_i%5E%7Bt%2B1%2Cn%7D%7D%7B%5Cpartial%20u_i%5E%7Bt%2Cn%7D%7D%2C%20%5Ctag%7B15%7D
其中o%5E%7Bt%2Cn%7Du%5E%7Bt%2Cn%7D表示第n层神经元在t时刻的脉冲和膜电位,由于脉冲活动不可微分,%5Cfrac%7B%5Cpartial%20o%5Et%7D%7B%5Cpartial%20u%5Et%7D实际上并不存在,为了解决这个问题,Wu3等人在2018年提出了导数曲线(derivative curve)来近似脉冲活动的导数,在本文中,我们使用矩形函数(rectangular function),它在梯度下降中是很有效的,可以由下式求得:
%5Cfrac%20%7B%5Cpartial%20o%5Et%7D%7B%5Cpartial%20u%5Et%7D%20%3D%20%5Cfrac%20%7B1%7D%7Ba%7Dsign%28%7Cu%5Et-V_%7Bth%7D%7C%20%3C%20%5Cfrac%20%7Ba%7D%7B2%7D%29.%20%5Ctag%7B16%7D
整体训练算法的代码同样见文末的补充材料A。

Theoretical Analysis

在本节中,我们将分析tdBN对STBP-tdBN训练的SNN的影响。利用神经网络中梯度规范理论(gradient norm theory)的理论工具,我们发现我们的tdBN可以缓解训练过程中梯度消失或爆炸的问题。我们还将解释归一化过程中加入的尺度因子%5CalphaV_%7Bth%7D的作用。

梯度规范理论(Gradient norm theory):

梯度范数理论是近年来发展较好的理论,它旨在克服各种神经网络结构中的梯度消失或爆炸问题。在本文中,我们采用Chen4等人在2020年提出的“Block Dynamical Isometry”来分析tdBN在直接训练的SNN中的效果,它把网络看作为一系列的块(block):
f%28x%29%20%3D%20f_%7Bi%2C%5Ctheta_i%7D%20%5Ccirc%20f_%7Bi-1%2C%5Ctheta_%7Bi-1%7D%7D%20%5Ccirc%20...%20%5Ccirc%20f_%7B1%2C%5Ctheta_1%7D%28x%29%2C%20%5Ctag%7B17%7D
其中函数f_%7Bj%2C%5Ctheta_j%7D表示第j个块,并将其输入输出雅可比矩阵定义为%5Cfrac%20%7B%5Cpartial%20f_j%7D%7B%5Cpartial%20f_%7Bj-1%7D%7D%20%3D%20J_j,使用%5Cphi%28J%29表示tr%28J%29的期望,%5Cvarphi%28J%29表示%5Cphi%28J%5E2%29%20-%20%5Cphi%20%5E2%28J%29,随后证明下面几个引理(lemma):

引理1.考虑一个神经网络,它可以如式(17)那样用一系列的块表示,第j个块的雅可比矩阵表示为J_j。如果%5Cforall%20j%2C%5Cphi%28J_jJ_j%5ET%29%5Capprox%201%5Cvarphi%28j_jJ_j%5ET%29%5Capprox%200,网络就达到了“块动态等距(Block Dynamical Isometry)”,并能避免梯度消失或梯度爆炸。

引理2.考虑一个神经网络块,包含有0均值的数据归一化(data normalization with 0-mean)、线性转换(linear transform )和整流激活rectifier activations(General Linear Transform)。设输入输出向量的二阶矩分别为%5Calpha_%7Bin%7D%5Calpha_%7Bout%7D,有%5Cphi%28JJ%5ET%29%3D%5Cfrac%7B%5Calpha_%7Bout%7D%7D%7B%5Calpha_%7Bin%7D%7D

在梯度规范理论框架的基础上,将梯度规范与脉冲神经元的特性相结合,进一步分析tdBN算法对SNN的有效性。
LIF模型有两个独特的超参数:%5Ctau_%7Bdecay%7DV_%7Bth%7D,其中%5Ctau_%7Bdecay%7D影响梯度在时间域中的传播,V_%7Bth%7D影响空间维度。在SNN的实验中,%5Ctau_%7Bdecay%7D通常设置为比较小的值(如0.25)。为了分析梯度变换,我们将模型简化,设置%5Ctau_%7Bdecay%7D为0,得到如下命题:

定理1:考虑一个时间步长为T的SNN,在T时刻的第j个雅可比矩阵记为J%5Et_j。当%5Ctau_%7Bdecay%7D%3D0时,将两个tdBN层之间每个块(block)的输入向量和输出向量的二阶矩固定为V_%7Bth%7D%5E2,有%5Cphi%28J_j%5Et%28J_j%5Et%29%5ET%29%5Capprox1并且可以在SNN的训练过程中避免梯度消失或梯度爆炸。
证明:定理1的证明基于引理1和引理2,详细的证明过程在补充材料B中。

膜电势衰减机制的影响(Influence of membrane decay mechanism):

下面分析%5Ctau_%7Bdecay%7D对梯度传播的影响。从等式(2)到等式(15),可以得到:
%5Cfrac%7B%5Cpartial%20L%7D%7B%5Cpartial%20u_i%5E%7Bt%2Cn%7D%7D%20%3D%20%5Cfrac%7B%5Cpartial%20L%7D%7B%5Cpartial%20o_i%5E%7Bt%2Cn%7D%7D%5Cfrac%7B%5Cpartial%20o_i%5E%7Bt%2Cn%7D%7D%7B%5Cpartial%20u_i%5E%7Bt%2Cn%7D%7D%20%2B%20%5Cfrac%7B%5Cpartial%20L%7D%7B%5Cpartial%20u_i%5E%7Bt%2B1%2Cn%7D%7D%5Ctau_%7Bdecay%7D%281-o_i%5E%7Bt%2Cn%7D%29.%20%5Ctag%7B18%7D
如果一个神经元发出一个脉冲%281-o_i%5E%7Bt%2Cn%7D%29%3D0,则梯度与%5Ctau_%7Bdecay%7D无关。另一方面,由于%5Ctau_%7Bdecay%7D是一个很小的常数,神经元在时间%5Ctau_%7Bdecay%7D的梯度对时间%5Ctau_%7Bdecay%7D的影响很小。
为了验证定理1和分析膜电势衰减机制的影响,我们在CIFAR-10数据集上用20层的plain spiking network评估tdBN。在图2中,我们展示了在训练过程中的前1/6个epoch中除了第一个编码层和最后一个输出层外的每一层梯度范数(gradient norm)的均值。结果表明,当%5Ctau_%7Bdecay%7D= 0时,梯度范数曲线表现较稳定,证实了我们的理论。但要注意的是,%5Ctau_%7Bdecay%7D不能设置为0,因为这样会阻碍信息沿时间维度传播从而造成严重的退化问题。因此我们在%5Ctau_%7Bdecay%7D%5Cneq0的条件下评估我们的方法。当%5Ctau_%7Bdecay%7D设置为0.25和0.5时,梯度范数随着网络的加深增长非常缓慢,也就不会影响训练过程。研究结果有力地支持了我们的结论,即在深层SNN中可以避免梯度消失或爆炸的问题。
在图2

图2 基于tdBN的plain spiking network的gradient norm

尺度因子(Scaling factors):

众所周知,SNN模型获得有竞争力的性能的关键是设置合适的阈值以保持发放率(firing rate)和减少信息丢失。为了实现这一点,我们在tdBN的归一化实现中引入两个尺度因子用来平衡预激活(pre-activation)和阈值。在早期的训练过程中,有%5CalphaV_%7Bth%7D两个尺度因子,通过将可训练参数%5Clambda%5Cbeta初始化为1和0来将预激活归一化为N%280%2CV_%7Bth%7D%5E2%29
首先,我们提出定理2来解释预激活和神经元膜电势之间的关系,以助于理解为什么我们的方法是有效的。

定理2采用迭代LIF模型,假设预激活x%5Et%20%5Csim%20N%280%2C%5Csigma_%7Bin%7D%5E2%29,则膜电势u%5Et%20%5Csim%20N%280%2C%5Csigma_%7Bout%7D%5E2%29%5Csigma_%7Bout%7D%5E2%20%5Cpropto%20%5Csigma_%7Bin%7D%5E2
证明定理2的证明同样在补充材料B中。

用可视化分析方法验证定理2,在实验中设置%5Ctau_%7Bdecay%7D%3D0.25并展示具有不同预激活方差%5Csigma_%7Bin%7D%5E2的膜电势分布情况,结果如图3所示。我们发现在预激活和膜电势的分布之间有高度的相似性,进一步支持了定理2。
图3

图3 具有不同预激活方差的膜电势分布
P ( u t > V t h ) P(u^t>V_{th}) P ( u t > V t h ​ ) 。不难看出, P ( u t > V t h ) P(u^t>V_{th}) P ( u t > V t h ​ ) 与膜电势的方差 σ i n 2 \sigma_{in}^2σ i n 2 ​σ o u t 2 \sigma_{out}^2σ o u t 2 ​正相关。因此,使用尺度因子来调整预激活的分布,以维持深层SNN的脉冲发放率。图4显示了当我们将预激活的方差 x t ∼ N ( 0 , σ i n 2 ) x^t\sim N(0,\sigma_{in}^2) x t ∼ N ( 0 , σ i n 2 ​ ) 设为不同值时神经元的发放率分布情况。

图4

图4 具有不同预激活方差的神经元发放率(firing rate)分布
N ( 0 , V t h 2 ) N(0,V_{th}^2) N ( 0 , V t h 2 ​ )

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中传统的基本块的形式。
图5

图5 传统的ResNet-ANN和带有tdBN的ResNet-SNN中不同的basic block

传统的基本块是在相对比较shallow的残差网络中实现的,它包含有两层的3%20%5Ctimes%203卷积、BN层和ReLU激活层。在此基础上,我们提出了我们的深度脉冲残差网络的基本块。如图5(b)所示,使用LIF模型替换掉ReLU激活并使用tdBN替换掉BN层。除此之外,还在最终的相加之前修改了shortcut connection并增加了一个tdBN层。然后将最终激活层或shortcut connection中的超参数%5Calpha设为%5Cfrac%20%7B1%7D%7B%5Csqrt%7B2%7D%7D,其它tdBN层的超参数%5Calpha设为1,这能保证每个激活的输入分布在训练开始时满足N%280%2CV_%7Bth%7D%5E2%29

网络架构(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中。
表1

表1 ResNet-SNN与别的SNN模型在CIFAR-10和ImageNet数据集上的比较结果

神经形态数据集的实验:

与非脉冲的静态数据集相比,神经形态数据集包含更多的时间信息,这更适合于SNN来体现其优势。这里采用了两个具有挑战性的神经形态数据集:DVS-Gesture和DVS-CIFAR10。在DVS-Gesture上面达到了直接训练的SNN(directly-trained SNN)的SOTA(state-of-the-art)结果,实验结果如下表1和表2。
表2
表3

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(算法代码)

算法代码1

图A-1 tdBN算法代码

算法代码2

图A-2 整体训练代码

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.↩︎

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
乘风的头像乘风管理团队
上一篇 2022年3月21日 下午8:10
下一篇 2022年3月22日 上午11:04

相关推荐