論文鏈接：https://arxiv.org/pdf/2408.16467
代碼鏈接：https://github.com/AndyCao1125/SDM

亮點(diǎn)直擊

• 本文提出了脈沖擴(kuò)散模型（Spiking Diffusion Model, SDM），一種高質(zhì)量的圖像生成器，在基于SNN的生成模型中實(shí)現(xiàn)了SOTA性能。
• 從生物學(xué)的角度出發(fā)，本文提出了一種時(shí)間脈沖機(jī)制（Temporal-wise Spiking Mechanism, TSM），使脈沖神經(jīng)元能夠捕捉更多的動(dòng)態(tài)信息，從而提高去噪圖像的質(zhì)量。
• 大量結(jié)果顯示，SDM在CIFAR-10數(shù)據(jù)集上的FID分?jǐn)?shù)上超越了SNN基線模型多達(dá)12倍，同時(shí)節(jié)省了約60%的能耗。此外，本文提出了一種閾值引導(dǎo)策略，以進(jìn)一步提高生成性能。

近年來，脈沖神經(jīng)網(wǎng)絡(luò)（Spiking Neural Networks, SNNs）因其超低能耗和高生物可塑性相比傳統(tǒng)人工神經(jīng)網(wǎng)絡(luò)（Artificial Neural Networks, ANNs）而受到關(guān)注。盡管SNNs具有獨(dú)特的優(yōu)勢，但其在計(jì)算密集型的圖像生成領(lǐng)域的應(yīng)用仍在探索之中。本文提出了脈沖擴(kuò)散模型（Spiking Diffusion Models, SDMs），一種創(chuàng)新的基于SNN的生成模型家族，能夠以顯著降低的能耗生成高質(zhì)量樣本。特別地，本文提出了一種時(shí)間脈沖機(jī)制（Temporal-wise Spiking Mechanism, TSM），使SNNs能夠從生物可塑性角度捕捉更多的時(shí)間特征。此外，本文提出了一種閾值引導(dǎo)策略，可以在不進(jìn)行額外訓(xùn)練的情況下將性能提高多達(dá)16.7%。本文還首次嘗試使用ANN-SNN方法進(jìn)行基于SNN的生成任務(wù)。大量實(shí)驗(yàn)結(jié)果表明，本文的方法不僅在少量脈沖時(shí)間步中表現(xiàn)出與其ANN對應(yīng)模型相當(dāng)?shù)男阅?，而且在很大程度上?yōu)于之前基于SNN的生成模型。此外，本文還展示了SDM在大規(guī)模數(shù)據(jù)集（例如LSUN臥室）上的高質(zhì)量生成能力。這一發(fā)展標(biāo)志著SNN基生成能力的一個(gè)重要進(jìn)步，為未來實(shí)現(xiàn)低能耗和低延遲的生成應(yīng)用開辟了新的研究途徑。

方法

A. 峰值前殘差學(xué)習(xí)

本文首先分析了先前脈沖神經(jīng)網(wǎng)絡(luò)（SNNs）中的殘差學(xué)習(xí)方法存在的局限性和概念不一致性，特別是SEW ResNet，其公式可以表示為：

受 [61], [62] 的啟發(fā)，本文在本文的脈沖UNet中采用了激活-卷積-批量歸一化（Activation-Conv-BatchNorm）結(jié)構(gòu)的預(yù)脈沖殘差學(xué)習(xí)方法，解決了基于卷積的SNNs中梯度爆炸/消失和性能下降的雙重挑戰(zhàn)。通過預(yù)脈沖塊，殘差和輸出通過浮點(diǎn)加法操作進(jìn)行求和，確保在進(jìn)入下一個(gè)脈沖神經(jīng)元之前表示是準(zhǔn)確的，同時(shí)避免了上述病態(tài)情況。整個(gè)預(yù)脈沖殘差學(xué)習(xí)過程在一個(gè)殘差塊內(nèi)可以表示如下：

總而言之，TSM允許膜電位在時(shí)間域內(nèi)動(dòng)態(tài)更新，從而提高捕捉潛在時(shí)間依賴特征的能力。后續(xù)實(shí)驗(yàn)表明，TSM機(jī)制優(yōu)于傳統(tǒng)的固定更新機(jī)制。

C. SDM 中的閾值指導(dǎo)

理論能耗計(jì)算

在本節(jié)中，本文描述了計(jì)算脈沖UNet架構(gòu)理論能耗的方法。該計(jì)算包括兩個(gè)主要步驟：確定架構(gòu)中每個(gè)模塊的突觸操作（SOPs），然后基于這些操作估算整體能耗。脈沖UNet每個(gè)模塊的突觸操作可以如下量化：

通過 ANN-SNN 轉(zhuǎn)換實(shí)施尖峰擴(kuò)散模型

