Going Deeper With Directly-Trained Larger Spiking Neural Networks

郑重声明：原文参见标题，如有侵权，请联系作者，将会撤销发布！

The Thirty-Fifth AAAI Conference on Artificial Intelligence (AAAI-21)

 

Abstract

　　脉冲神经网络（SNN）在时空信息和事件驱动信号处理的生物可编程编码中很有前途，非常适合在神经形态硬件中实现节能。然而，SNN独特的工作模式使其比传统网络更难训练。目前，有两条主要路线来探索高性能深度SNN的训练。第一种是将预先训练的ANN模型转换为其SNN版本，SNN版本通常需要很长的编码窗口才能收敛，并且不能在训练过程中利用时空特征来解决时间任务。另一种是在时空域中直接训练SNN。但是，由于发放函数的二值脉冲活动和梯度消失或爆炸的问题，当前的方法仅限于浅层架构，因此难以利用大规模数据集（例如ImageNet）。为此，我们提出了一种基于新兴时空反向传播的阈值相关批归一化（tdBN）方法，称为“STBP-tdBN”，能够直接训练非常深的SNN，并在神经形态硬件上有效实现其推理。通过所提出的方法和详细的跳连，我们将直接训练的SNN从浅层结构（<10层）显著扩展到非常深的结构（50层）。此外，我们还从理论上分析了基于“块体动力学等距”（Block Dynamical Isometry）理论的方法的有效性。最后，我们报告了优越的准确率结果，包括在CIFAR-10上的93.15%，在DVS-CIFR10上的67.8%，以及在ImageNet上的67.05%，时间步长非常少。据我们所知，这是第一次在ImageNet上探索直接训练的高性能深度SNN。我们相信这项工作将为充分利用SNN的优势铺平道路，并吸引更多的研究人员在该领域做出贡献。

 

Introduction

　　受人类神经元工作模式的启发，脉冲神经网络（SNN）被认为是人工智能和理论神经科学中一种很有前途的模型（Roy, Jaiswal, and Panda 2019）。得益于内在的神经元动力学和事件驱动的脉冲通信范式，SNN在连续的时空信息处理中显示出巨大的潜力，具有更低的能耗和更好的鲁棒性（Stromatias et al. 2015）。此外，SNN可以很容易地应用于一些专门的神经形态硬件（Merolla et al. 2014; Davies et al. 2018），这可能被视为下一代神经网络。

　　有两种主要的方法来训练具有高性能的SNN。一种是将预先训练的ANN转换为SNN模型，通常需要数百个时间步骤（Sengupta et al. 2019; Hu et al. 2018）。因此，尽管这些SNN实现了与具有类似结构的ANN相当的精度，但大量的时间步骤会导致严重的信号延迟并增加计算量。另一种是基于梯度下降方法直接训练神经网络，该方法独立于预先训练的神经网络，可以减少时间步长。最近的一项工作（Wu et al. 2018）提出了一种称为“时空反向传播”（STBP）的学习算法，该算法具有在面向人工神经网络的编程框架（例如Pytorch）上直接训练SNN的方法，为我们探索更深且更大的直接训练的SNN提供了机会。然而，该算法训练的SNN仅限于浅层架构，无法在ImageNet等大规模数据集上获得令人满意的性能。因此，在上述算法下，我们明确了直接训练更深的SNN需要解决的两个问题。

　　第一个问题是梯度消失或爆炸。由于时空信息处理的特殊机制和不可微分的脉冲信号，当直接训练SNN时，梯度传播表现得非常不稳定，并且在大多数情况下趋于消失，这阻止了SNN变深。到目前为止，在直接训练的SNN中还没有一种有效的方法来很好地处理这个问题。我们之前的著名工作（Lee et al. 2020; Lee, Delbruck, and Pfeiffer 2016; Wu et al. 2019）未能直接训练深度SNN，并且它们的所有模型都小于10层，这严重影响了它们方法的性能。

　　另一个问题是，我们需要平衡每个神经元的阈值和输入，以在SNN中获得适当的发放率。当输入与阈值相比太小时，神经元发出很少的脉冲，神经元膜电位保持不变，因此神经元处理的信息无法充分表达。当输入过大时，神经元会一直发放，并且对输入的变化不敏感。对于直接训练的SNN，随着二值脉冲逐层传播，突触前输入的分布将在训练过程中发生变化，使得输入的大小不合适。已经提出了许多方法来处理它，例如阈值正则化（Lee, Delbruck, and Pfeiffer 2016）和NeuNorm（Wu et al. 2019)。

