脉冲神经网络（SNN）论文阅读（五）-----AAAI-2024 时间步长逐渐收缩的SNN

原文链接：CSDN-脉冲神经网络（SNN）论文阅读（五）-----AAAI-2024 时间步长逐渐收缩的SNN

目录

论文信息

论文地址
arXiv版本（带有附录）地址
由电子科技大学（左琳教授团队）研究人员发表于AAAI 2024。

主要贡献

• 第一个具有异质性时间步长的SNN
• 论文提出了Shrinking SNN (SSNN)，SNN的时间步长逐渐收缩，相比于普通的SNN，SSNN的平均时间步长更低，从而实现低延迟、高性能的神经形态目标识别；Timestep Shrinkage方法
• 论文在SNN的多个阶段后引入额外的分类器辅助训练，推理时则无需辅助分类器。从而缓解了代理梯度和真实梯度的误差以及梯度消失/爆炸问题。Early Classifier方法
  

Timestep Shrinkage

将SNN划分为多个阶段，其中每个阶段的时间步长逐渐收缩降低。
初始阶段的时间步长稍大，以使得SNN能够提取到充分有价值的信息；后面的阶段时间步长较小，以降低整体的推理时延。
MLF方法基于LIF神经元，其主要思路为使用MLF unit（可以理解为新的脉冲神经元）替换掉SNN中的LIF神经元。

Temporal Transformer转换时间步长收缩前后两个阶段的信息维度

• 时间步长收缩需要解决的问题：假设前后两个阶段的时间步长分别为$T_1$和$T_2$，则两个阶段传递的数据维度分别为$T_1\times C\times H\times W$和$T_2\times C\times H\times W$。需要在时间步长收缩的同时转换前一个阶段的输出维度，以使其能够用作后一个阶段的输入。

• Temporal Transformer：设计了轻量化的Temporal Transformer来转换信息的维度，并最大化地保留有价值的信息。

  • 对于前一个阶段的输出${\mathbf{O}}_1\in {\mathbb{R}}^{T_1\times C\times H\times W}$先计算$T_1$个时间步长内每个时间步长的总信息${\mathbf{O}}_1^{avg}\in {\mathbb{R}}^{T_1\times 1}$：
    $${\mathbf{O}}_1^{avg}=\frac{1}{C\times H\times W}\sum _{i=1}^C\sum _{j=1}^H\sum _{k=1}^W{\mathbf{O}}_{1,i,j,k},$$
    然后使用非线性转换以及softmax计算$T_2$个时间步长的重要性分数$\mathbf{d}\in {\mathbb{R}}^{T_2\times 1}$：
    $$\mathbf{d}=\text{softmax}(\mathbf{W}{\mathbf{O}}_1^{avg}),$$
    其中$\mathbf{W}\in {\mathbb{R}}^{T_2\times T_1}$是非线性转换的可学习权重。softmax函数保证$T_2$个时间步长的重要性分数之和为1，确保后续的信息分配的完整性。
    另一方面，计算$O_1$在$T_1$个时间步长内的信息之和（所有时间步长的数据累加）${\mathbf{O}}_1^{total}\in {\mathbb{R}}^{C\times H\times W}$，并依据$d$将其分配至$T_2$个时间步长内，得到用于后续阶段的输入${\mathbf{I}}_2\in {\mathbb{R}}^{T_2\times C\times H\times W}$：
    $${\mathbf{I}}_{2,t}={\mathbf{O}}_1^{total}\odot {\mathbf{d}}_t=\sum _{t^{{}^{\prime }}=1}^{T_1}{\mathbf{O}}_{1,t^{{}^{\prime }}}\odot {\mathbf{d}}_t,$$
    从而在实现时间步长收缩的同时保留有效信息。

• 平均时间步长

  • 假设SNN被分为$n$个阶段，每个阶段有$n_i$个计算单元（一个卷积层以及脉冲神经元层），每个阶段的时间步长是$T_i$，则使用时间步长收缩得到的SNN的平均时间步长可计算为：
    $$T_{avg}=\frac{{\sum }_i^n{n}_i{T}_i}{{\sum }_i^n{n}_i}.$$
  • 由于上式不包含用于分类的全连接层，而全连接层以最小的时间步长$T_n$运行，因此实际的时间步长要比上式计算得到的结果更小。

• 额外开销

  • 每次时间步长收缩时的Temporal Transformer仅需要一个线性层来实现，所需的开销极小。

Early Classifier

• 基于代理梯度训练SNN遭受着代理梯度和真实梯度不匹配的问题，从而限制了SNN的性能。另外，梯度消失/爆炸问题始终影响着SNN的性能。这些问题对高性能SNN的训练造成了极大的困扰。
• 为了缓和以上问题，受到ANN领域部分方法的启发，提出在训练时候在SNN的每个阶段后添加一个early classifier，early classifier与标签计算损失并且在反向传播时传递梯度。
• 由于这些early classifier比SNN最终的分类器更靠近网络的前面层，因此这些early classifier传递的梯度受到梯度不匹配、梯度爆炸/消失的影响更小，从而促进了SNN的训练。
• 每个early classifier都由卷积层、脉冲神经元层和全连接层组成。进一步地，可以为每一个early classifier设置不同的结构以探索更优的性能增益或是设置一个全局共享的early classifier降低训练期间的开销。
• SNN训练的最终损失由多个early classifier和最终的输出与标签计算损失并加权得到：
  $${\mathcal{L}}_{total}=\sum _i^n{\lambda }_i{\mathcal{L}}_i(\frac{1}{T_i}\sum _t^{T_i}{\mathbf{Y}}_{i,t},\hat{\mathbf{Y}})$$
• 额外开销：early classifier仅在推理时引入了部分计算及参数开销，不影响推理时SNN的效率。

完整的SSNN训练算法

实验验证及对比

• 在CIFAR10-DVS、N-Caltech101和DVS-Gesture三个神经形态数据集上使用ResNet、VGG架构进行消融实验，平均时间步长设置为5，其性能比baseline大幅提升：
  
• 在消融实验部分探究了控制多个early classifier的loss权重对性能的影响，实验结果表明只要权重在合理的范围内，SSNN对不同的权重组合并不敏感，始终具有较好的性能：
  * 在消融实验部分探究并验证了所提出的Temporal Transformer的效果：
  
• 基于VGG-9，探究了SSNN划分阶段数量和每个阶段的时间步长对性能的影响，实验结果表明SSNN对阶段划分数量和每个阶段的时间步长设置并不敏感：
  
• 在不同的平均时间步长下，SSNN始终比普通的SNN具有更好的效果：
  
• 与现有的SNN方法或架构对比，SSNN基于VGG-9、ResNet架构都取得了更好的效果，甚至超出了Spikformer：
  
• SSNN与普通的SNN进行脉冲发射率的可视化对比，SSNN能够精确地关注到与识别最相关的区域（DVS-Gesture中的手势部分）：
  

本文由CSDN-lan人啊原创，转载请注明！