論文地址：

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??

背景

在神經(jīng)形態(tài)芯片上運(yùn)行脈沖神經(jīng)網(wǎng)絡(luò)（Spiking Neural Network，SNN）的神經(jīng)形態(tài)計(jì)算，是類(lèi)腦計(jì)算領(lǐng)域的主流發(fā)展方向。CNN-based SNN 是神經(jīng)形態(tài)計(jì)算領(lǐng)域內(nèi)的主流架構(gòu)，幾乎所有的神經(jīng)形態(tài)芯片都支持這一架構(gòu)。

最近，Transformer-based SNN 架構(gòu)在 SNN 領(lǐng)域內(nèi)興起，并表現(xiàn)出了相當(dāng)?shù)母?jìng)爭(zhēng)力，目前尚未有神經(jīng)形態(tài)芯片專(zhuān)門(mén)針對(duì)這一新興架構(gòu)設(shè)計(jì)。此外，在性能方面，現(xiàn)有 spiking Transformer 架構(gòu)僅能夠與 state-of-the-art 的 spiking CNN 算法相持平，并未表現(xiàn)出明顯的優(yōu)勢(shì)。

本文的目標(biāo)是探索 SNN 的 meta 架構(gòu)設(shè)計(jì)，以期在算法和硬件兩個(gè)層面上推動(dòng) SNN 領(lǐng)域的發(fā)展。在算法層面，展示 spiking Transformer 架構(gòu)相對(duì)于 spiking CNN 在任務(wù)性能和通用性方面的優(yōu)勢(shì)，縮小 SNN 和 ANN 之間的性能差距。

在硬件層面，spiking Transformer 的元架構(gòu)設(shè)計(jì)有助于減少 Transformer-based SNN 神經(jīng)形態(tài)芯片設(shè)計(jì)時(shí)所需要的算法探索成本，推動(dòng)神經(jīng)形態(tài)計(jì)算領(lǐng)域中下一代基于 spiking Transformer 的神經(jīng)形態(tài)芯片的發(fā)展。

本文貢獻(xiàn)

本文將 Spike-driven Transformer [1] 擴(kuò)展成一種元 SNN 架構(gòu)，命名為 "Meta-SpikeFormer". 本文的主要貢獻(xiàn)包括：

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

方法

3.1 ANN中Meta Transformer架構(gòu)的概念

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

在這項(xiàng)工作中，我們的目標(biāo)是從種類(lèi)繁多的 ViT 變體中參考和探索 meta spiking Transformer 架構(gòu)設(shè)計(jì)，以快速縮小 SNN 和 ANN 在架構(gòu)設(shè)計(jì)上的差距。

▲ 圖1 Meta Transformer Block.

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

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

3.2 Meta-SpikeFormer設(shè)計(jì)

基于 ANN 中 meta Transformer 架構(gòu)的理解，本文將 spike-driven Transformer [1] 與 [4]中的 CAFormer 對(duì)齊。其中，CAFormer 由兩個(gè) CNN 階段和兩個(gè) Transformer 階段構(gòu)成，這種前卷積/后自注意力的架構(gòu)是計(jì)算機(jī)視覺(jué)中的經(jīng)典架構(gòu)，被認(rèn)為有助加速網(wǎng)絡(luò)收斂和提升任務(wù)性能 [5]。

具體地，本文將 spike-driven Transformer 中的脈沖編碼層擴(kuò)充為 4 個(gè) Conv-based SNN 塊，并對(duì) Transformer-based SNN 塊進(jìn)行了重新設(shè)計(jì)，如圖 2 所示。

圖2 Meta-SpikeFormer架構(gòu)設(shè)計(jì)

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

▲ 圖3 具有不同計(jì)算復(fù)雜度的脈沖驅(qū)動(dòng)自注意力算子

通過(guò)調(diào)節(jié)本文所提出的 Conv-based 和 Transformer-based SNN 塊的內(nèi)部設(shè)計(jì)細(xì)節(jié)，可以實(shí)現(xiàn)脈沖神經(jīng)網(wǎng)絡(luò)在任務(wù)精度、參數(shù)量、能耗等方面的 trade-off。圖 2 中的架構(gòu)為本文推薦的架構(gòu)，在 ImageNet 上具有最高的任務(wù)精度。進(jìn)一步地，表 1 給出了不同設(shè)計(jì)條件下 Meta-SpikeFormer 的表現(xiàn)。研究者們可以根據(jù)自身的應(yīng)用場(chǎng)景基于 Meta-SpikeFormer 對(duì)模型進(jìn)行調(diào)整。

▲ 表1 Meta-SpikeFormer架構(gòu)設(shè)計(jì)的消融實(shí)驗(yàn)

實(shí)驗(yàn)

