論文地址：

Spike-driven Transformer V2: Meta Spiking Neural Network Architecture Inspiring the Design of Next-generation Neuromorphic Chips

代碼地址：

??https://github.com/BICLab/Spike-Driven-Transformer-V2??

背景

在神經形態芯片上運行脈沖神經網絡（Spiking Neural Network，SNN）的神經形態計算，是類腦計算領域的主流發展方向。CNN-based SNN 是神經形態計算領域內的主流架構，幾乎所有的神經形態芯片都支持這一架構。

最近，Transformer-based SNN 架構在 SNN 領域內興起，并表現出了相當的競爭力，目前尚未有神經形態芯片專門針對這一新興架構設計。此外，在性能方面，現有 spiking Transformer 架構僅能夠與 state-of-the-art 的 spiking CNN 算法相持平，并未表現出明顯的優勢。

本文的目標是探索 SNN 的 meta 架構設計，以期在算法和硬件兩個層面上推動 SNN 領域的發展。在算法層面，展示 spiking Transformer 架構相對于 spiking CNN 在任務性能和通用性方面的優勢，縮小 SNN 和 ANN 之間的性能差距。

在硬件層面，spiking Transformer 的元架構設計有助于減少 Transformer-based SNN 神經形態芯片設計時所需要的算法探索成本，推動神經形態計算領域中下一代基于 spiking Transformer 的神經形態芯片的發展。

本文貢獻

本文將 Spike-driven Transformer [1] 擴展成一種元 SNN 架構，命名為 "Meta-SpikeFormer". 本文的主要貢獻包括：

• SNN 架構設計。本文設計了一種僅包含稀疏加法的 meta Transformer-based SNN 架構，主要包括了宏觀層面的 Conv-based 和 Transformer-base SNN 塊設計，以及微觀層面上的設計，包括幾種新的脈沖卷積方法，脈沖形式 Q，K，V 矩陣的生成方法，以及三種具有不同計算復雜度的脈沖驅動自注意力（Spike-Driven Self-Atttention，SDSA）算子等。
• 性能。所提出的 Meta-SpikeFormer 使得 SNN 領域首次在 ImageNet-1K 上達到 80%，比當前的 state-of-the-art 基線 Spike-driven Transformer [1] 性能提升了 3.7%，但參數量減少了 17%。
• 通用性。Meta-SpikeFormer 是首個可以同時處理分類、檢測、分割的直接訓練 SNN 架構。Meta-SpikeFormer 在 ImageNet-1K，HAR-DVS（目前最大的 DVS 動作分類數據集），COCO，ADE20K 等數據集上進行了實驗，并在所有的數據集上取得了 SNN 領域內的最佳性能。這也是 SNN 領域首次在 ADE20K 數據集上報道性能。
• 神經形態芯片設計。本文系統性地研究了 Transformer-based SNN 的通用組成，包括架構、殘差連接、脈沖驅動自注意力算子等。與 CNN-based SNN 相比，Meta-SpikeFormer 表現出明顯的性能和通用性優勢。這無疑將啟發和引導神經形態計算領域開發基于 Transformer 的神經形態芯片。

方法

3.1 ANN中Meta Transformer架構的概念

在 Vision Transformer [2] 展現出令人鼓舞的性能之后，對 ViT 的改進和討論逐成為架構設計的中流砥柱。各種改進層出不窮，一些典型的工作包括架構設計（PVT 、MLP-Mixer 等）、自注意力增強（Swin、Twins 等）、訓練優化（DeiT、T2T-ViT等)、高效 ViT 設計等。

在這項工作中，我們的目標是從種類繁多的 ViT 變體中參考和探索 meta spiking Transformer 架構設計，以快速縮小 SNN 和 ANN 在架構設計上的差距。

▲ 圖1 Meta Transformer Block.

自注意力機制被認為是 Transformer 架構取得成功的核心，其主要用于聚合不同位置 token 之間的信息（即，token mixer）。基于這一理解，研究者們擴展了許多 token mixer 的替代方案，比如 MLP，傅里葉變換等。其中，Yu 等人 [3,4] 認為，相較于關注某一個具體的 token mixer，一個通用的 meta Transformer 架構也非常重要。

如圖 1 所示，一個基礎的 meta Transformer 塊可以分為兩部分，一部分是 token mixer（混合空間信息，例如 self-attention），一部分是 channel mixer（混合通道信息，例如 channel MLP）。

3.2 Meta-SpikeFormer設計

基于 ANN 中 meta Transformer 架構的理解，本文將 spike-driven Transformer [1] 與 [4]中的 CAFormer 對齊。其中，CAFormer 由兩個 CNN 階段和兩個 Transformer 階段構成，這種前卷積/后自注意力的架構是計算機視覺中的經典架構，被認為有助加速網絡收斂和提升任務性能 [5]。

具體地，本文將 spike-driven Transformer 中的脈沖編碼層擴充為 4 個 Conv-based SNN 塊，并對 Transformer-based SNN 塊進行了重新設計，如圖 2 所示。

