技術文章 » 文獻精讀1：SpikTransformer

文獻精讀1：SpikTransformer

2023-06-27 15:01:13

Spikformer

code source(pku)：GitHub - ZK-Zhou/spikformer: ICLR 2023, Spikformer: When Spiking Neural Network Meets Transformer

摘要

本文結合了兩種生物學上的合理結構（biologically plausible structures），尖峰神經網路(spiking neural network,snn)和自注意力機制（self-attention mechanism）。

  • SNN：提供節能和事件驅動（event-driven）的正規化
  • AT:捕獲特徵依賴關係，實現更好的效能

基於兩種特性的結合，提出了尖峰自注意力（Spiking Self Attention,SSA）為基礎的框架即尖峰神經網路（Spiking Transformer，Spikformer），通過使用spike-form查詢/鍵/值（代替softmax）來對稀疏視覺特徵進行建模。由於計算係數，避免了乘法，SSA是高效的，具有低計算能耗。

實驗結果表明，Spikformer在神經形態和靜態資料集上的影象分類由於先進的SNNs-like框架。

名詞解釋

SOPs:突觸運算（synaptic operations）

原始自注意力機制和尖峰注意力機制

介紹

尖峰自注意力（Spiking Self Attention,SSA）機制通過引入尖峰序列建模相互依賴性（interdependence）。

在SSA中，尖峰自注意力機制主要有以下特點：

  • 輸入和值均為二值化形式，僅包含0和1（二進位制）。與VSA的浮點輸入和值相比包含較少的細粒度特徵，浮點QKV對尖峰序列的建模是冗餘的。
  • 非負特性。解耦了SOFTMAX的影響。【softmax的作用可能僅僅是保證非負的意義，先前的TRM變體如Performer採用隨機正特徵來逼近softmax，cosformer使用RELU和cos-func代替了softmax。】

基於上述特性，可以使用與門（AND）和加法器來實現乘法。Spikformer的架構如下圖所示，其提高了在靜態資料集和神經形態資料集上訓練的效能。這是首次探索SNN中的自注意力機制和直接訓練的TRM，本文的貢獻體現如下方面：

  • 我們設計了一個新的尖峰形式的自我注意命名為尖峰自我注意（SSA）的SNNs的屬性。使用稀疏尖峰形式的QKV而不使用softmax，SSA的計算避免了乘法運算，效率很高。
  • 基於SSA，提出了尖峰TRANSFORMER，並且使用直接訓練的SNN模型在ImageNet上以4個時間步長實現了超過74%的準確率。所提出的架構優於靜態和神經形態資料集上最先進的SNN。
Spiformer架構圖

【架構中包括了Spiking patchspliting（SPS）模組，Spikformer編碼器和線性分類頭（Linear classification head）組成。LayerNorm(LN)不適合用於SNN，採用了BN代替】

方法

相關工作

SNN:

區別於傳統使用連續十進位制資料(continuous decimal values)傳遞資訊的深度學習模型，SNN使用離散尖峰序列來計算和傳輸資訊。尖峰神經網路可以接收連續值(continuous values)，並將其轉換為尖峰序列（spike sequences）。【相關工作有Leaky Integrate-and-Fire,LIF neuron和PLIF】

有兩種方法可以獲得深度SNN模型：ANN-TO-SNN轉換和直接訓練（direct training）。

  • ANN-TO-SNN：可以通過使用尖峰神經元（sping neurous）替換RELU啟用層，將高效能預訓練的ANN轉換為SNN。轉換後的SNN需要很大的時間步來準確近似RELU啟用，這會造成很大的延遲。
  • 直接訓練：SNN在模擬時間步長上展開（unfold），並以時間反向傳播的方式訓練？【尖峰神經元的事件觸發機制是不可微的，代理梯度（surrogate gradient）可用於反向傳播、並採用隱式微分的平衡狀態(implicit differentiation on the equailibrium state)來訓練SNN】

