TPAMI2025頂刊！基於脈衝驅動Transformer的高效且可擴充套件訓練方法

©作者 | 中科院自動化所李國齊課題組

脈衝神經網路（SNN）因其生物合理性和低功耗特性，被視為人工神經網路（ANN）的潛在替代方案。然而，SNN 在實際應用中面臨效能差距和訓練成本高的挑戰。

為此我們提出了一種脈衝發放近似（SFA）方法，透過整數訓練和脈衝驅動推理最佳化脈衝神經元的發放模式。此外，我們開發了高效的脈衝驅動 Transformer 架構和脈衝掩碼自動編碼器，以防止 SNN 在模型引數量擴充套件過程中效能下降。

在 ImageNet-1k 資料集上，我們的模型（10M、19M、83M、173M 引數）分別實現了 78.5%、79.8%、84.0% 和 86.2% 的準確率。例如，10M 引數模型比現有最佳 SNN 準確率高 7.2%，訓練速度提升 4.5 倍，推理能效提高 3.9 倍。

我們在目標檢測、語義分割和神經形態視覺等任務中驗證了方法的有效性和高效性。這項工作使 SNN 在保持低功耗的同時達到與 ANN 相當的效能，標誌著 SNN 向通用視覺主幹網路邁出重要一步。

論文標題：

Scaling Spike-driven Transformer with Efficient Spike Firing Approximation Training

論文連結：

https://arxiv.org/pdf/2411.16061

程式碼連結：

https://github.com/BICLab/Spike-Driven-Transformer-V3

背景

脈衝神經網路（SNNs）因其受大腦啟發的時空動態特性和高效的脈衝驅動計算正規化而備受關注。脈衝神經元遵循“積分–發放”模型，透過二進位制脈衝進行通訊，將空間和時間資訊整合到膜電位中。當膜電位超過閾值時，神經元發出脈衝並重置膜電位，僅在接收到脈衝時觸發稀疏加法運算。

當部署 SNN 到神經形態晶片（如 Speck 晶片 [1]）上時，其會展現出顯著的低功耗優勢，靜息功耗低至 0.42mW 同時在典型邊緣視覺場景中功耗僅為 0.7mW。

在深度學習時代，模型規模對現代機器智慧的成功至關重要，脈衝神經網路（SNN）也在不斷擴充套件規模。目前，大規模 SNN 的訓練方法主要有兩種：ANN2SNN [2] 和直接訓練 [3][4]。

ANN2SNN 透過脈衝發放率近似人工神經網路的啟用值，實現了較高的準確率，但需要大量推理時間步長，且放棄了 SNN 的時空動態特性。

直接訓練則利用時空反向傳播（STBP）[3] 和替代梯度 [4] 進行訓練，保留了 SNN 的時空動態特性，但高時間步長模擬需要大量計算資源。例如，對 10 個時間步長的 ResNet-19 進行直接訓練，所需記憶體是傳統 ResNet-19 的 20 倍 [5]。

一些文章 [2] 透過延長時間步長可以減少量化誤差，從而提升效能，但高時間步長的模擬需要大量 GPU 記憶體，增加了訓練難度，限制了 SNN 的可擴充套件性和效能提升。

現有觀點認為，量化誤差是主要問題，但我們進一步證明，二進位制脈衝發放是根本機制缺陷，削弱了脈衝神經元評估輸入訊號重要性的能力（空間表徵）和記憶與遺忘能力（時間動態特性）。現有方法（如 ANN2SNN 和直接訓練）透過速率編碼解決空間表徵問題，但訓練效率低下，且無法克服時間動態特性的限制。

為解決上述問題，我們提出了脈衝發放近似（SFA）方法。SFA 在訓練期間使用整數啟用值，推理期間將其轉換為脈衝序列，形成非同步連續的脈衝發放模式。與 ANN2SNN 和直接訓練的同步隨機發放不同，SFA 的發放模式顯著提升了訓練效率、推理功耗、效能和網路擴充套件能力。

為進一步提升 SNN 效能，我們改進了 Meta-SpikeFormer [6] 的架構設計，升級了卷積和脈衝驅動的自注意力運算元，提出了 E-SpikeFormer。E-SpikeFormer 在效能和能效上均有提升。

為解決擴充套件時效能下降的問題，我們設計了脈衝稀疏卷積（SSC），並結合掩膜影像建模（MIM）預訓練策略，顯著提升了模型的表現。

所提出的方法在多個任務上得到驗證，包括靜態影像分類、目標檢測、語義分割和神經形態動作識別。我們成功擴充套件了 SNN 的規模，並在所有測試資料集上取得了當前最優的結果，為 SNN 的大規模應用奠定了堅實基礎。

