回复
SDM: 第三代神经网络和扩散模型强强联合！FID最多超基线12倍，能耗省60%，实力SOTA！
论文链接：https://arxiv.org/pdf/2408.16467
代码链接：https://github.com/AndyCao1125/SDM
亮点直击
- 本文提出了脉冲扩散模型（Spiking Diffusion Model, SDM），一种高质量的图像生成器，在基于SNN的生成模型中实现了SOTA性能。
- 从生物学的角度出发，本文提出了一种时间脉冲机制（Temporal-wise Spiking Mechanism, TSM），使脉冲神经元能够捕捉更多的动态信息，从而提高去噪图像的质量。
- 大量结果显示，SDM在CIFAR-10数据集上的FID分数上超越了SNN基线模型多达12倍，同时节省了约60%的能耗。此外，本文提出了一种阈值引导策略，以进一步提高生成性能。
近年来，脉冲神经网络（Spiking Neural Networks, SNNs）因其超低能耗和高生物可塑性相比传统人工神经网络（Artificial Neural Networks, ANNs）而受到关注。尽管SNNs具有独特的优势，但其在计算密集型的图像生成领域的应用仍在探索之中。本文提出了脉冲扩散模型（Spiking Diffusion Models, SDMs），一种创新的基于SNN的生成模型家族，能够以显著降低的能耗生成高质量样本。特别地，本文提出了一种时间脉冲机制（Temporal-wise Spiking Mechanism, TSM），使SNNs能够从生物可塑性角度捕捉更多的时间特征。此外，本文提出了一种阈值引导策略，可以在不进行额外训练的情况下将性能提高多达16.7%。本文还首次尝试使用ANN-SNN方法进行基于SNN的生成任务。大量实验结果表明，本文的方法不仅在少量脉冲时间步中表现出与其ANN对应模型相当的性能，而且在很大程度上优于之前基于SNN的生成模型。此外，本文还展示了SDM在大规模数据集（例如LSUN卧室）上的高质量生成能力。这一发展标志着SNN基生成能力的一个重要进步，为未来实现低能耗和低延迟的生成应用开辟了新的研究途径。
方法
A. 峰值前残差学习
本文首先分析了先前脉冲神经网络（SNNs）中的残差学习方法存在的局限性和概念不一致性，特别是SEW ResNet，其公式可以表示为：
受 [61], [62] 的启发，本文在本文的脉冲UNet中采用了激活-卷积-批量归一化（Activation-Conv-BatchNorm）结构的预脉冲残差学习方法，解决了基于卷积的SNNs中梯度爆炸/消失和性能下降的双重挑战。通过预脉冲块，残差和输出通过浮点加法操作进行求和，确保在进入下一个脉冲神经元之前表示是准确的，同时避免了上述病态情况。整个预脉冲残差学习过程在一个残差块内可以表示如下：
总而言之，TSM允许膜电位在时间域内动态更新，从而提高捕捉潜在时间依赖特征的能力。后续实验表明，TSM机制优于传统的固定更新机制。
C. SDM 中的阈值指导
理论能耗计算
在本节中，本文描述了计算脉冲UNet架构理论能耗的方法。该计算包括两个主要步骤：确定架构中每个模块的突触操作（SOPs），然后基于这些操作估算整体能耗。脉冲UNet每个模块的突触操作可以如下量化：
通过 ANN-SNN 转换实施尖峰扩散模型
在本文中，本文首次利用ANN-SNN方法成功实现了SNN扩散。本文采用了Fast-SNN方法来构建量化ANN与SNN之间的转换。由于这一实现并不是本文论文的主要贡献，本文将简要描述ANN-SNN的原理，更多细节可以在[67]中找到。
实验
A. 实验设置
数据集和评估指标
为了展示所提算法的有效性和效率，本文在以下数据集上进行实验：32×32 MNIST、32×32 FashionMNIST 、32×32 CIFAR-10和 64×64 CelebA。定性结果根据Frechet Inception Distance（FID，越低越好）和Inception Score（IS，越高越好）进行比较。FID分数是通过比较50,000张生成图像与数据集的相应参考统计数据来计算的。
实现细节
B. 与最先进的方法比较
在下表 I 中，本文展示了本文的脉冲扩散模型（SDMs）与当前最先进的生成模型在无条件生成任务中的比较分析。本文还包括了ANN的结果作为参考。定性结果展示在下图4中。本文的结果表明，SDMs在所有数据集上均显著优于SNN基线，即使在较少的脉冲仿真步数（4/8）下也是如此。特别是，SDDPM在CelebA数据集上相比FSVAE和SGAD（两者均为16个时间步）分别有4倍和6倍的FID提升，而在CIFAR-10数据集上则有11倍和12倍的提升。正如预期的那样，随着时间步的增加，样本质量也会提高。本文还注意到，结合TSM后，性能进一步提升，而模型参数仅有微小增加（2e-4 M）。SDMs还可以处理快速采样求解器，并在较少的步骤中获得更高的采样质量（见下表VI）。重要的是，SDMs在使用相同的UNet架构下获得了与ANN基线相当的质量，甚至超过了一些ANN模型（例如，15.45 vs. 19.04）。这一结果突显了本文模型中使用的SNN的卓越表达能力。
C. 与 ANN-SNN 方法的比较
为了验证SDM在ANN-SNN方法下的生成能力，本文在32×32 CIFAR-10和64×64 FFHQ数据集上进行了实验。如下表II所示，ANN-SNN方法在CIFAR-10上表现出色（即51.18 FID），并且在微调策略后显著提高了图像质量（即29.53 FID）。然而，ANN-SNN的结果与直接训练的结果之间仍存在差距。尽管ANN-SNN方法在基于分类的任务中表现出与ANN相当的性能，但在生成任务方面仍缺乏深入的研究。ANN-SNN方法的定性结果展示在下图7中。
D. 时间性尖峰机制的有效性
为了更好地可视化TSM模块带来的性能提升，本文提供了使用SDDIM生成的CIFAR-10图像结果，分别展示了有和没有TSM模块的情况。这里本文使用DDIM而不是DDPM进行比较，因为DDIM基于常微分方程（ODEs）操作，确保了确定性和一致的生成结果。相比之下，DDPM依赖于随机微分方程（SDEs），在生成过程中引入了随机性，导致输出图像的可变性，从而使直接比较变得具有挑战性。
下图5中的结果显示，带有TSM模块的生成图像质量有显著提高。与没有TSM模块的图像相比，这些图像的轮廓更加清晰，背景更为清晰，纹理细节更丰富，从而证明了TSM的有效性。
E. 阈值指导的有效性
在前文中，本文提出了一种无需训练的方法：阈值引导（Threshold Guidance，TG），旨在通过在推理阶段仅略微调整脉冲神经元的阈值水平来提高生成图像的质量。如下表III所示，通过阈值调整进行抑制性引导显著提升了图像质量，在两个关键指标上都有所改善：FID分数从19.73下降到19.20，阈值减少0.3%；IS分数从7.44上升到7.55，阈值减少0.2%。相反，在某些条件下，兴奋性引导同样可以提高采样质量。这些发现强调了阈值引导作为一种在训练后显著提高模型效果的方法的潜力，而无需额外的训练资源。本文在附录中提供了更多关于阈值引导的解释。
F. TSM方法分析
G. 计算成本评估
为了进一步强调本文SDM的低能耗特性，本文对比分析了所提出的SDDPM与其对应的ANN模型在FID和能耗方面的表现。如表IV所示，当时间步长设定为4时，SDDPM的能耗显著降低，仅为其ANN对应模型的37.5%。此外，SDDPM的FID也提高了0.47，表明本文的模型在有效减少能耗的同时保持了竞争力的性能。当本文将分析扩展到不同的时间步长增量时，可以观察到一个明显的模式：随着时间步长的增加，FID分数有所改善，但代价是能耗的增加。这一观察结果表明，随着时间步长的增加，FID改善与能耗之间存在权衡。
H. 消融研究
不同组件对SDM的影响
本文首先在CIFAR-10数据集上进行消融研究，以探讨时间脉冲模块（TSM）和阈值引导（TG）的影响。如下表V所示，本文发现TSM和TG都对图像质量的提升有贡献。通过同时使用TSM和TG，本文获得了最佳的FID结果，相较于原始的SDDIM提升了18.4%。
SDM在不同求解器上的有效性
在下表VI中，本文验证了SDM在各种扩散求解器上的可行性和有效性。SDDIM在采样步骤上表现出更稳定的性能，而Analytic-SDPM展示了卓越的能力，达到了新的最先进性能，超越了ANN-DDIM的结果。总之，本文的SDM证明了其在处理任何扩散求解器方面的高效性，并且本文相信利用本文的SDM还有很大的潜力进一步提升FID性能。
讨论 & 结论
本研究提出了一种新的基于SNN的扩散模型家族，称为脉冲扩散模型（SDMs），它结合了SNN的能效优势和卓越的生成性能。SDMs在SNN基线中以更少的脉冲时间步长达到了最先进的结果，并且与ANNs相比，能耗更低。SDMs主要受益于两个方面：(1) 时间脉冲机制（TSM），它使去噪网络SNN-UNet的突触电流在每个时间步长中能够收集更多的动态信息，而不是像传统SNN那样由固定的突触权重控制；(2) 无需训练的阈值引导（TG），通过调整脉冲阈值进一步提高采样质量。
然而，本文工作的一个限制是SNN-UNet的时间步长相对较小，未能充分挖掘SDMs的全部潜力。此外，还应考虑在更高分辨率的数据集（如ImageNet）上进行测试。在未来的研究中，本文计划探索SDMs在生成领域的进一步应用，例如文本-图像生成，并尝试将其与先进的语言模型结合，以实现更有趣的任务。
本文转自 AI生成未来 ，作者：Jiahang Cao等
赞
收藏
回复
相关推荐