目前人工神經網路的各類模型已被移植到SNN上，但現前自注意力在SNN上的研究仍為空白。有以下嘗試性工作：

  • 使用時間注意力（temporal）來減少冗餘的時間步長。
  • 使用ANN-TRM來處理尖峰資料（spike data），儘管方法標題中提到了spiking trm。
  • ANN-SNN轉換TRM，仍然保持了原始的自注意力）—未證實SNN的特性。

SNN基本單元，即尖峰神經元（spike neuron）會接收所得到的電流並累計模電位（membrane potential），該膜電位用於和閾值比較以確定是否生成尖峰。本文使用了LIF尖峰神經元，描述如下：

\[\begin{aligned} &H[t]=V[t-1]+\frac{1}{\tau}\left(X[t]-\left(V[t-1]-V_{r e s e t}\right)\right), \\ &S[t]=\Theta(H[t]-V_{t h}), \\ &V[t]=H[t](1-S[t])+V_{r e s e t}S[t], \end{aligned} \]

其中，\(\tau\)是膜時間常數（membrane time constant），\(X[t]\)為時間步\(t\)的輸入電流。當膜電位\(H(t)\)超過激發閾值（firing threshold，\(V_{th}\)），尖峰神經元將觸發尖峰\(S[t]\)。\(\Theta(v)\)是Heaviside階躍函數（\(v\ge0\)時為1，否則為0）；\(V[t]\)表示觸發事件之後的膜電位，如果不產生尖峰，則其等於\(H[t]\)，否則會等於復位電位\(V_{reset}\)。

架構解釋

總體框架

給定二維影象序列\(I\in \mathbb R^{T\times C\times H\times W}\)，Spiking Patch Spliting(SPS)模組將其線性投影到一個D維尖峰狀特徵向量上，並將其分割為\(N\)個展開的尖峰形狀塊（flattened spike-form patches）\(x\)。

浮點形式的位置嵌入（position embedding）不可用於SNN中，我們採用了條件位置嵌入生成器（conditional position embedding generator）來生成尖峰形式的相對位置嵌入（relative position embedding,RPE），並將RPE新增到塊序列\(x\)來得到\(X_0\)。條件位置編碼生成器包含核心大小為\(3\)的二維折積層（Conv2D）、批歸一化層（BN）和尖峰神經元層（SN）。

將\(X_0\)傳遞到\(L\)塊Spikformer編碼器中，該編碼器由尖峰自注意力（SSA）和MLP塊組成，在SSA和MLP中應用了殘差連線。

作為Spikformer編碼器塊中的主要元件，SSA提供了一種有效的方法來使用Spike-form的QKV，對影象的區域性-全域性資訊進行建模，而無需softmax。全域性平均池化（global average pooling）被用於從Spikformer編碼器處理的特徵中輸出D維特徵，該D維特徵會被送入全連線層分類頭（classification head,CH）來輸出預測值Y。

\[\begin{aligned} &x=\mathrm{SPS}\left(I\right),&& I\in\mathbb{R}^{T\times C\times H\times W},x\in\mathbb{R}^{T\times N\times D}, \\ &\mathrm{RPE}={\mathcal{S N}}(\mathrm{BN}((\mathrm{Conv2d}(x)))),&& \mathrm{RPE}\in\mathbb{R}^{T\times N\times D} \\ &X_{0}=x+\mathrm{RPE},&& X_0 \in \mathbb{R}^{T \times N \times D} \\ & \\ &X_l'=\mathrm{SSA}(X_{l-1}) + X_{l-1},&& X_l^{\prime}\in\mathbb{R}^{T\times N\times D},l=1...L \\ &X_{l}=\mathrm{MLP}(X_{l}^{\prime}) + X_{l}^{\prime},&& X_l\in\mathbb{R}^{T\times N\times D},l=1...L \\ &Y=\mathrm{CH}(\mathrm{GAP}(X_{L})) \end{aligned} \]

SPS

尖峰塊分離模組（SPS）用於將影象線性投影到\(D\)維度的尖峰形式特徵（spike-form feature），並將特徵拆分為固定大小的塊（patches）。與Vision TRM中的折積骨幹（convolution stem）類似，本文在每個SPS塊中應用了折積層來引入歸納偏置到Spikformer中。給定影象序列\(I\in \mathbb R ^{T\times C\times H\times W}\)：

\[x=\mathcal{MP}\left(\mathcal{SN}(\mathrm{BN}((\mathrm{Conv2d}(I)))))\right) \]