本文主要貢獻

2.1 重新思考脈衝神經元機制

我們先深入剖析普通脈衝神經元的積分、發放及重置機制，發現二元脈衝發放是影響脈衝神經元空間表示與時間動態的關鍵機制缺陷。具體表現為：

其一，普通脈衝神經元無法評估當前空間輸入的重要性。不管膜電位大小，輸出皆為單個脈衝，這相當於在膜電位向脈衝轉換時存在顯著量化誤差。

其二，普通脈衝神經元無法執行與資料相關的記憶或遺忘操作。因其脈衝發放值僅能為 0 或 1，意味著軟復位只能遺忘固定值。

以圖 1 中脈衝神經元為例，假設第個時間步長的空間輸入極為重要（數值大）。硬復位時，脈衝神經元輸出一個脈衝，膜電位清零至重置電位，這顯然存在較大量化誤差，無法體現輸入的重要性。而軟復位時，脈衝神經元輸出一個脈衝，傳遞到下一時間步長的膜電位依然很大。

如此一來，後續脈衝神經元可能長時間持續發放，致使無法判斷第 t 個時間步長的發放是由第個時間步長的空間輸入，還是當前空間輸入所致（如圖 1（d）所示）。所以，軟復位可能在每個時間步長產生相同的記憶和遺忘問題。

▲ 圖1

2.2 脈衝發放近似（SFA）方法

為解決二值脈衝發放機制的缺陷，我們提出了脈衝發放近似（SFA）方法。如圖 2 所示，SFA 利用整數訓練與脈衝序列推理的等價性，最佳化 SNN 的脈衝發放模式。

SFA 方法分為三步：首先，假設脈衝神經元在空間維度傳輸脈衝發放率；其次，訓練階段採用整數值啟用近似脈衝發放率，減少量化誤差，提升模型學習能力；最後，推理階段將整數啟用轉換為多時間步的二元脈衝，擴充套件時間步長以實現高效推理。

▲ 圖 2（a）

SFA 方法透過獨特設計在靜態和動態視覺任務中展現出顯著優勢。在靜態任務中，SFA 僅需單步訓練，大幅降低訓練成本，並在推理時無需重複輸入影像，提升效率並減少計算資源消耗。

在動態任務中，SFA 透過擴充套件時間步長克服了傳統二元脈衝發放機制的缺陷，賦予脈衝神經元動態遺忘能力，顯著提升模型效能。同時 SFA 方法在神經形態晶片上既可同步實現，適用於時間同步要求高的場景，也可非同步實現，充分發揮低功耗和低延遲優勢，尤其適合非同步連續脈衝發放特性。

▲ 圖 2（b）

同時，SFA 方法改變了脈衝的發射模式（如上圖 2（b）所示）。假設近似值為 0.3（=0.3），即在十個時間步（D=10）內需要發射三個脈衝。對於 ANN2SNN 和直接訓練 SNN，三個脈衝出現的時間步是同步或隨機的，其中同步發射意味著必須完成所有十個時間步，才能計算脈衝發射率。

相比之下，SFA 訓練在前三個時間步發射脈衝後，不會在其餘七個時間步發射脈衝。因此，SFA 中的脈衝發射可以以非同步的方式實現。

2.3 E-SpikeFormer 架構

Meta-SpikeFormer [5] 探索了脈衝神經網路在脈衝驅動自注意力機制（SDSA）、架構和快捷連線方面的元設計。

在此基礎上，我們致力於設計一種更高效的脈衝神經網路，將其命名為 E-SpikeFormer （見下圖3）。E-SpikeFormer 結合了卷積和基於 Transformer 的 SNN 模組，每個模組包含令牌混合器和通道混合器。以下是我們的主要改進：

▲ 圖3

1. 去除 RepConv：我們移除了 Meta-SpikeFormer 中高能耗的重引數化卷積（RepConv），代之以輕量級的脈衝可分離卷積（SpikeSepConv），從而顯著降低了能耗。

2. 基於卷積的 SNN 模組：我們將令牌混合器替換為 SpikeSepConv，最佳化後的結構如下：

其中，SpikeSepConv 透過深度卷積和逐點卷積結合脈衝神經元層實現。

3. 基於 Transformer 的 SNN 模組：我們在 SDSA 模組前添加了 SpikeSepConv，以增強區域性特徵提取，最佳化後的結構為：

4. 高效脈衝驅動自注意力（E-SDSA）：我們用線性操作替代 RepConv 生成、、：