圖2 Meta-SpikeFormer架構設計

通過在 Conv-based SNN 中引入深度可分離卷積，利用重參數化生成脈沖形式 Q,K,V，設計三種具有不同計算復雜度的脈沖驅動自注意力算子（如圖 3）所示，Spike-driven Transformer 被全面升級。

▲ 圖3 具有不同計算復雜度的脈沖驅動自注意力算子

通過調節本文所提出的 Conv-based 和 Transformer-based SNN 塊的內部設計細節，可以實現脈沖神經網絡在任務精度、參數量、能耗等方面的 trade-off。圖 2 中的架構為本文推薦的架構，在 ImageNet 上具有最高的任務精度。進一步地，表 1 給出了不同設計條件下 Meta-SpikeFormer 的表現。研究者們可以根據自身的應用場景基于 Meta-SpikeFormer 對模型進行調整。

▲ 表1 Meta-SpikeFormer架構設計的消融實驗

實驗

Meta-SpikeFormer 在圖像分類，event-based 動作分類，目標檢測，語義分割任務上進行了測試。本文利用在 ImageNet 上訓練好的模型，在目標檢測和語義分割任務進行微調測試。這是直接訓練 SNN 首次能夠以統一的方式處理分類、檢測、分割這三種任務。

4.1 靜態圖像分類（ImageNet-1K）

結果如表 2 所示，所提出的 Meta-SpikeFormer 在 ImageNet-1K 上取得了 80% 的性能，是目前 SNN 領域內的最佳結果。

▲ 表2 Meta-SpikeFormer在ImageNet中的實驗結果

4.2 Event-based動作分類（HAR-DVS）

動態視覺場景（即，神經形態視覺）是 SNN 的典型應用場景之一。DVS 相機只在亮度發生變化時異步輸出稀疏事件流。SNN 也能夠以異步事件驅動的方式運行，因此天然能夠以低功耗、低時延的方式處理 DVS 數據。本文采用目前最大的 DVS 數據集 HAR-DVS 來測試所提出的 Meta-SpikeFormer，結果如表 3 所示。

▲ 表3 Meta-SpikeFormer在HAR-DVS中的實驗結果

4.3 目標檢測（COCO）

我們首先將 mmdetection 代碼庫轉換為 spiking mmdetection 版本。然后在 Meta-SpikeFormer 中加入 Mask R-CNN，在 ImageNet 上預訓練的模型被用來初始化骨干網絡。為對比公平，本文采取了預訓練+微調，直接訓練兩種方式來評估 Meta-SpikeFormer 在 COCO 上的性能，結果如表 4 所示。

▲ 表4 Meta-SpikeFormer在COCO中的實驗結果

4.4 語義分割（ADE20K）

本文報道了 SNN 在 ADE20K 上的首個結果。我們將 mmsegmentation 代碼庫轉為相應的 spiking 版本，在 ImageNet 上預訓練的模型被用來初始化骨干網絡。結果如表 5 所示。

▲ 表5 Meta-SpikeFormer在ADE20K中的實驗結果

總之，Meta-SpikeFormer 這些數據集上都取得了 SNN 領域的最佳結果。

結論

5.1 Meta-SpikeFormer在技術上如何啟發未來的神經形態芯片設計？

• Conv+ViT 混合架構設計。這種混合漸進式局部-全局建模可以同時利用 CNN 和 Transformer 的優勢，其中前者進行低層次特征建模，后者捕獲長距依賴關系。本文通過充分的實驗驗證了這種設計有利于 SNN 的性能和通用性。
• 脈沖驅動自注意力（SDSA）算子設計。SDSA 算子是 spiking Transformer 中進行長距離依賴建模的核心設計，這也是當前神經形態芯片所缺乏的設計。
• SNN 的 meta 架構設計。基于本文所設計的 meta Conv-based 和 Transformer-based SNN 模塊，研究人員可以根據自己在精度、參數和功耗方面的需求（表 1），對 meta SNN 模塊內部的設計細節進行針對性的優化。

5.2 Meta-SpikeFormer對神經形態芯片設計的意義

• 算法-硬件協同設計。大多數神經形態芯片設計都是從計算堆棧的底部開始的 [6]，也就是，從材料和設備的角度出發。Meta-SpikeFormer 所表現出的高精度和通用性可能會吸引和啟發算法驅動神經形態芯片設計。 
• 對大規模神經形態計算的信心。小規模神經形態計算已經在邊緣計算端展現出了明確的功耗和性能優勢 [7,8,9]。Meta-SpikeFormer 展示了大規模 SNN 在性能和通用性方面的潛力。
• 降低下一代神經形態芯片設計成本。Meta 設計有利于大幅縮小 SNN 和 ANN 在架構設計上的差距，降低算法驅動神經形態芯片設計之前所需的算法探索成本。

?本文轉載自 ???PaperWeekly????，作者：  李國齊課題組

原文鏈接：??https://mp.weixin.qq.com/s/ZP1v1ALH5IwOZ7g_Ve-3vQ??????