降低百倍時間步,精度媲美傳統神經網路:上交等機構提出ANN-SNN 轉換框架SpikeConverter

語言: CN / TW / HK

論文連結:https://www.aaai.org/AAAI22Papers/AAAI-364.LiuF.pdf

脈衝神經網路 (Spiking Neural Network, SNN)被譽為第三代的神經網路,以其豐富的時空領域的神經動力學特性、多樣的編碼機制、事件驅動的優勢引起了學者的關注。由於生物合理性與任務需求間存在不同程度的取捨,對網路使用的神經元模型、編碼方式等各有差異,這造成了 SNN 的多樣化。與向量乘法為基礎的傳統人工神經網路(ANN)不同,SNN 通過模擬大腦的功能結構,以神經元為計算、儲存單元,通過累積在膜電壓上由相鄰神經元發出的脈衝訊號,並在超過閾值電位時發射脈衝的方式傳遞資訊。由於其計算可以完全由加法完成,需要的計算量和功耗大幅減少[1]。此外,由於 SNN 中神經元的執行、學習機制,使得 SNN 與一些新興器件天然結合在一起[2]。

脈衝神經網路 的目標是在基於非同步電路的實現,每個神經元可以非同步接收和發射脈衝訊號。然而這種大規模的非同步電路目前還無法實現, 脈衝神經網路 中的所有脈衝訊號需要對齊至時鐘訊號。這樣做會帶來兩個問題:其一,為了精確表達脈衝發射的時間,需要將一次推理分割成大量的時間步(time step),使得脈衝序列變的十分冗長;其二,脈衝序列中的脈衝數量沒有變化,使得脈衝序列出現不規則的稀疏脈衝。現有的同步 脈衝神經網路 晶片,例如 Loihi,在一定程度上利用了稀疏性。但是,時間步增加導致的推理時間延長問題依然存在。

現有的 脈衝神經網路 的訓練方法分為兩種方法:直接訓練和通過傳統神經網路轉換得來。

直接訓練的 SNN:可以進一步劃分為有監督訓練和無監督訓練。有監督訓練的 SNN 是基於脈衝的反向傳播演算法,然而神經元函式不可微,導致傳統的反向傳播演算法沒辦法直接應用在 SNN 上。無監督訓練的 SNN 是利用突觸可塑性規則的仿生學習演算法。例如:脈衝時序依賴的突觸可塑性(Spike-Timing Dependent Plasticity, STDP)是最具代表性的無監督 SNN,根據神經元激發的先後順序,調整神經元之間連線的強弱。

基於轉換的 SNN:為了在已發展的深度學習成果上進一步利用 脈衝神經網路 低能耗的特點,從 ANN 的視角出發,將連續的啟用函式值通過 SNN 中頻率編碼下的平均脈衝發放率進行替代,並且使用神經元函式替換 ANN 中的 ReLU 函式,再通過特定手段將其轉換為 SNN。從本質上說,基於轉換的 SNN 其訓練仍然依賴於 ANN 中的反向傳播演算法,因此它避免了對 SNN 進行直接訓練的困難。就效能表現而言,基於轉換的 SNN 保持著與 ANN 發展相匹配的程序。

然而,現有的 SNN 在實現超大規模的網路結構上,仍然面臨梯度消失、訓練資源開銷大甚至演算法收斂性的問題,一般只能應用在淺層的 SNN 中,並且精度還距離 ANN 有一定的差距。

而現有的轉換方法應用在更深的網路或更復雜的資料集時,需要大量的時間步長表示脈衝序列(例如:在 ImageNet 中,VGG轉換而來的 SNN 需要的時間步數量高達 2048[3,4])。相比於進行一次 ANN 的推理只需要執行一次前向傳遞過程;SNN 在每個時間步都需要執行一次前向傳遞。因此,時間步的數量過多會引入極大的推理延遲,導致了轉換後的 SNN 網路難以勝任對實時性有要求的應用。

方法介紹

本文提出一種基於高效且準確的 ANN-SNN 轉換框架,不僅可以應用到傳統神經網路的最新發展,而且大幅度降低了時間步的數量。只需要 8-16 個時間步就可以達到和傳統神經網路幾乎一致的精度,節省了百倍的時間延遲。

轉換中的一致性關係

首先,我們在神經元發射脈衝之後採用了軟重置的策略,在當前的膜電位中減去閾值電位的值,從而保留了超出閾值的部分。在之前的大部分文章中,採用的是硬重置策略,當膜電壓超過閾值時發射脈衝並將膜電壓直接重置為 0。這樣做的好處是儘可能多的保留了膜電壓中包含的資訊,並且能夠在輸入的微小擾動下保持相對一致的輸出結果。軟重置方法可以用公式表達為

