[论文解读] Membrane Potential Batch Normalization for Spiking Neural Networks

Yufei Guo, Yuhan Zhang+6 more|arXiv (Cornell University)|Aug 16, 2023

Advanced Memory and Neural Computing被引用 5

一句话总结

该论文提出膜电位批归一化（MPBN），一种用于脉冲神经网络（SNNs）的新批归一化技术，通过在放电函数之前对膜电位进行归一化，以缓解残余电位更新带来的干扰。通过将MPBN重新参数化为可学习的放电阈值，该方法消除了推理阶段的开销，同时支持通道维和元素维归一化，在CIFAR-10、CIFAR-100和ImageNet上以极少的时间步数实现了最先进（SoTA）的准确率。

ABSTRACT

As one of the energy-efficient alternatives of conventional neural networks (CNNs), spiking neural networks (SNNs) have gained more and more interest recently. To train the deep models, some effective batch normalization (BN) techniques are proposed in SNNs. All these BNs are suggested to be used after the convolution layer as usually doing in CNNs. However, the spiking neuron is much more complex with the spatio-temporal dynamics. The regulated data flow after the BN layer will be disturbed again by the membrane potential updating operation before the firing function, i.e., the nonlinear activation. Therefore, we advocate adding another BN layer before the firing function to normalize the membrane potential again, called MPBN. To eliminate the induced time cost of MPBN, we also propose a training-inference-decoupled re-parameterization technique to fold the trained MPBN into the firing threshold. With the re-parameterization technique, the MPBN will not introduce any extra time burden in the inference. Furthermore, the MPBN can also adopt the element-wised form, while these BNs after the convolution layer can only use the channel-wised form. Experimental results show that the proposed MPBN performs well on both popular non-spiking static and neuromorphic datasets. Our code is open-sourced at \href{https://github.com/yfguo91/MPBN}{MPBN}.

研究动机与目标

• 解决卷积层中批归一化后残余膜电位更新导致的SNN数据流干扰问题。
• 通过在膜电位阶段引入第二个归一化层，提升深度SNN的训练稳定性和收敛性。
• 通过将MPBN重参数化为放电阈值，消除额外归一化带来的推理时开销。
• 通过阈值折叠技术，实现SNN中此前在CNN中不可行的元素维归一化。
• 在静态和神经形态数据集上均实现一致的性能提升，且时间步数极少。

提出的方法

• 提出MPBN：一种在脉冲神经元的膜电位更新（MPU）之后、放电函数之前应用的批归一化层。
• 设计了一种训练-推理解耦的重参数化技术，将训练得到的MPBN统计量（均值、方差、缩放、偏移）折叠到放电阈值中。
• 推导出重参数化后的阈值公式：$ (\tilde{V}_{\text{th}})_{i,j,k} = \frac{(V_{\text{th}} - \beta_{i,j,k}) \sqrt{\sigma_{i,j,k}^2}}{\lambda_{i,j,k}} + \mu_{i,j,k} $，实现元素维阈值自适应。
• 支持通道维和元素维MPBN，克服了CNN中仅通道维BN可折叠到权重的局限性。
• 在CIFAR和ImageNet基准上，结合标准数据增强和使用余弦退火的SGD优化方法应用MPBN。
• 采用两阶段训练流程：反向传播中使用MPBN进行训练，随后重参数化以在推理中移除MPBN。

Figure 1 : The difference between our SNN with MPBN and the vanilla SNN. We add another BN layer after membrane potential updating (MPU) operation in the training. The MPBN can be folded into the firing threshold and then the homogenous firing threshold will be transformed into different ones.

实验结果

研究问题

• RQ1 在膜电位更新后添加第二个批归一化层，是否能提升SNN的训练稳定性和收敛性？
• RQ2 能否在不损失性能的前提下，消除MPBN在推理阶段的计算开销？
• RQ3 在重参数化为放电阈值的前提下，SNN中实现元素维归一化是否可行且有益？
• RQ4 在静态和动态数据集上，MPBN与现有SNN批归一化方法相比表现如何？
• RQ5 MPBN能否在少于先前SNN模型的时间步数下实现最先进性能？

主要发现

• 在CIFAR-10上，MPBN仅用2个时间步即达到96.47%的top-1准确率，优于Real Spike（6个时间步）和RecDis-SNN。
• 在CIFAR-100上，MPBN使用2个时间步达到79.51%的top-1准确率，比TEBN、TET和RecDis-SNN高出1.66%至3.99%，即使后者使用更多时间步。
• 在1个时间步下，MPBN使用ResNet19在CIFAR-10上达到96.06%的准确率，优于其他SOTA方法（这些方法需4–6个时间步）。
• 元素维MPBN变体在CIFAR-10上使用4个时间步达到94.42%的准确率，略高于通道维版本（94.28%），表明其具有更强的表征能力。
• 在ImageNet上，MPBN表现强劲，尽管略低于SEW ResNet，后者使用非标准、非事件驱动的脉冲模式，牺牲了SNN的优势。
• 重参数化技术成功将MPBN折叠到放电阈值中，完全消除了推理阶段的开销，同时保留了训练阶段的优势。

(a) CIFAR-10