[論文レビュー] Tensor Decomposition Based Attention Module for Spiking Neural Networks

Haoyu Deng, Ruijie Zhu+4 more|arXiv (Cornell University)|Oct 23, 2023

Advanced Memory and Neural Computing被引用数 4

ひとこと要約

本稿では、テンソル分解を用いて、パrameterが線形に増加する低ランクで三次元（時間的・チャネル的・空間的）の注意機構を実現する、効率的なスパikingニューラルネットワーク（SNN）向けの「プロジェクティドフルアテンション（PFA）」モジュールを提案する。CP分解とその逆プロセスである線形射影とアテンションマップの合成を活用することで、計算コストを最小限に抑えつつ、静的および動的SNNベンチマークで最先端の性能を達成する。

ABSTRACT

The attention mechanism has been proven to be an effective way to improve spiking neural network (SNN). However, based on the fact that the current SNN input data flow is split into tensors to process on GPUs, none of the previous works consider the properties of tensors to implement an attention module. This inspires us to rethink current SNN from the perspective of tensor-relevant theories. Using tensor decomposition, we design the extit{projected full attention} (PFA) module, which demonstrates excellent results with linearly growing parameters. Specifically, PFA is composed by the extit{linear projection of spike tensor} (LPST) module and extit{attention map composing} (AMC) module. In LPST, we start by compressing the original spike tensor into three projected tensors using a single property-preserving strategy with learnable parameters for each dimension. Then, in AMC, we exploit the inverse procedure of the tensor decomposition process to combine the three tensors into the attention map using a so-called connecting factor. To validate the effectiveness of the proposed PFA module, we integrate it into the widely used VGG and ResNet architectures for classification tasks. Our method achieves state-of-the-art performance on both static and dynamic benchmark datasets, surpassing the existing SNN models with Transformer-based and CNN-based backbones.

研究の動機と目的

• 既存のSNNでは、イベントデータをテンソルとして静的処理するが、テンソル構造を活用しないことが一般的であるため、そのようなテンソルに意識的なアテンション機構の欠如に対処すること。
• 時間的・チャネル的・空間的次元の全アテンションを可能にしつつ、パrameterの増加を抑えた軽量なアテンションモジュールを設計すること。
• 特にCP分解とその逆プロセスであるテンソル分解理論を活用し、SNNにおける効率的なアテンションマップ生成を実現すること。
• 神経形態的コンピューティングおよびGPU加速フレームワークと互換性がある理論的根拠に基づいた、解釈可能でスケーラブルなアテンション機構を提供すること。

提案手法

• 入力スパイクテンソルを学習可能な次元特化線形射影を用いて、3つの低次元のプロジェクティドテンソルに圧縮するための「スパイクテンソルの線形射影（LPST）」モジュールを提案する。
• CP分解プロセスの逆操作を用いて接続要因を用いてアテンションマップを再構成する「アテンションマップ合成（AMC）」モジュールを導入する。
• テンソルランク解析とデータセット固有の性質に基づくランク選択基準を採用し、最適な分解ランク $ r $ を決定することで、性能と効率の両立を実現する。
• 洗練されたアテンションマップと元の入力テンソルをヘダムアルド積により融合することで、動的特徴再重み付けを可能にする。
• VGGおよびResNetアーキテクチャにPFAモジュールを組み込み、静的（CIFAR10/100）および動的（CIFAR10DVS）データセットでエンドツーエンド学習を実行する。
• バックプロパゲーションと互換性のある微分可能フレームワークを採用し、アテンションを備えたSNNのエンドツーエンド学習を可能にする。

Figure 3: The detailed workflow of PFA. The input tensor is first sent to Linear Projection of Spike Tensor (LPST) unit to generate three projections and split the three projections into corresponding vectors. In Attention Map Composing (AMC) unit, these vectors are composed into the final attention

実験結果

リサーチクエスチョン

• RQ1 テンソル分解を効果的に活用することで、スパキングニューラルネットワークにおける軽量でパrameter効率の良いアテンション機構を設計できるか？
• RQ2 入力テンソルのランクと接続要因 $ r $ の選択が、SNNにおけるアテンションマップの品質とモデル性能に与える影響はいかほどか？
• RQ3 時間的・チャネル的・空間的アテンションの各成分が、SNNにおける分類精度にどの程度寄与しているか？
• RQ4 テンソルベースのアテンション機構は、静的および動的ベンチマークデータセットにおいて、既存のTransformerおよびCNNベースのSNNモデルを上回る性能を発揮できるか？

主な発見

• PFAはCIFAR10（94.85％の正確度）およびCIFAR100（78.65％の正確度）で最先端の性能を達成し、TransformerおよびCNNベースのバックボーンを用いた既存のSNNモデルを上回った。
• アブレーションスタディの結果、時間的・チャネル的・空間的次元のいずれのアテンションをも削除しても性能が低下し、特に空間的アテンションの欠落が最も顕著な影響（79.60％に低下）を及ぼした。
• 単一のPFAレイヤーを挿入するだけでもネットワークの正確度が向上し、PFAレイヤーを増やすことで性能が単調に向上する傾向を示しており、スケーラビリティと一貫した利点が確認された。
• 疎なSNNにおいてもアテンションマップのランクは低く保たれていることから、疎表現における未だ解き放たれていない最適化の余地が示唆された。
• 可視化により、PFAが3次元すべての次元で意味のあるアテンション分布を学習していることが確認され、空間的アテンションが最も顕著で、次に時間的およびチャネル的次元の順に影響が大きかった。
• テンソルランクと接続要因の分析に基づく提案されたランク選択基準は、最適な $ r $ を効果的に予測できており、実験的性能と整合的であった。

Figure 8: Attention visualization on CIFAR10DVS dataset. Temporal and channel parts are taken from the deep layer of VGGSNN. The spatial layer visualizations are extracted from the shallow layer of the network due to the inherent complexity of the information contained in the deeper layers, which ma