　　归一化似乎是解决这两个问题的合适方法。根据（Chen et al. 2020），它们稳定了网络和梯度传播。此外，他们将突触前输入的分布归一化为相同的期望和方差，这有助于通过减少内部协变量偏移来平衡阈值和输入。然而，现有的归一化方法不适合于SNN的训练。对于额外的时间维度和特殊的激活机制，直接训练的SNN需要专门设计的归一化方法。

　　在本文中，我们开发了一种新的算法来直接训练深度SNN。这项工作的主要贡献总结如下：

• 我们提出了依赖阈值的批归一化来解决梯度消失或爆炸问题，并调整发放率。此外，我们采用了剩余的网络结构，并修改了适用于深度SNN的跳连。
• 在此基础上，我们研究了非常深的直接训练SNN（将其从不到10层扩展到50层），并在大规模非脉冲数据集（CIFAR-10，ImageNet）和神经形态数据集（DVS-Gesture，DVS-CIFR10）上对其进行了测试。
• 据我们所知，在CIFAR-10和ImageNet上，我们全面验证了不同的SNN架构（ResNet-18, 34, 50），并报告了与时间步长少得多（不超过6个时间步长）的类似SNN相比具有竞争力的结果，这是直接训练的全脉冲SNN首次在ImageNet上报告相当高的准确性。在神经形态数据集上，我们的模型在DVS-Gesture和DVS-CFAR10上都实现了最先进的性能，这表明了SNN在处理时空信息方面的优势。

 

Related Work

Learning algorithm of SNNs 在过去的几年里，许多学习算法探索了如何训练高性能的深度神经网络，包括：(1) 将预先训练的神经网络转换为神经网络的一些方法；(2) 基于梯度下降的算法。

　　第一种被称为“ANN-SNN转换方法”（Sengupta et al. 2019; Han, Srinivasan, and Roy 2020），被视为训练高性能深度SNN的最流行方法，它将ReLU函数的实值输出转换为SNN中的二值脉冲。与ANN相比，这种方法在大规模数据集上成功地报告了具有竞争力的结果，而没有严重退化。然而，它忽略了脉冲神经元丰富的时间动态行为，通常需要数百或数千个时间步长才能接近预先训练的神经网络的准确性。

　　基于梯度下降的算法利用误差反向传播来训练SNN。通过梯度下降优化学习算法，一些SNN模型（Lee, Delbruck, and Pfeiffer 2016; Jin, Zhang, and Li 2018; Lee et al. 2020）在CIFAR-10和其他神经形态数据集上实现了高性能。其中，（Wu et al. 2019）将LIF模型（Hunsberger and Eliasmith 2015）改进为迭代LIF模型，并开发了STBP学习算法，使其在面向ANN的编程框架上友好，并加快了训练过程。此外，训练SNN在处理空间和时间信息方面直接显示出巨大的潜力，并在很少的时间步长内报告高精度。然而，由于梯度消失和内部协变量偏移，它无法直接训练非常深的SNN，这正是我们想要克服的。

Gradient vanishing or explosion in the deep neural network (DNN) 当DNN是动态等距时，它可以避免梯度消失或爆炸，这意味着它的输入输出雅可比矩阵的每个奇异值都保持接近1。（Chen et al.2020）提出了一种度量——“块动态等距”，作为所有复杂串-并行DNN的通用统计工具。它研究了神经网络中每个块的一阶矩和二阶矩，并分析了它们对梯度分布的影响。此外，它还对神经网络中权重初始化、批量归一化和跳连的功能进行了理论解释，这有助于我们开发算法。

Normalization 归一化技术能够训练表现良好的神经网络。对于人工神经网络，归一化，如批归一化（Ioffe and Szegedy 2015）、组归一化（Wu and He 2018）和层归一化（Ba, Kiros, and Hinton 2016），已成为常见的方法。批归一化（BN）通过减少内部协变量偏移来加速深度网络训练，从而实现更高的学习率并使模型正则化。虽然BN会导致高学习延迟并增加计算量，但它可以使网络更深，避免梯度消失或爆炸。对于SNN，研究人员提出了其他归一化技术，如基于数据的归一化（Diehl et al. 2015）、Spike-Norm（Sengupta et al. 2019）和NeuNorm（Wu et al. 2019）。这些归一化方法旨在平衡输入和阈值，以避免严重的信息丢失，但它们对我们直接训练的深度SNN无效，因为它们仍然忽略了梯度消失的问题。我们注意到了ANN中BN的影响以及SNN中输入分布的重要性，因此我们对BN进行了修改，以满足SNN模型的训练和推理。

 