透過這些最佳化，我們的 E-SpikeFormer 在降低能耗的同時大幅度提升了模型精度，使其可以適用於各種視覺任務。

2.4 脈衝稀疏卷積

我們發現直接對視覺 Transformer（ViT）[7] 進行規模擴充套件會導致效能下降，這一問題同樣出現在我們提出的 E-SpikeFormer 中。

為了解決這一問題，我們採用了預訓練-微調的掩碼影像建模（MIM）[7] 正規化，並結合脈衝稀疏卷積（SSC）編碼器，提出了 MIM-SNN 方法（見下圖 4）。該方法流程如下：

▲ 圖4

首先進行掩碼操作，我們首先對輸入影像進行補丁塊級別的掩碼操作。將影像重新組織為一系列展平的二維補丁塊，生成隨機二進位制掩碼，從而獲得兩個互補的掩碼檢視和。

接下來我們介紹編碼器。掩碼影像與神經形態資料類似，僅在亮度變化時觸發事件，因此包含大量空白區域。由於脈衝驅動的特性，SNN 中的卷積本質上是稀疏卷積。然而，普通脈衝卷積（VSC）在 MIM-SNN 中會導致資訊洩露，因為它在所有位置執行卷積，導致深層脈衝神經元幾乎不再發放脈衝（脈衝退化）。

為解決這一問題，我們提出了脈衝稀疏卷積（SSC），它僅在未被掩碼的區域執行卷積。具體來說，我們將輸入特徵 \(X_1\) 在 \(d\) 維空間中重塑，並使用核權重進行 SSC 處理。SSC 的卷積過程表示為：

其中，是選擇器，決定是否參與計算。我們在預訓練階段使用 SSC，微調階段則將權重轉換回VSC。這一方法有效避免了資訊洩露，提升了模型效能。。

之後我們介紹解碼器。編碼器的輸出與可學習的掩碼令牌一起提供給解碼器，解碼器將掩碼令牌重建為影像。在預訓練期間，我們採用普通的 ANN（Transformer）作為解碼器，在微調時將其移除。

最後是損失函式在預訓練中，我們使用均方誤差（MSE）損失作為度量函式：

其中，是解碼器的輸出，表示範數。同時在微調階段，我們採用了蒸餾（distillation）技術，這與 Meta-SpikeFormer [6] 中的方法一致。

實驗結果

3.1 影像分類

在 ImageNet-1k 資料集上，E-SpikeFormer 在不同模型規模下均取得優異成績。例如，如圖 5，6 所示，10M 引數模型的 Top-1 準確率達到 78.5%，比之前最好的 SNN（71.3%）高 7.2%，訓練時間加速 4.5 倍（從 63 小時縮短到 14 小時），推理能效提高 3.9 倍（從 11.9mJ 降低到 3.7mJ）。

在 172M 引數規模下，E-SpikeFormer 的 Top-1 準確率高達 86.2%，超越了許多經典 ANN 模型（如 ConvNeXt-B 的 83.8% 和 CrossFormer-L 的 84.0%），同時功耗僅為 19.1mJ，遠低於 ANN 模型的 70.8mJ（ConvNeXt-B）和 74.1mJ（CrossFormer-L）。

這些結果表明，E-SpikeFormer 在影像分類任務中不僅效能優異，還具備顯著的能效優勢。

▲ 圖5

▲ 圖6

3.2 目標檢測

在目標檢測任務中，E-SpikeFormer 在 COCO 資料集上達到 SOTA 結果。例如，E-SpikeFormer 的引數數量為 3870萬，[email protected] 為 58.8%，功耗為 119.5mJ，優於 Meta-SpikeFormer（引數 7500 萬，[email protected] 為 51.2%，功耗 140.8mJ）和 PVT-Tiny（引數 3290 萬，[email protected] 為 56.7%，功耗 1040.5mJ）。

E-SpikeFormer 在 [email protected] 上超越 Meta-SpikeFormer，同時功耗顯著低於 PVT-Tiny，展示了其在目標檢測任務中的高效性和效能優勢。

3.3 語義分割

在語義分割任務中，E-SpikeFormer 在 ADE20K 資料集上表現優異。例如，E-SpikeFormer 的引數數量為 11 M，MIoU 為 41.4%，功耗為 33.6mJ，相比 Meta-SpikeFormer（引數 59.8 M，MIoU 為 35.3%，功耗 183.6mJ）在 MIoU上提升 6.1%，同時功耗降低 81.7%。這一結果證明了 E-SpikeFormer 在語義分割任務中的高效性和準確性。

3.4 神經形態人類動作識別