在本文中，本文首次利用ANN-SNN方法成功實(shí)現(xiàn)了SNN擴(kuò)散。本文采用了Fast-SNN方法來構(gòu)建量化ANN與SNN之間的轉(zhuǎn)換。由于這一實(shí)現(xiàn)并不是本文論文的主要貢獻(xiàn)，本文將簡要描述ANN-SNN的原理，更多細(xì)節(jié)可以在[67]中找到。

實(shí)驗(yàn)

A. 實(shí)驗(yàn)設(shè)置

數(shù)據(jù)集和評估指標(biāo)

為了展示所提算法的有效性和效率，本文在以下數(shù)據(jù)集上進(jìn)行實(shí)驗(yàn)：32×32 MNIST、32×32 FashionMNIST 、32×32 CIFAR-10和 64×64 CelebA。定性結(jié)果根據(jù)Frechet Inception Distance（FID，越低越好）和Inception Score（IS，越高越好）進(jìn)行比較。FID分?jǐn)?shù)是通過比較50,000張生成圖像與數(shù)據(jù)集的相應(yīng)參考統(tǒng)計(jì)數(shù)據(jù)來計(jì)算的。

實(shí)現(xiàn)細(xì)節(jié)

B. 與最先進(jìn)的方法比較

在下表 I 中，本文展示了本文的脈沖擴(kuò)散模型（SDMs）與當(dāng)前最先進(jìn)的生成模型在無條件生成任務(wù)中的比較分析。本文還包括了ANN的結(jié)果作為參考。定性結(jié)果展示在下圖4中。本文的結(jié)果表明，SDMs在所有數(shù)據(jù)集上均顯著優(yōu)于SNN基線，即使在較少的脈沖仿真步數(shù)（4/8）下也是如此。特別是，SDDPM在CelebA數(shù)據(jù)集上相比FSVAE和SGAD（兩者均為16個(gè)時(shí)間步）分別有4倍和6倍的FID提升，而在CIFAR-10數(shù)據(jù)集上則有11倍和12倍的提升。正如預(yù)期的那樣，隨著時(shí)間步的增加，樣本質(zhì)量也會(huì)提高。本文還注意到，結(jié)合TSM后，性能進(jìn)一步提升，而模型參數(shù)僅有微小增加（2e-4 M）。SDMs還可以處理快速采樣求解器，并在較少的步驟中獲得更高的采樣質(zhì)量（見下表VI）。重要的是，SDMs在使用相同的UNet架構(gòu)下獲得了與ANN基線相當(dāng)?shù)馁|(zhì)量，甚至超過了一些ANN模型（例如，15.45 vs. 19.04）。這一結(jié)果突顯了本文模型中使用的SNN的卓越表達(dá)能力。

C. 與 ANN-SNN 方法的比較

為了驗(yàn)證SDM在ANN-SNN方法下的生成能力，本文在32×32 CIFAR-10和64×64 FFHQ數(shù)據(jù)集上進(jìn)行了實(shí)驗(yàn)。如下表II所示，ANN-SNN方法在CIFAR-10上表現(xiàn)出色（即51.18 FID），并且在微調(diào)策略后顯著提高了圖像質(zhì)量（即29.53 FID）。然而，ANN-SNN的結(jié)果與直接訓(xùn)練的結(jié)果之間仍存在差距。盡管ANN-SNN方法在基于分類的任務(wù)中表現(xiàn)出與ANN相當(dāng)?shù)男阅?，但在生成任?wù)方面仍缺乏深入的研究。ANN-SNN方法的定性結(jié)果展示在下圖7中。

D. 時(shí)間性尖峰機(jī)制的有效性

為了更好地可視化TSM模塊帶來的性能提升，本文提供了使用SDDIM生成的CIFAR-10圖像結(jié)果，分別展示了有和沒有TSM模塊的情況。這里本文使用DDIM而不是DDPM進(jìn)行比較，因?yàn)镈DIM基于常微分方程（ODEs）操作，確保了確定性和一致的生成結(jié)果。相比之下，DDPM依賴于隨機(jī)微分方程（SDEs），在生成過程中引入了隨機(jī)性，導(dǎo)致輸出圖像的可變性，從而使直接比較變得具有挑戰(zhàn)性。

下圖5中的結(jié)果顯示，帶有TSM模塊的生成圖像質(zhì)量有顯著提高。與沒有TSM模塊的圖像相比，這些圖像的輪廓更加清晰，背景更為清晰，紋理細(xì)節(jié)更豐富，從而證明了TSM的有效性。

E. 閾值指導(dǎo)的有效性

