论文 Spiking Diffusion Models 主要内容是提出了“脉冲扩散模型(Spiking Diffusion Models, SDMs)”,一种基于脉冲神经网络(SNN)的生成模型,旨在解决传统人工神经网络(ANN)在图像生成领域中的高能耗和计算强度问题。脉冲神经网络因其生物学上的合理性和低能耗特性而受到广泛关注。
对最先进的SNN模型进行比较。FID采用log2尺度,标记的大小对应于IS指标。与其他SNN生成模型相比,我们的模型在使用较少时间步长的情况下表现出更好的FID。
一、引言
本论文提出了一种名为“脉冲扩散模型(Spiking Diffusion Models, SDMs)”的新型生成模型。SDMs结合了脉冲神经网络(Spiking Neural Networks, SNNs)的低能耗特性与扩散模型(Diffusion Models)的强大生成能力,旨在解决传统人工神经网络(Artificial Neural Networks, ANNs)在图像生成任务中面临的高能耗和计算强度问题。脉冲神经网络因其生物学上的合理性和高效的能量利用率,近年来在神经计算领域受到广泛关注。然而,现有的SNN生成模型在图像质量和生成性能方面尚未达到ANN模型的水平。为此,本文提出了两项核心创新,以提升SNN在生成任务中的表现。
二、相关工作
2.1 脉冲神经网络(SNNs)
SNNs模拟生物神经元的脉冲传递机制,通过离散的脉冲(spikes)进行信息传递,具有低能耗和高效处理时序信息的优势。近年来,随着硬件技术的发展,SNNs在图像识别、目标检测等任务中取得了一定的成果。
2.2 扩散模型(Diffusion Models)
扩散模型是一类基于概率模型的生成方法,通过逐步去噪生成高质量的图像。近年来,扩散模型在图像生成任务中表现出色,成为GANs的有力竞争者。然而,扩散模型通常需要大量的计算资源,限制了其在低能耗设备上的应用。
2.3 SNN与扩散模型的结合
尽管SNN在能耗和时序信息处理上具有优势,但将其应用于生成任务尚处于起步阶段。现有工作主要集中在SNN与GANs的结合上,而将SNN应用于扩散模型的研究较为有限。因此,本文旨在填补这一研究空白,提出适用于扩散模型的SNN架构。
三、方法
3.1 SDMs的总体架构
SDMs结合了扩散模型的去噪过程与SNN的脉冲传递机制。具体来说,SDMs使用SNN-UNet作为去噪网络,通过脉冲神经元的时序活动来逐步生成图像。为了提升生成质量和能效,本文引入了两项关键机制:时间脉冲机制(Temporal-wise Spiking Mechanism, TSM)和阈值引导策略(Threshold Guidance, TG)。
我们脉冲扩散模型的概述。SDM的学习过程分为两个阶段:(1)训练阶段,和(2)微调阶段。在训练阶段,我们的脉冲UNet采用标准的前脉冲残差块(左下,见第IV-A节),然后在微调阶段将前脉冲块转换为TSM块(右下,见第IV-B节)。给定一个随机高斯噪声输入\(x_t\),它首先通过脉冲编码器转换为脉冲表示,然后与时间嵌入一起输入到脉冲UNet中。该网络仅传递由0/1向量(属于\(Z{0,1}\))表示的脉冲。最终,输出的脉冲通过解码器获得预测的噪声ϵ,并计算损失来更新网络。在微调阶段,我们加载训练阶段的权重,并用TSM块替换前脉冲块,此时时间参数p初始化为1.0。此阶段将继续优化网络参数,以获得更好的生成性能。
3.2 时间脉冲机制(TSM)
TSM模块旨在增强SNN在时间维度上的信息处理能力。传统SNN的脉冲传递通常依赖于固定的突触权重,限制了信息的动态传递。TSM通过动态调整神经元的膜电位,使得SNN-UNet能够在每个时间步捕捉更多的动态信息,从而提升生成图像的细节和质量。
时间脉冲机制概述。当脉冲神经元触发脉冲后,脉冲将在突触前被转换为输入电流\(I\)。为了获取更多动态信息,时间参数\(P\)将作用于电流,形成时间自适应电流\(Iˆ\)。
3.3 阈值引导策略(TG)
TG是一种无需额外训练的策略,通过在推理阶段微调神经元的阈值来优化生成过程。具体而言,TG通过调整脉冲神经元的激活阈值,可以在不增加训练成本的情况下,显著提升生成图像的质量。实验表明,适当的阈值调整能够有效降低生成图像的FID分数,同时提升IS分数。
四、实验
4.1 实验设置
-
数据集:MNIST、Fashion-MNIST、CIFAR-10、CelebA和LSUN-bed。
-
评价指标:FID(Frechet Inception Distance)和IS(Inception Score),以及能耗测量。
-
模型对比:与传统ANN扩散模型以及现有SNN生成模型进行对比,评估SDMs的生成性能和能耗优势。
-
不同数据集的结果。在所有数据集中,SDMs(我们的模型)在样本质量方面优于所有SNN基线模型,甚至超过了一些ANN模型,样本质量主要通过FID和IS进行评估。▽的结果来自文献[30],♮的结果来自文献[46]。ema表示使用了EMA [75]方法。为了公平比较,我们重新评估了DDPM [16]和DDIM [32]的结果,使用了与SDMs相同的UNet架构。∗表示仅对MNIST、Fashion-MNIST和CelebA使用FID评分,因为它们的数据分布与ImageNet差异较大,使得Inception Score无意义。排名前两位的结果分别用粗体和下划线表示。
4.2 无条件图像生成结果
在MNIST、Fashion-MNIST、CIFAR-10、CelebA和LSUN-bed数据集上,SDDMs通过直接训练方法展示了优异的生成能力。表中展示了不同时间步长下SDDPM的FID和IS分数,结果表明,随着时间步长的增加,生成质量提升明显。
使用ANN-SNN方法在CIFAR-10和FFHQ64上的无条件图像生成结果。
4.3 时间脉冲机制的有效性
通过对比使用和不使用TSM模块的CIFAR-10生成结果,实验结果(图5)显示,TSM模块显著提升了生成图像的质量,表现为图像轮廓更清晰、背景更明确、纹理细节更丰富。这证明了TSM在增强SNN动态信息处理能力方面的有效性。
4.4 阈值引导策略的有效性
表III展示了不同阈值引导下CIFAR-10的生成结果。通过轻微调整脉冲神经元的阈值,TG策略显著降低了FID分数并提升了IS分数,证明了TG在推理阶段优化生成质量的有效性。
4.5 TSM方法的分析
图6展示了各层时间脉冲机制参数p[t]的平均值,结果表明,随着时间步长的增加,p[t]呈上升趋势,表明后期时间步长在信息传递过程中具有更高的重要性。TSM通过动态调整这些参数,有效提升了SNN对时序动态的理解和生成图像的质量。
4.6 计算成本评估
表IV对比了SDDPM与其对应的ANN模型在FID分数和能耗上的表现。结果显示,SDDPM在保持竞争性FID分数的同时,能耗显著降低。例如,在时间步长为4时,SDDPM的能耗仅为ANN模型的37.5%,且FID分数仅略高于ANN模型,表明SDDPM在能效和生成性能之间实现了良好的平衡。
4.7 消融实验
-
表V:在CIFAR-10上,分别评估了TSM和TG对SDDIM的影响。结果表明,TSM和TG各自提升了生成质量,二者结合使用时,FID分数进一步显著下降,达到最佳效果。
-
表VI:评估了SDM在不同扩散求解器上的表现,结果显示SDDIM和Analytic-SDPM在不同时间步长下均表现出色,尤其是Analytic-SDPM在FID分数上达到了新的最优水平,证明了SDM在多种扩散求解器上的通用性和优越性。
相关信息
-
代码: https://github.com/andycao1125/sdm
-
论文: https://arxiv.org/abs/2408.16467v1