Materials and Methods

Iterative LIF Model

　　迭代LIF模型最早由（Wu et al. 2019）提出，他利用欧拉方法求解LIF模型的一阶微分方程，并将其转换为迭代表达式：

其中τ_decay是一个常数，用来描述膜电位衰减的速度，u^t是膜电位，它是突触前输入。设V_th表示给定的阈值。当u^t＞V_th时，神经元发出脉冲，u^t将重置为0。突触前输入是来自最后一层的其他神经元的累积脉冲。因此，它可以用来表示，其中w_j是权重，o^t(j)表示在 t 时刻来自其他人的二值脉冲输出。考虑到空间结构并设置u_reset＝0，整个空间和时间域的迭代LIF模型可以由下式定义：

其中u^t,n是第n层神经元在 t 时刻的膜电位，o^t,n是二值脉冲，衰减是电压衰减常数。

　　迭代LIF模型能够在空间和时间维度上实现前向和后向传播，这使得它对通用机器学习编程框架很友好。

 

Threshold-dependent Batch Normalization

　　作为DNN的一个常规组成部分，批处理归一化（BN）一直是当前神经网络的一种常见方法，它允许稳定的收敛和更深层次的神经网络。然而，由于额外的时间维度和特殊的激活机制，直接训练的SNN需要专门设计的归一化方法。这促使我们提出了阈值相关的批归一化（tdBN）。

　　我们考虑一种脉冲卷积神经网络（SCNN）。设o^t表示一层中所有神经元在时间步骤 t 的脉冲输出。对于卷积核W和偏差B，我们得到：

其中x^t ∈ R^N×C×H×W表示时间步骤 t 处的突触前输入，其中N作为批处理轴，C作为通道轴，(H, W)作为空间轴。

　　在我们的tdBN中，高维突触前输入将沿着通道维度进行归一化（图1）。设表示x^t的第k个通道特征图。那么将通过下式归一化：

其中，V_th表示阈值，α是取决于网络结构的超参数，ε是微小常数，λ_k和β_k是两个可训练参数，E[x_k], Var[x_k]是在Mini-Batch中统计估计的x_k的均值和方差。图1显示了如何计算E[x_k], Var[x_k]，其定义如下：

因此，在训练过程中，y_k ∈ R^T×N×H×W正是由下一层的第k个通道神经元在时间步长T期间接收的归一化突触前输入。

　　在推断中，我们遵循作为标准批归一化的模式来估计μ_inf和，它们分别表示整个数据集上E[x_k]和Var[x_k]的期望，可以在训练过程中通过移动平均解来计算。

　　此外，对于具有tdBN的SNN，在推断中需要进行批范数尺度融合。它消除了推断过程中的批归一化操作，从而保持网络完全脉冲，并使其能够在神经形态平台上实现。设W_c,k和B_c,k表示一层中的第c个特征图和下一层中第k个特征图之间的卷积核和偏差。架构由下式决定：

其中和偏差表示在分批范数尺度融合之后的变换权重。因此，在推断期间，脉冲通过变换的权重和偏差逐层传播，而不进行批归一化操作。因此，我们的tdBN只影响训练期间的计算成本，而不影响已经训练的SNN的运行机制。

　　简言之，我们的tdBN与标准BN有两个主要区别。首先，与ANN不同，SNN不仅逐层传播信息，而且从最后一刻到下一刻都在传播信息。因此，tdBN应该在时间和空间维度上对特征输入进行归一化。其次，我们使归一化方差依赖于阈值。在tdBN中，预激活被归一化为N(0, (αV_th)²)，而不是N(0, 1)。我们将初始化可训练参数λ和β为1和0。在串行神经网络中，超参数α为1。对于具有n个分支的局部并行网络结构，α将为。它在早期训练时进行均值为0和标准差为V_th的预激活。tdBN的代码可以在补充材料A中找到。

 

Overall Training Algorithm

　　在本节中，我们介绍了STBP-tdBN的整体训练算法，用于用我们的tdBN从头开始训练深度SNN。

　　在误差反向传播中，我们将最后一层视为解码层，最终输出Q可以通过下式计算：

其中o^n,t是最后一个输出层发放的脉冲，M是解码层的矩阵，t表示时间步数。

　　然后我们使输出通过一个softmax层。损失函数被确定为交叉熵。考虑到输出Q = (q₁, q₂, ... , q_n)和标签向量Y = (y₁, y₂, ... , y_n)，损失函数L被定义为：