在前文中，本文提出了一種無需訓(xùn)練的方法：閾值引導(dǎo)（Threshold Guidance，TG），旨在通過在推理階段僅略微調(diào)整脈沖神經(jīng)元的閾值水平來提高生成圖像的質(zhì)量。如下表III所示，通過閾值調(diào)整進(jìn)行抑制性引導(dǎo)顯著提升了圖像質(zhì)量，在兩個(gè)關(guān)鍵指標(biāo)上都有所改善：FID分?jǐn)?shù)從19.73下降到19.20，閾值減少0.3%；IS分?jǐn)?shù)從7.44上升到7.55，閾值減少0.2%。相反，在某些條件下，興奮性引導(dǎo)同樣可以提高采樣質(zhì)量。這些發(fā)現(xiàn)強(qiáng)調(diào)了閾值引導(dǎo)作為一種在訓(xùn)練后顯著提高模型效果的方法的潛力，而無需額外的訓(xùn)練資源。本文在附錄中提供了更多關(guān)于閾值引導(dǎo)的解釋。

F. TSM方法分析

G. 計(jì)算成本評估

為了進(jìn)一步強(qiáng)調(diào)本文SDM的低能耗特性，本文對比分析了所提出的SDDPM與其對應(yīng)的ANN模型在FID和能耗方面的表現(xiàn)。如表IV所示，當(dāng)時(shí)間步長設(shè)定為4時(shí)，SDDPM的能耗顯著降低，僅為其ANN對應(yīng)模型的37.5%。此外，SDDPM的FID也提高了0.47，表明本文的模型在有效減少能耗的同時(shí)保持了競爭力的性能。當(dāng)本文將分析擴(kuò)展到不同的時(shí)間步長增量時(shí)，可以觀察到一個(gè)明顯的模式：隨著時(shí)間步長的增加，F(xiàn)ID分?jǐn)?shù)有所改善，但代價(jià)是能耗的增加。這一觀察結(jié)果表明，隨著時(shí)間步長的增加，F(xiàn)ID改善與能耗之間存在權(quán)衡。

H. 消融研究

不同組件對SDM的影響

本文首先在CIFAR-10數(shù)據(jù)集上進(jìn)行消融研究，以探討時(shí)間脈沖模塊（TSM）和閾值引導(dǎo)（TG）的影響。如下表V所示，本文發(fā)現(xiàn)TSM和TG都對圖像質(zhì)量的提升有貢獻(xiàn)。通過同時(shí)使用TSM和TG，本文獲得了最佳的FID結(jié)果，相較于原始的SDDIM提升了18.4%。

SDM在不同求解器上的有效性

在下表VI中，本文驗(yàn)證了SDM在各種擴(kuò)散求解器上的可行性和有效性。SDDIM在采樣步驟上表現(xiàn)出更穩(wěn)定的性能，而Analytic-SDPM展示了卓越的能力，達(dá)到了新的最先進(jìn)性能，超越了ANN-DDIM的結(jié)果。總之，本文的SDM證明了其在處理任何擴(kuò)散求解器方面的高效性，并且本文相信利用本文的SDM還有很大的潛力進(jìn)一步提升FID性能。

討論 & 結(jié)論

本研究提出了一種新的基于SNN的擴(kuò)散模型家族，稱為脈沖擴(kuò)散模型（SDMs），它結(jié)合了SNN的能效優(yōu)勢和卓越的生成性能。SDMs在SNN基線中以更少的脈沖時(shí)間步長達(dá)到了最先進(jìn)的結(jié)果，并且與ANNs相比，能耗更低。SDMs主要受益于兩個(gè)方面：(1) 時(shí)間脈沖機(jī)制（TSM），它使去噪網(wǎng)絡(luò)SNN-UNet的突觸電流在每個(gè)時(shí)間步長中能夠收集更多的動(dòng)態(tài)信息，而不是像傳統(tǒng)SNN那樣由固定的突觸權(quán)重控制；(2) 無需訓(xùn)練的閾值引導(dǎo)（TG），通過調(diào)整脈沖閾值進(jìn)一步提高采樣質(zhì)量。

然而，本文工作的一個(gè)限制是SNN-UNet的時(shí)間步長相對較小，未能充分挖掘SDMs的全部潛力。此外，還應(yīng)考慮在更高分辨率的數(shù)據(jù)集（如ImageNet）上進(jìn)行測試。在未來的研究中，本文計(jì)劃探索SDMs在生成領(lǐng)域的進(jìn)一步應(yīng)用，例如文本-圖像生成，并嘗試將其與先進(jìn)的語言模型結(jié)合，以實(shí)現(xiàn)更有趣的任務(wù)。

本文轉(zhuǎn)自 AI生成未來 ，作者：Jiahang Cao等

原文鏈接:??https://mp.weixin.qq.com/s/1JUjAoC1xZ0lZrj4S7X-ag??