在神經形態人類動作識別任務中，E-SpikeFormer 在 HAR-DVS 資料集上驗證了其有效性。例如，E-SpikeFormer 的準確率為 92.5%，相比 Meta-SpikeFormer（90.8%）提升 1.7%。在時間步長 T=8 的設定下，E-SpikeFormer 在直接訓練方式下表現出色，進一步驗證了其在神經形態任務中的高效性和效能提升。

3.5 消融實驗

透過消融實驗，我們驗證了 SFA 訓練和 E-SpikeFormer 架構設計的有效性。例如，基線模型在加入 SFA 後，時間步設定從 4×1 改為 1×4，功耗從 11.9mJ 降低到 7.3mJ（降低 4.6mJ），準確率從 71.3% 提升到 77.8%（提升 6.5%）。

進一步結合 E-SpikeFormer 架構後，功耗從 7.3mJ 降低到 3.7mJ（降低 4.3mJ），準確率從 77.8% 提升到 78.5%（提升 0.7%）。這表明 SFA 策略比現有的直接訓練 SNN 方法更有效、更高效，同時所提出的高效架構設計在降低功耗的同時進一步提高了準確率。

討論

4.1 SFA 在功率和效能方面的優勢

脈衝發放近似（SFA）方法透過引入一種全新的脈衝發放模式，顯著優化了脈衝神經網路（SNNs）的功耗和效能。我們透過分析網路脈衝發放率（NSFR）和脈衝分佈情況，揭示了 SFA 的最佳化機制。

首先，SFA 方法在功耗最佳化方面表現出色。實驗表明，採用 SFA 訓練的 SNN 在不同時間步長下的脈衝發放率逐漸降低（圖 7），而常規訓練的脈衝發放率則保持較高且穩定（圖 8）。

圖 7 中的脈衝圖進一步驗證了這一點：SFA 訓練的脈衝圖從左到右顏色由紅變藍，表明脈衝發放率隨時間步長增加而減少，而常規訓練的脈衝圖顏色變化不明顯。

此外，SFA 訓練的最高脈衝發放率為 0.19，遠低於常規訓練的 0.78，表明 SFA 顯著降低了脈衝發放密度，從而減少了計算能耗。

其次，SFA 方法在效能最佳化方面也有顯著優勢。透過將脈衝圖在時間和通道維度上平均，SFA 訓練的脈衝特徵更加清晰，且發放率更低（圖 7 最右側的脈衝圖）。這表明 SFA 方法能夠更有效地提取關鍵資訊，同時減少冗餘脈衝，從而提升模型的效能和效率。

▲ 圖7

▲ 圖8

4.2 為什麼二元脈衝發射機制不利於SNN的擴充套件

本中，我們從編碼器設計的角度分析了直接擴充套件 E-SpikeFormer 導致效能下降的原因。這一問題的根源在於，掩膜影像建模（MIM）中的掩膜操作會使得常規脈衝卷積洩露資訊。為此，我們提出了脈衝稀疏卷積（SSC）來解決這一問題。此外，我們還觀察到二進位制脈衝發放會干擾脈衝神經網路（SNNs）的擴充套件。

透過實驗，我們從兩個正交角度分析了 SNN 的擴充套件情況：是否採用 MIM-SNN 預訓練，以及使用脈衝發放近似（SFA）方法還是直接訓練。

如圖 9（a）所示，當使用 SFA 方法並結合 MIM 策略時，SNN 能夠有效擴充套件（橙色曲線）。然而，在常規訓練中，即使採用 MIM 策略，效能仍會下降（紅色曲線），這表明二進位制脈衝發放限制了模型的擴充套件能力。

我們認為，二進位制脈衝發放影響了 MIM-SNN 編碼器的特徵提取能力。透過有效秩（Rank）來衡量特徵提取能力，我們發現：在常規訓練的 MIM-SNN 中，有效秩在訓練過程中逐漸下降（圖 9（b）），表明特徵表示坍縮到低維空間，損害了模型的泛化能力和任務效能。

相比之下，採用 SFA 方法的 MIM-SNN 在訓練過程中有效秩增加並趨於穩定，表明模型具有更穩健的特徵表示和更好的微調效能。

▲ 圖9

4.3 SFA在神經形態晶片上的實現

常規訓練和 SFA 訓練的脈衝發放模式存在顯著差異。在 ANN2SNN 和常規直接訓練中，多個時間步長的脈衝發放是隨機的，而 SFA 訓練中的脈衝則是集中且連續的。這種差異決定了它們在神經形態晶片上的實現機制不同。圖 10 展示了 SFA 方法在神經形態晶片上的兩種實現方式：

