SDM: 第三代神經網絡和擴散模型強強聯(lián)合！FID最多超基線12倍，能耗省60%，實力SOTA！
論文鏈接：https://arxiv.org/pdf/2408.16467
代碼鏈接：https://github.com/AndyCao1125/SDM
亮點直擊
- 本文提出了脈沖擴散模型（Spiking Diffusion Model, SDM），一種高質量的圖像生成器，在基于SNN的生成模型中實現(xiàn)了SOTA性能。
- 從生物學的角度出發(fā)，本文提出了一種時間脈沖機制（Temporal-wise Spiking Mechanism, TSM），使脈沖神經元能夠捕捉更多的動態(tài)信息，從而提高去噪圖像的質量。
- 大量結果顯示，SDM在CIFAR-10數(shù)據(jù)集上的FID分數(shù)上超越了SNN基線模型多達12倍，同時節(jié)省了約60%的能耗。此外，本文提出了一種閾值引導策略，以進一步提高生成性能。
近年來，脈沖神經網絡（Spiking Neural Networks, SNNs）因其超低能耗和高生物可塑性相比傳統(tǒng)人工神經網絡（Artificial Neural Networks, ANNs）而受到關注。盡管SNNs具有獨特的優(yōu)勢，但其在計算密集型的圖像生成領域的應用仍在探索之中。本文提出了脈沖擴散模型（Spiking Diffusion Models, SDMs），一種創(chuàng)新的基于SNN的生成模型家族，能夠以顯著降低的能耗生成高質量樣本。特別地，本文提出了一種時間脈沖機制（Temporal-wise Spiking Mechanism, TSM），使SNNs能夠從生物可塑性角度捕捉更多的時間特征。此外，本文提出了一種閾值引導策略，可以在不進行額外訓練的情況下將性能提高多達16.7%。本文還首次嘗試使用ANN-SNN方法進行基于SNN的生成任務。大量實驗結果表明，本文的方法不僅在少量脈沖時間步中表現(xiàn)出與其ANN對應模型相當?shù)男阅?，而且在很大程度上?yōu)于之前基于SNN的生成模型。此外，本文還展示了SDM在大規(guī)模數(shù)據(jù)集（例如LSUN臥室）上的高質量生成能力。這一發(fā)展標志著SNN基生成能力的一個重要進步，為未來實現(xiàn)低能耗和低延遲的生成應用開辟了新的研究途徑。
方法
A. 峰值前殘差學習
本文首先分析了先前脈沖神經網絡（SNNs）中的殘差學習方法存在的局限性和概念不一致性，特別是SEW ResNet，其公式可以表示為：
受 [61], [62] 的啟發(fā)，本文在本文的脈沖UNet中采用了激活-卷積-批量歸一化（Activation-Conv-BatchNorm）結構的預脈沖殘差學習方法，解決了基于卷積的SNNs中梯度爆炸/消失和性能下降的雙重挑戰(zhàn)。通過預脈沖塊，殘差和輸出通過浮點加法操作進行求和，確保在進入下一個脈沖神經元之前表示是準確的，同時避免了上述病態(tài)情況。整個預脈沖殘差學習過程在一個殘差塊內可以表示如下：
總而言之，TSM允許膜電位在時間域內動態(tài)更新，從而提高捕捉潛在時間依賴特征的能力。后續(xù)實驗表明，TSM機制優(yōu)于傳統(tǒng)的固定更新機制。
C. SDM 中的閾值指導
理論能耗計算
在本節(jié)中，本文描述了計算脈沖UNet架構理論能耗的方法。該計算包括兩個主要步驟：確定架構中每個模塊的突觸操作（SOPs），然后基于這些操作估算整體能耗。脈沖UNet每個模塊的突觸操作可以如下量化：
通過 ANN-SNN 轉換實施尖峰擴散模型
在本文中，本文首次利用ANN-SNN方法成功實現(xiàn)了SNN擴散。本文采用了Fast-SNN方法來構建量化ANN與SNN之間的轉換。由于這一實現(xiàn)并不是本文論文的主要貢獻，本文將簡要描述ANN-SNN的原理，更多細節(jié)可以在[67]中找到。
實驗
A. 實驗設置
數(shù)據(jù)集和評估指標
為了展示所提算法的有效性和效率，本文在以下數(shù)據(jù)集上進行實驗：32×32 MNIST、32×32 FashionMNIST 、32×32 CIFAR-10和 64×64 CelebA。定性結果根據(jù)Frechet Inception Distance（FID，越低越好）和Inception Score（IS，越高越好）進行比較。FID分數(shù)是通過比較50,000張生成圖像與數(shù)據(jù)集的相應參考統(tǒng)計數(shù)據(jù)來計算的。
實現(xiàn)細節(jié)
B. 與最先進的方法比較
在下表 I 中，本文展示了本文的脈沖擴散模型（SDMs）與當前最先進的生成模型在無條件生成任務中的比較分析。本文還包括了ANN的結果作為參考。定性結果展示在下圖4中。本文的結果表明，SDMs在所有數(shù)據(jù)集上均顯著優(yōu)于SNN基線，即使在較少的脈沖仿真步數(shù)（4/8）下也是如此。特別是，SDDPM在CelebA數(shù)據(jù)集上相比FSVAE和SGAD（兩者均為16個時間步）分別有4倍和6倍的FID提升，而在CIFAR-10數(shù)據(jù)集上則有11倍和12倍的提升。正如預期的那樣，隨著時間步的增加，樣本質量也會提高。本文還注意到，結合TSM后，性能進一步提升，而模型參數(shù)僅有微小增加（2e-4 M）。SDMs還可以處理快速采樣求解器，并在較少的步驟中獲得更高的采樣質量（見下表VI）。重要的是，SDMs在使用相同的UNet架構下獲得了與ANN基線相當?shù)馁|量，甚至超過了一些ANN模型（例如，15.45 vs. 19.04）。這一結果突顯了本文模型中使用的SNN的卓越表達能力。
C. 與 ANN-SNN 方法的比較
為了驗證SDM在ANN-SNN方法下的生成能力，本文在32×32 CIFAR-10和64×64 FFHQ數(shù)據(jù)集上進行了實驗。如下表II所示，ANN-SNN方法在CIFAR-10上表現(xiàn)出色（即51.18 FID），并且在微調策略后顯著提高了圖像質量（即29.53 FID）。然而，ANN-SNN的結果與直接訓練的結果之間仍存在差距。盡管ANN-SNN方法在基于分類的任務中表現(xiàn)出與ANN相當?shù)男阅?，但在生成任務方面仍缺乏深入的研究。ANN-SNN方法的定性結果展示在下圖7中。
D. 時間性尖峰機制的有效性
為了更好地可視化TSM模塊帶來的性能提升，本文提供了使用SDDIM生成的CIFAR-10圖像結果，分別展示了有和沒有TSM模塊的情況。這里本文使用DDIM而不是DDPM進行比較，因為DDIM基于常微分方程（ODEs）操作，確保了確定性和一致的生成結果。相比之下，DDPM依賴于隨機微分方程（SDEs），在生成過程中引入了隨機性，導致輸出圖像的可變性，從而使直接比較變得具有挑戰(zhàn)性。
下圖5中的結果顯示，帶有TSM模塊的生成圖像質量有顯著提高。與沒有TSM模塊的圖像相比，這些圖像的輪廓更加清晰，背景更為清晰，紋理細節(jié)更豐富，從而證明了TSM的有效性。
E. 閾值指導的有效性
在前文中，本文提出了一種無需訓練的方法：閾值引導（Threshold Guidance，TG），旨在通過在推理階段僅略微調整脈沖神經元的閾值水平來提高生成圖像的質量。如下表III所示，通過閾值調整進行抑制性引導顯著提升了圖像質量，在兩個關鍵指標上都有所改善：FID分數(shù)從19.73下降到19.20，閾值減少0.3%；IS分數(shù)從7.44上升到7.55，閾值減少0.2%。相反，在某些條件下，興奮性引導同樣可以提高采樣質量。這些發(fā)現(xiàn)強調了閾值引導作為一種在訓練后顯著提高模型效果的方法的潛力，而無需額外的訓練資源。本文在附錄中提供了更多關于閾值引導的解釋。
F. TSM方法分析
G. 計算成本評估
為了進一步強調本文SDM的低能耗特性，本文對比分析了所提出的SDDPM與其對應的ANN模型在FID和能耗方面的表現(xiàn)。如表IV所示，當時間步長設定為4時，SDDPM的能耗顯著降低，僅為其ANN對應模型的37.5%。此外，SDDPM的FID也提高了0.47，表明本文的模型在有效減少能耗的同時保持了競爭力的性能。當本文將分析擴展到不同的時間步長增量時，可以觀察到一個明顯的模式：隨著時間步長的增加，F(xiàn)ID分數(shù)有所改善，但代價是能耗的增加。這一觀察結果表明，隨著時間步長的增加，F(xiàn)ID改善與能耗之間存在權衡。
H. 消融研究
不同組件對SDM的影響
本文首先在CIFAR-10數(shù)據(jù)集上進行消融研究，以探討時間脈沖模塊（TSM）和閾值引導（TG）的影響。如下表V所示，本文發(fā)現(xiàn)TSM和TG都對圖像質量的提升有貢獻。通過同時使用TSM和TG，本文獲得了最佳的FID結果，相較于原始的SDDIM提升了18.4%。
SDM在不同求解器上的有效性
在下表VI中，本文驗證了SDM在各種擴散求解器上的可行性和有效性。SDDIM在采樣步驟上表現(xiàn)出更穩(wěn)定的性能，而Analytic-SDPM展示了卓越的能力，達到了新的最先進性能，超越了ANN-DDIM的結果?？傊疚牡腟DM證明了其在處理任何擴散求解器方面的高效性，并且本文相信利用本文的SDM還有很大的潛力進一步提升FID性能。
討論 & 結論
本研究提出了一種新的基于SNN的擴散模型家族，稱為脈沖擴散模型（SDMs），它結合了SNN的能效優(yōu)勢和卓越的生成性能。SDMs在SNN基線中以更少的脈沖時間步長達到了最先進的結果，并且與ANNs相比，能耗更低。SDMs主要受益于兩個方面：(1) 時間脈沖機制（TSM），它使去噪網絡SNN-UNet的突觸電流在每個時間步長中能夠收集更多的動態(tài)信息，而不是像傳統(tǒng)SNN那樣由固定的突觸權重控制；(2) 無需訓練的閾值引導（TG），通過調整脈沖閾值進一步提高采樣質量。
然而，本文工作的一個限制是SNN-UNet的時間步長相對較小，未能充分挖掘SDMs的全部潛力。此外，還應考慮在更高分辨率的數(shù)據(jù)集（如ImageNet）上進行測試。在未來的研究中，本文計劃探索SDMs在生成領域的進一步應用，例如文本-圖像生成，并嘗試將其與先進的語言模型結合，以實現(xiàn)更有趣的任務。
本文轉自 AI生成未來 ，作者：Jiahang Cao等
