[論文レビュー] Temporal Efficient Training of Spiking Neural Network via Gradient Re-weighting

Shikuang Deng, Yuhang Li+2 more|arXiv (Cornell University)|Feb 24, 2022

Advanced Memory and Neural Computing被引用数 102

ひとこと要約

tldr: 本論文は、Spiking Neural Networks (SNNs) の一般化性能と時間的スケーラビリティを向上させる勾配再重み付けを用いた Temporal Efficient Training (TET) を導入し、ニューロモルフィックデータと静的データセットの最先端結果を達成し、DVS-CIFAR10 において top-1 の 83.17% を含む。

ABSTRACT

Recently, brain-inspired spiking neuron networks (SNNs) have attracted widespread research interest because of their event-driven and energy-efficient characteristics. Still, it is difficult to efficiently train deep SNNs due to the non-differentiability of its activation function, which disables the typically used gradient descent approaches for traditional artificial neural networks (ANNs). Although the adoption of surrogate gradient (SG) formally allows for the back-propagation of losses, the discrete spiking mechanism actually differentiates the loss landscape of SNNs from that of ANNs, failing the surrogate gradient methods to achieve comparable accuracy as for ANNs. In this paper, we first analyze why the current direct training approach with surrogate gradient results in SNNs with poor generalizability. Then we introduce the temporal efficient training (TET) approach to compensate for the loss of momentum in the gradient descent with SG so that the training process can converge into flatter minima with better generalizability. Meanwhile, we demonstrate that TET improves the temporal scalability of SNN and induces a temporal inheritable training for acceleration. Our method consistently outperforms the SOTA on all reported mainstream datasets, including CIFAR-10/100 and ImageNet. Remarkably on DVS-CIFAR10, we obtained 83$\%$ top-1 accuracy, over 10$\%$ improvement compared to existing state of the art. Codes are available at \url{https://github.com/Gus-Lab/temporal_efficient_training}.

研究の動機と目的

• 直接訓練での surrogate gradients (SG) を用いた SNN の一般化の限界を動機づける。
• Gradient signal を各時間ステップの出力に合わせる Temporal Efficient Training (TET) を提案し、より flatter minima を促進する。
• TET が時間的スケーラビリティを高め、訓練時間を短縮する Time Inheritance Training (TIT) を可能にする。
• 静的データセット（CIFAR-10/100、ImageNet）とニューロモルフィックデータ（DVS-CIFAR10）で最先端の性能を示す。

提案手法

• SNN のために iterative Leaky Integrate-and-Fire (LIF) モデルと Spike/Membrane ダイナミクスを採用する。
• 標準の direct training loss を L_TET に置き換え、各時間 t におけるクロスエントロピーで per-time-step outputs を最適化する。
• per-time-step outputs を制約し訓練を安定化させる正則化項 L_MSE を組み込む。
• SDT と TET は異なる loss ランプを生み出すことを示し、TET は flatter minima に収束しやすい。
• Time Inheritance Training (TIT) を導入して短いシミュレーション長で訓練をブートストラップし、徐々に T を拡張する。
• 従来の SDT loss (L_SDT) を L_TET と Jensen/AM-GM 不等式を用いて上界する損失定式を提供する。
• 推論ルールは変更せず、訓練損失のみを変更する。

実験結果

リサーチクエスチョン

• RQ1 surrogate-gradient-based direct training (SDT) は TET と比較して SNN を鋭い局所極値に閉じ込めてしまうか？
• RQ2 per-time-step trainings losses (L_TET) は SNN の flatter minima と一般化をもたらすか？
• RQ3 Time Inheritance Training (TIT) は訓練時間を短縮しつつ精度を維持または向上させるか？
• RQ4 静的データセット（CIFAR-10/100, ImageNet）とニューロモルフィックデータ（DVS-CIFAR10）における TET の性能は、最先端手法と比較してどうか？

主な発見

• TET は複数のネットワークとデータセットで SDT に比べて精度を大幅に向上させる（例：CIFAR-100 の ResNet-19）。
• TET の下では局所極値での損失ランドスケープが SDT より flatter で、一般化が改善される。
• 学習中盤で SDT から TET に切替えると鋭い局所極値から抜け出し、テスト性能が向上する。
• Time Inheritance Training (TIT) は訓練時間を約半分に短縮し、特にシミュレーション長が大きくなるにつれて精度を維持・向上させる。
• DVS-CIFAR10 に VGGSNN を用い augmentation なしで、TET は top-1 精度 83.17% を達成し、従来の SOTA より 10 ポイント超を記録；augmentation を用いても 83.17% が報告される。