利用迭代LIF模型，STBP-tdBN方法在空间和时间域上反向传播损失L的梯度。通过应用链式规则，和可以通过以下公式计算：

其中o^t,n和u^t,n表示n层神经元在时间 t 的脉冲和膜电位。由于不可微分的脉冲活动，实际上并不存在。为了解决这个问题，（Wu et al. 2018）提出了导数曲线来近似脉冲活动的导数。在本文中，我们使用矩形函数，它被证明在梯度下降中是有效的，并且可以由下式决定：

整个训练算法的代码如补充材料A所示。

 

Theoretical Analysis

　　在本节中，我们将分析tdBN对STBP-tdBN训练的SNN的影响。利用人工神经网络中梯度范数理论的理论工具，我们发现我们的tdBN可以缓解训练过程中梯度消失或爆炸的问题。我们还解释了我们在归一化过程中添加的比例因子α和V_th的函数。

 

Gradient Norm Theory

　　梯度范数理论近年来得到了很好的发展，旨在克服各种神经网络结构中的梯度消失或爆炸。在本文中，我们采用（Chen et al.2020）提出的“块动态等距”来分析直接训练的SNN中tdBN的影响。它将网络视为一系列块：

其中函数表示第 j 个块，并将其输入-输出雅可比矩阵定义为。设Φ(J)表示tr(J)的期望，φ(J)代表Φ(J²) - Φ²(J)所示。然后他们证明了以下引理。

　　基于梯度范数的理论框架，我们将其与脉冲神经元的独特性质相结合，进一步分析了我们提出的tdBN算法对SNN的有效性。

　　LIF模型有两个特殊的超参数：τ_decay和V_th，其中τ_decay影响时域中的梯度传播，V_th影响空间维度。对于SNN的实验，τ_decay通常被设置为小值（例如0.25）。为了分析梯度变换，我们简化模型并将τ_decay设置为零，我们可以得到以下命题。

证据. 定理1的证明是基于引理1和引理2。详细信息见补充材料B。

Influence of membrane decay mechanism 我们分析了衰减对梯度传播的影响。根据公式(2)和(15)，我们得到：

也就是说，如果一个神经元发出一个脉冲，等于零，并且梯度与τ_decay无关。另一方面，由于τ_decay是一个很小的常数，神经元在时间t+1的梯度对时间 t 的梯度几乎没有影响。

　　为了验证定理1和我们对膜衰变机制影响的分析，我们在CIFAR-10上评估了20层脉冲网络中的tdBN。在图2中，我们显示了在训练过程的第一个1/6 epoch期间每一层（尽管有第一个编码层和最后一个解码层）中的梯度范数的平均值。我们发现，当τ_decay=0时，梯度范数的曲线表现得非常稳定，这完全支持了我们的理论。需要强调的是，τ_decay不能设置为0，因为它会阻止信息沿时间维度传播并导致严重退化。因此，我们在τ_decay≠0的条件下评估我们的方法。例如，当τ_decay=0.25和0.5时，随着网络的加深，梯度范数的增加非常缓慢，这不会影响训练过程。这些结果有力地支持了我们的结果，可以避免深度SNN中的梯度消失或爆炸。

 

Scaling Factors

　　众所周知，SNN模型获得竞争性能的关键是设置合适的阈值以保持发放率并减少信息损失。为了实现这一点，我们在tdBN中的归一化实现中引入了两个缩放因子，这意味着平衡预激活和阈值。在早期训练中，使用两个缩放因子α和V_th，我们通过初始化可训练参数λ, β为1和0将预激活归一化为。

　　首先，我们提出定理2来解释神经元的每次激活和膜电位之间的关系，这有助于理解我们的方法为什么有效。

 

 

证据. 定理2的证明见补充材料B。

（续）

 

 

Deep Spiking Residual Network

　　ResNet是解决网络深入时降级问题的最流行的架构之一。（He et al. 2016）增加了层之间的恒等映射，这使得能够训练非常深的神经网络。受残差学习的启发，我们提出了我们的深度脉冲残差网络，它用我们的tdBN取代了BN层，并改变了跳连。

Basic block

Network architecture 我们用我们的基本块构建了深度脉冲残差网络。第一编码层接收输入，并使用卷积核和2的步长直接执行下采样。然后，脉冲活动通过基本块传播。与ResNet-ANN类似，当特征图减半时，我们将通道数量增加一倍。在平均池化或全连接层（如有必要）之后，最后一个解码层是一个全连接层，后面跟着softmax函数。

 

Experiment

 

Experiment on Static Datasets

 

Experiment on Neuromorphic Datasets

 

Conclusion