1. 同步實現：整數啟用值被同步轉換為脈衝。同步神經形態晶片利用全域性時鐘，將 T 個時間步長擴充套件為 T×D 個時間步長。這種方式適用於靜態和動態視覺任務，且“天機”晶片支援此類操作。

2. 非同步實現：非同步神經形態晶片在極短的時間視窗內發出一系列脈衝（最多 D 個）。由於沒有全域性時鐘，脈衝神經元獨立工作，僅在事件輸入時啟用，其餘時間處於空閒狀態。Speck 晶片支援這種實現方式。

SFA 方法在神經形態晶片上具有顯著優勢：其相容性使其可透過同步或非同步晶片實現，而常規訓練的脈衝神經網路因依賴全域性時鐘難以在非同步晶片上執行；在能效與延遲方面，當 D 值較小且脈衝發放率較低時，SFA 的非同步實現顯著降低功耗和延遲，儘管存在通訊頻寬超載風險，但風險可控。

此外，D 值對能耗的影響因任務而異，例如在影像分類任務中，D 值從 1×4 調整為 1×8 導致的能耗增加幅度高於目標檢測和語義分割任務。綜上所述，SFA 方法的獨特脈衝發放模式使其與非同步神經形態晶片高度相容，為演算法驅動的晶片設計提供了重要參考。

▲ 圖10

4.4 神經形態晶片自然支援MIM-SNN預訓練

脈衝卷積在神經形態晶片上以脈衝驅動的方式執行，透過逐個事件的稀疏處理實現低延遲和低功耗的優勢。系統狀態隨每個脈衝輸入動態更新，SNN 核心透過地址搜尋識別權重和目標神經元位置，並更新目標神經元狀態。為解決 MIM-SNN 預訓練中的資訊洩露問題，我們提出了脈衝稀疏卷積（SSC）。

如圖 11 所示，與常規脈衝卷積（VSC）相比，SSC 僅在卷積核中心位置有脈衝輸入時執行卷積操作，確保資訊傳遞到下一層的相應位置。這一特性在 MIM 編碼器中尤為重要，因為大部分資訊需要被掩膜處理。

SSC 透過記錄掩膜位置，僅在未掩膜區域操作，從而避免資訊洩露。由於 VSC 和 SSC 的共性，只需對 VSC 的定址對映函式稍作修改，即可在神經形態晶片上實現 SSC，支援基於 MIM 預訓練的片上學習。

▲ 圖11

參考文獻

[1] Yao, M., Richter, O., Zhao, G., Qiao, N., Xing, Y., Wang, D., Hu, T., Fang, W., Demirci, T., De Marchi, M., Deng, L., Yan, T., Nielsen, C., Sheik, S., Wu, C., Tian, Y., Xu, B., Li, G.: Spike-based dynamic computing with asynchronous sensing computing neuromorphic chip. Nature Communications, 2024, 15(1): 4464.

[2] Bu, T., Fang, W., Ding, J., DAI, P., Yu, Z., & Huang, T. Optimal ANN-SNN Conversion for High-accuracy and Ultra-low-latency Spiking Neural Networks. In International Conference on Learning Representations (2022), https://openreview.net/forum?id=7B3IJMM1k_M.

[3] Wu, Y., Deng, L., Li, G., Zhu, J., Shi, L. Spatio-temporal backpropagation for training high-performance spiking neural networks. Frontiers in neuroscience, 2018, 12: 331.

[4] Neftci E O, Mostafa H, Zenke F. Surrogate gradient learning in spiking neural networks: Bringing the power of gradient-based optimization to spiking neural networks. IEEE Signal Processing Magazine, 2019, 36(6): 51-63.

[5] Fang, W., Chen, Y., Ding, J., Yu, Z., Masquelier, T., Chen, D., Huang, L., Zhou, H., Li. G, Tian, Y. Spikingjelly: An open-source machine learning infrastructure platform for spike-based intelligence. Science Advances, 2023, 9(40): eadi1480.

[6] Yao, M., Hu, J., Hu, T., Xu, Y., Zhou, Z., Tian, Y., XU, B., Li, G.: Spike-driven transformer v2: Meta spiking neural network architecture inspiring the design of next-generation neuromorphic chips. In International Conference on Learning Representations (2024), https://openreview.net/forum?id=1SIBN5Xyw7

[7] Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Houlsby, N. An Image is Worth 16×16 Words: Transformers for Image Recognition at Scale[C]//International Conference on Learning Representations (2020), https://openreview.net/forum?id=YicbFdNTTy

更多閱讀