[論文レビュー] Spikformer V2: Join the High Accuracy Club on ImageNet with an SNN Ticket

Zhaokun Zhou, Kaiwei Che+6 more|arXiv (Cornell University)|Jan 4, 2024

Advanced Memory and Neural Computing被引用数 12

ひとこと要約

Spikformer V2 は Spiking Self-Attention と Spiking Convolutional Stem を導入し、SNN による自己教師付きプリトレーニングで ImageNet の精度を 80% 超に達成、16 層で 81.10% を含む SSL。

ABSTRACT

Spiking Neural Networks (SNNs), known for their biologically plausible architecture, face the challenge of limited performance. The self-attention mechanism, which is the cornerstone of the high-performance Transformer and also a biologically inspired structure, is absent in existing SNNs. To this end, we explore the potential of leveraging both self-attention capability and biological properties of SNNs, and propose a novel Spiking Self-Attention (SSA) and Spiking Transformer (Spikformer). The SSA mechanism eliminates the need for softmax and captures the sparse visual feature employing spike-based Query, Key, and Value. This sparse computation without multiplication makes SSA efficient and energy-saving. Further, we develop a Spiking Convolutional Stem (SCS) with supplementary convolutional layers to enhance the architecture of Spikformer. The Spikformer enhanced with the SCS is referred to as Spikformer V2. To train larger and deeper Spikformer V2, we introduce a pioneering exploration of Self-Supervised Learning (SSL) within the SNN. Specifically, we pre-train Spikformer V2 with masking and reconstruction style inspired by the mainstream self-supervised Transformer, and then finetune the Spikformer V2 on the image classification on ImageNet. Extensive experiments show that Spikformer V2 outperforms other previous surrogate training and ANN2SNN methods. An 8-layer Spikformer V2 achieves an accuracy of 80.38% using 4 time steps, and after SSL, a 172M 16-layer Spikformer V2 reaches an accuracy of 81.10% with just 1 time step. To the best of our knowledge, this is the first time that the SNN achieves 80+% accuracy on ImageNet. The code will be available at Spikformer V2.

研究の動機と目的

• Spiking Neural Networks (SNNs) で高性能な画像分類を動機づけ、実現する。
• 乗算を必要とせずエネルギー効率の高い Spiking Self-Attention (SSA) を導入する。
• アーキテクチャを強化し SSL pre-training を可能にする Spiking Convolutional Stem (SCS) を開発する。
• Self-Supervised Learning を活用してより大規模で深い Spikformer V2 を訓練する。
• 純粋な SNN に対して最先端の ImageNet 精度を実証する。

提案手法

• softmax ベースのアテンションを、スパイク形式の Q、K、V と疎で乗算不要な計算に置換する Spiking Self-Attention (SSA) を提案する。
• SPS を置換し、追加の畳み込み層を介して帰納的バイアスを注入する Spiking Convolutional Stem (SCS) を導入する。
• SCS とより深いアーキテクチャを組み込むことで Spikformer を Spikformer V2 に拡張し、トレーニングの安定性を監視する。
• 不対称な SNN-ANN デコーダを用いた Mask Autoencoder 風の Self-Supervised Learning (SSL) を適用して、大規模な Spikformer V2 モデルを事前学習する。
• SSL 後に ImageNet でファインチューニングして、限られた時間ステップで高精度を達成する。

実験結果

リサーチクエスチョン

• RQ1 ソフトマックスと乗算を使わずに、自己注意は SNNs に効果的に実装できるか？
• RQ2 畳み込み幹は Spikformer におけるスパイクベースの特徴抽出と ImageNet の性能を向上させるか？
• RQ3 SSL pre-training はより大規模で深い Spikformer モデルを可能にし、ImageNet で 80% の精度を超えるか？
• RQ4 ImageNet における Spikformer V2 の時間ステップ、モデルの深さ、精度のトレードオフは？
• RQ5 ANN-Decoder は pre-training のみの利点で、SNN の推論には不要ですか？

主な発見

• 8 層と 4 タイムステップを備えた Spikformer V2 は、監督付きトレーニングで ImageNet で 80.38% の精度を達成。
• 自己教師付き事前学習により 16 層の Spikformer V2 が 1 つの時間ステップのみで 81.10% の精度に達する。
• Spikformer V2 は、SNN の中で初めて ImageNet の精度 80% を超えた。
• SSA は softmax を排除しスパイク形 Q, K, V を活用することで計算を削減し、乗算なしの演算を実現する。
• 畳み込みブロックを備えた SCS は SPS より性能を向上させ、効果的な SSL 事前学習を可能にする。
• SSL 事前学習は大規模モデルの訓練を安定化させ、ImageNet での性能向上をもたらす。