其中 V[t]是 t 時刻的膜電壓,k 是每個時間步之間的膜電壓衰減係數。Xi 是第 i 個相鄰神經元的脈衝序列,0 表示 t 時沒有脈衝傳輸。Y 是當前神經元的輸出脈衝序列。將公式 (1) 在 t=1 到 T 累加,我們可以得到

在理想情況下,最後神經元內沒有膜電壓存留,即 V[T]=0。此時,公式 (2) 可以簡化為

不難發現,公式 (3) 與傳統神經網路的 MAC 計算模式 十分相似,而事實上之前最為常用的基於 ANN 轉換的 SNN,將 ANN 的啟用值作為 SNN 中的神經元發射脈衝的頻率,正是公式 (3) 在 k=1 的特例。

然而,公式 (3) 在很多情況下並不成立。最明顯的一種情況,是膜電壓在最後並不一定會等於 0 或一個相對較小的數。事實上,由於神經元之間的權重有負值,導致很多時候膜電壓最終是一個負數。SNN 的神經元模型的機制決定了其只會對正的電壓做出反應而累積的負膜電位則會被保留,從而導致公式 (3) 的等號並不成立,資訊無法準確表達,使得 脈衝神經網路 的精度低於傳統神經網路。例如當 k=1 以及閾值電位 = 1 的情況下,如果輸入總膜電位分別為 1,1,-1,-1,則會在前兩個時間步產生脈衝訊號;而後兩個 - 1 則被累積在膜電位中無法消耗。此時公式左邊 = 0 而右邊 = 2,公式並不成立。

神經元計算中的時域分離

針對這一點,文章採用了時域分離的方法,將膜電壓的累積過程和膜電壓釋放脈衝的過程在時域上分離開,避免膜電壓累積資訊丟失的情況出現。為了實現時域分離的效果,即發射脈衝在累積電壓之後,應使得較晚時間的膜電壓在沒有額外輸入的情況下要大於之前的膜電壓。這在之前的 leaky integrate-and-fire 模型中是無法做到的。文章提出了與傳統 LIF 相反的模型,iLIF,在每個時間步結束後增幅而不是減少模電壓,即在公式中的模電壓衰減係數 k 現在是一個大於 1 的數,被稱作為電壓增幅係數。

然而,時域分離的方法使得完成一個神經元計算的時間翻倍,文章通過流水線的方法,在前一層的脈衝輸出階段直接將訊號傳輸至下一層作為輸入訊號。在消除了時域分離帶來的額外執行時間的同時,也減少了需要臨時儲存的內容。

以上示意圖展示了圖片推理的執行示意圖。核心是利用層間並行,實現不同圖片的兩個執行階段重疊,降低推理時延。

引數選擇方面,主要是電壓增幅係數的選擇。文章考慮了不同的時間步數量和增幅係數的情況下,平均的轉化誤差。實驗表明在不同的時間步選擇下,k=2 都是最優選擇。

以上示意圖展示了在不同引數選擇下,轉化誤差的變化情況。K=2 時,在各個時間步下,轉換誤差最小,可以膜電壓累積表示的資訊最準確。

實驗結果

文章使用 CIFAR-10/100,以及 ImageNet 資料集進行實驗。實驗結果表明文章提出的方法在僅需要 16 個時間步的情況下,可以達到和傳統神經網路幾乎一致的精度。極大節省空間儲存和計算代價的情況下取得了最佳的效能,具體實驗結果參見論文。部分結果如下。

未來方向

脈衝神經網路 還有較大的挖掘空間,還有許多領域 SNN 未能成功應用,可以在視訊識別、機器人控制等方面繼續研究。

蔣力老師團隊主要從事計算機體系架構與晶片設計自動化(EDA)研究。作者劉方鑫主要研究方向包括存內計算,神經模態計算,以及神經網路軟硬體協同設計。作為第一作者發表在TCAD, DAC, ICCAD, ICCD, ICCV, AAAI, SIGIR等頂級期刊,會議上發表論文10餘篇。

[1] Roy, K.; Jaiswal, A.; and Panda, P. 2019. Towards spike based machine intelligence with neuromorphic computing. Nature, 575(7784): 607–617.

[2] Woźniak, Stanisław, et al. "Deep learning incorporating biologically inspired neural dynamics and in-memory computing." Nature Machine Intelligence 2.6 (2020): 325-336.

[3] Han, Bing, Gopalakrishnan Srinivasan, and Kaushik Roy. "Rmp-snn: Residual membrane potential neuron for enabling deeper high-accuracy and low-latency spiking neural network." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2020.

[4] Han, Bing, and Kaushik Roy. "Deep spiking neural network: Energy efficiency through time based coding." European Conference on Computer Vision. Springer, Cham, 2020.

「其他文章」