Meta-SpikeFormer 在圖像分類(lèi)，event-based 動(dòng)作分類(lèi)，目標(biāo)檢測(cè)，語(yǔ)義分割任務(wù)上進(jìn)行了測(cè)試。本文利用在 ImageNet 上訓(xùn)練好的模型，在目標(biāo)檢測(cè)和語(yǔ)義分割任務(wù)進(jìn)行微調(diào)測(cè)試。這是直接訓(xùn)練 SNN 首次能夠以統(tǒng)一的方式處理分類(lèi)、檢測(cè)、分割這三種任務(wù)。

4.1 靜態(tài)圖像分類(lèi)（ImageNet-1K）

結(jié)果如表 2 所示，所提出的 Meta-SpikeFormer 在 ImageNet-1K 上取得了 80% 的性能，是目前 SNN 領(lǐng)域內(nèi)的最佳結(jié)果。

▲ 表2 Meta-SpikeFormer在ImageNet中的實(shí)驗(yàn)結(jié)果

4.2 Event-based動(dòng)作分類(lèi)（HAR-DVS）

動(dòng)態(tài)視覺(jué)場(chǎng)景（即，神經(jīng)形態(tài)視覺(jué)）是 SNN 的典型應(yīng)用場(chǎng)景之一。DVS 相機(jī)只在亮度發(fā)生變化時(shí)異步輸出稀疏事件流。SNN 也能夠以異步事件驅(qū)動(dòng)的方式運(yùn)行，因此天然能夠以低功耗、低時(shí)延的方式處理 DVS 數(shù)據(jù)。本文采用目前最大的 DVS 數(shù)據(jù)集 HAR-DVS 來(lái)測(cè)試所提出的 Meta-SpikeFormer，結(jié)果如表 3 所示。

▲ 表3 Meta-SpikeFormer在HAR-DVS中的實(shí)驗(yàn)結(jié)果

4.3 目標(biāo)檢測(cè)（COCO）

我們首先將 mmdetection 代碼庫(kù)轉(zhuǎn)換為 spiking mmdetection 版本。然后在 Meta-SpikeFormer 中加入 Mask R-CNN，在 ImageNet 上預(yù)訓(xùn)練的模型被用來(lái)初始化骨干網(wǎng)絡(luò)。為對(duì)比公平，本文采取了預(yù)訓(xùn)練+微調(diào)，直接訓(xùn)練兩種方式來(lái)評(píng)估 Meta-SpikeFormer 在 COCO 上的性能，結(jié)果如表 4 所示。

▲ 表4 Meta-SpikeFormer在COCO中的實(shí)驗(yàn)結(jié)果

4.4 語(yǔ)義分割（ADE20K）

本文報(bào)道了 SNN 在 ADE20K 上的首個(gè)結(jié)果。我們將 mmsegmentation 代碼庫(kù)轉(zhuǎn)為相應(yīng)的 spiking 版本，在 ImageNet 上預(yù)訓(xùn)練的模型被用來(lái)初始化骨干網(wǎng)絡(luò)。結(jié)果如表 5 所示。

▲ 表5 Meta-SpikeFormer在ADE20K中的實(shí)驗(yàn)結(jié)果

總之，Meta-SpikeFormer 這些數(shù)據(jù)集上都取得了 SNN 領(lǐng)域的最佳結(jié)果。

結(jié)論

5.1 Meta-SpikeFormer在技術(shù)上如何啟發(fā)未來(lái)的神經(jīng)形態(tài)芯片設(shè)計(jì)？

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

5.2 Meta-SpikeFormer對(duì)神經(jīng)形態(tài)芯片設(shè)計(jì)的意義

• 算法-硬件協(xié)同設(shè)計(jì)。大多數(shù)神經(jīng)形態(tài)芯片設(shè)計(jì)都是從計(jì)算堆棧的底部開(kāi)始的 [6]，也就是，從材料和設(shè)備的角度出發(fā)。Meta-SpikeFormer 所表現(xiàn)出的高精度和通用性可能會(huì)吸引和啟發(fā)算法驅(qū)動(dòng)神經(jīng)形態(tài)芯片設(shè)計(jì)。 
• 對(duì)大規(guī)模神經(jīng)形態(tài)計(jì)算的信心。小規(guī)模神經(jīng)形態(tài)計(jì)算已經(jīng)在邊緣計(jì)算端展現(xiàn)出了明確的功耗和性能優(yōu)勢(shì) [7,8,9]。Meta-SpikeFormer 展示了大規(guī)模 SNN 在性能和通用性方面的潛力。
• 降低下一代神經(jīng)形態(tài)芯片設(shè)計(jì)成本。Meta 設(shè)計(jì)有利于大幅縮小 SNN 和 ANN 在架構(gòu)設(shè)計(jì)上的差距，降低算法驅(qū)動(dòng)神經(jīng)形態(tài)芯片設(shè)計(jì)之前所需的算法探索成本。

?本文轉(zhuǎn)載自 ???PaperWeekly????，作者：  李國(guó)齊課題組

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