其中，Conv2D和MP表示步長1、核大小3的二維折積層和最大池化，SPS的數目可以大於1。當使用多個SPS塊時，折積層的輸出通道數目增加，並最終匹配塊的嵌入維度（embedding dimension）。給定嵌入維度D和四塊SPS模組，四個折積層中的輸出通道數目為\(D/8,D/4,D/2,D\)。二維最大池化層被應用於在具有固定大小的SPS塊後對特徵大小進行下取樣。在SPS處理後，I被分割為影象塊序列\(x\in \mathbb R^{T\times N \times D}\)。

尖峰自注意力機制

Vanllia Self-attention難以應用於SNNs中，主要存在以下兩個原因：浮點矩陣乘法\(Q_F,K_F\)和softmax函數（包含指數計算和除法運算，不符合SNN的計算規則）；VSA序列長度的二次空間複雜度和時間複雜度不滿足SNN的高效計算要求。

首先通過可學習矩陣計算查詢、鍵和值，然後通過不同的尖峰神經元層來成為尖峰序列（spikin sequences）：

\[Q=\mathcal{SN}_Q(\mathrm{BN}(XW_Q)),K=\mathcal{SN}_K(\mathrm{BN}(XW_K)),V=\mathcal{SN}_V(\mathrm{BN}(XW_V)) \]

其中，\(Q,K,V\in \mathbb R^{T\times N \times D}\)，本文認為注意力矩陣的計算過程應使用純尖峰形式的Q和K(只包含0和1)。受原始自注意力啟發，我們加入了縮放因子\(s\)來控制矩陣乘法結果的大值，\(s\)不影響SSA的屬性。spike-friendly SSA定義如下：

\[\begin{aligned} &\mathrm{SSA}^{'}(Q,K,V)=\mathcal{S N}\left(QK^{\mathrm{T}}V*s\right) \\ &\mathrm{SSA}(Q,K,V)={\mathcal{S N}}(\mathrm{BN}(\mathrm{Linear}(\mathrm{SSA}^{'}(Q,K,V)))). \end{aligned} \]

在上式由尖峰神經元輸出的\(Q,K,V\)為非負的，因此會生成非負注意力圖。SSA只聚合這些相關的特徵，而忽略不相干的資訊。因此，SSA不需要通過softmax來保證注意力圖的非負性。此外，SSN中輸入\(X\)和自注意力值\(V\)為尖峰形式，包含有限資訊；浮點形式的\(QK\)和注意力softmax對於建模尖峰形式的\(X,V\)時冗餘的，不能從\(X,V\)中獲得更多資訊，即SSA比VSA更適合於SNN。

實驗

實驗資料集：使用靜態資料集CIFAR,ImageNet和神經形態資料集（neuromorphic datasets）如CIFAR10-DVS，DVS128 Gesture來評估Spikformer。

靜態資料集

ImageNet：使用130萬張1000類的影象用於訓練，5萬張影象用於驗證。我們在ImageNet上的模型的輸入大小被設定為預設的224×224。優化器是AdamW，在310個訓練階段中，批次大小設定為128或256，餘弦衰減學習率的初始值為0.0005。在ImageNet和CIFAR上訓練時，縮放因子為0.125。四塊SPS將影象分割為196個16×16的塊。【訓練實驗中使用了標準資料增廣方法，如隨機增強，混合和剪下混合等。

CIFAR：提供50，000張訓練影象和10，000張測試影象，解析度為32×32。批次大小設定為128。四塊SPS（前兩個塊不包含最大池化層）將影象分割為64個4 × 4塊。