[論文レビュー] 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

ひとこと要約

本稿では、スパイikingニューラルネットワーク（SNNs）向けの新しいバッチ正則化手法である膜電位バッチ正則化（MPBN）を提案する。MPBNは、発火関数の直前に膜電位を正則化することで、畳み込み層におけるバッチ正則化後の残留膜電位更新によるデータフローの乱れを是正する。MPBNを学習可能な発火閾値に再パrameter化することにより、推論時のオーバーヘッドを排除するとともに、チャネル単位および要素単位の正則化を可能とし、最小限の時間ステップでCIFAR-10、CIFAR-100、ImageNetにおいて最先端の精度を達成する。

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のデータフローに与える乱れを是正すること。
• 膜電位段階に第2の正則化層を導入することで、深層SNNにおける学習安定性と収束性を向上させること。
• MPBNを発火閾値に再パラメータ化することで、追加の正則化に起因する推論時のオーバーヘッドを排除すること。
• 閾値畳み込みを活用することで、従来のCNNでは実現不可能だったSNNにおける要素単位の正則化を可能とすること。
• 静的および神経形状データセットにおいて、最小限の時間ステップで一貫した性能向上を示すこと。

提案手法

• スパイキングニューロンにおける膜電位更新（MPU）の直後で発火関数の直前に適用するMPBNというバッチ正則化レイヤーを提案する。
• 訓練・推論の分離を図る再パラメータ化技術を設計し、学習済みの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を適用する。
• 2段階の訓練プロセスを採用：バックプロパゲーション中は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 膜電位更新の直後に第2のバッチ正則化レイヤーを追加することで、SNNにおける学習安定性と収束性が向上するか？
• RQ2 MPBNの計算オーバーヘッドを性能を損なわず推論時に排除できるか？
• RQ3 閾値への再パラメータ化を踏まえた場合、SNNにおける要素単位の正則化は実現可能で有益であるか？
• RQ4 静的および動的データセットにおいて、MPBNは既存のBN手法と比較してどのように性能を発揮するか？
• RQ5 MPBNは、従来のSNNモデルと比較して、より少ない時間ステップで最先端の性能を達成できるか？

主な発見

• CIFAR-10では2時間ステップで96.47％のトップ-1精度を達成し、Real Spike（6時間ステップ）やRecDis-SNNを上回る。
• CIFAR-100では2時間ステップで79.51％のトップ-1精度を達成し、TEBN、TET、RecDis-SNNを1.66％～3.99％上回るが、それらはより多くの時間ステップを要する。
• 1時間ステップでMPBNはResNet19を用いてCIFAR-10で96.06％の精度を達成し、4～6時間ステップを要する他のSoTA手法を上回る。
• 要素単位MPBNのバリエーションは4時間ステップでCIFAR-10で94.42％の精度を達成し、チャネル単位バージョン（94.28％）をわずかに上回る。これは、より豊かな表現能力を示している。
• ImageNetでは強力な性能を発揮するが、非標準的でイベント駆動でないスパイクパターンを用いるSEW ResNetよりわずかに低い。
• 再パラメータ化技術により、MPBNが発火閾値にうまく統合され、推論時のコストが完全に排除された一方で、学習の利点は保持された。

(a) CIFAR-10