[论文解读] Spiking-Diffusion: Vector Quantized Discrete Diffusion Model with Spiking Neural Networks

Mingxuan Liu, Gan, Jie+4 more|arXiv (Cornell University)|Aug 20, 2023

Advanced Memory and Neural Computing被引用 4

一句话总结

Spiking-Diffusion 是首个基于脉冲神经网络（SNN）的扩散模型，结合了用于离散潜在表征的向量量化脉冲变分自编码器（VQ-SVAE）与用于生成建模的脉冲扩散图像解码器（SDID）。其在 FID 分数上达到最先进水平：MNIST 为 37.50，FMNIST 为 91.98，KMNIST 为 59.23，Letters 为 67.41，CIFAR-10 为 120.5，相较于先前基于 SNN 的生成模型，FID 分数最高降低 58.60%。

ABSTRACT

Spiking neural networks (SNNs) have tremendous potential for energy-efficient neuromorphic chips due to their binary and event-driven architecture. SNNs have been primarily used in classification tasks, but limited exploration on image generation tasks. To fill the gap, we propose a Spiking-Diffusion model, which is based on the vector quantized discrete diffusion model. First, we develop a vector quantized variational autoencoder with SNNs (VQ-SVAE) to learn a discrete latent space for images. In VQ-SVAE, image features are encoded using both the spike firing rate and postsynaptic potential, and an adaptive spike generator is designed to restore embedding features in the form of spike trains. Next, we perform absorbing state diffusion in the discrete latent space and construct a spiking diffusion image decoder (SDID) with SNNs to denoise the image. Our work is the first to build the diffusion model entirely from SNN layers. Experimental results on MNIST, FMNIST, KMNIST, Letters, and Cifar10 demonstrate that Spiking-Diffusion outperforms the existing SNN-based generation model. We achieve FIDs of 37.50, 91.98, 59.23, 67.41, and 120.5 on the above datasets respectively, with reductions of 58.60\%, 18.75\%, 64.51\%, 29.75\%, and 44.88\% in FIDs compared with the state-of-art work. Our code will be available at \url{https://github.com/Arktis2022/Spiking-Diffusion}.

研究动机与目标

• 为弥合脉冲神经网络（SNN）在复杂图像生成任务中的应用差距，特别是超越分类任务的限制。
• 开发一种完全基于 SNN 的扩散模型，利用事件驱动脉冲动态的能效优势进行生成建模。
• 解决基于 SNN 的 VQ-VAE 所面临的挑战，包括内存高效的脉冲序列编码以及通过自适应脉冲生成实现无损特征重建。
• 通过在潜在空间中实现离散扩散，完全使用 SNN 层实现高保真图像生成。

提出的方法

• 提出一种向量量化脉冲变分自编码器（VQ-SVAE），通过脉冲发放率（SFR）和突触后电位（PSP）联合编码图像为离散潜在码，实现内存高效的表征。
• 引入可学习的权重算子 $k$，在码书内平衡 SFR 与 PSP，以实现最优特征建模。
• 设计自适应脉冲生成器（ASG），将嵌入特征重建为脉冲序列，最大限度减少解码过程中的信息损失。
• 采用吸收态马尔可夫转移矩阵，在前向扩散过程中逐步掩码离散潜在码。
• 构建基于 SNN 的脉冲扩散图像解码器（SDID），通过自回归过程去噪掩码后的潜在码，学习反向去噪分布。
• 使用变分界 $\mathbb{E}_{q(\bm{h}_{0})}[\sum_{t_{d}=1}^{T_{d}}\frac{1}{t_{d}}\mathbb{E}_{q(\bm{h}_{t_{d}}|\bm{h}_{0})}[\sum_{\bm{h}_{t_{d}}(i,j)=m}\log{p_{\theta}}(\bm{h}_{0}(i,j)|\bm{h}_{t_{d}})]]$ 稳定训练过程。

实验结果

研究问题

• RQ1 是否可以使用完全脉冲神经网络（SNN）架构实现用于图像生成的扩散模型？
• RQ2 在 VQ-VAE 框架中，如何在不造成过多内存消耗或信息损失的前提下，有效编码和解码脉冲序列？
• RQ3 在离散潜在空间中，结合 SNN 基础去噪器时，吸收态马尔可夫链是否能够实现高质量图像生成？
• RQ4 所提出的 Spiking-Diffusion 模型是否在图像重建与生成质量方面优于现有的基于 SNN 的生成模型？

主要发现

• Spiking-Diffusion 在 MNIST 上实现 37.50 的 Fréchet Inception Distance（FID）分数，相较于先前最先进（SOTA）的基于 SNN 的模型，FID 分数降低 58.60%。
• 在 FMNIST 上，该模型实现 91.98 的 FID 分数，相较于先前 SOTA 模型提升 18.75%。
• 在 KMNIST 上，FID 得分为 59.23，相较于先前最佳的基于 SNN 的方法降低 64.51%。
• 在 Letters 数据集上，Spiking-Diffusion 实现 67.41 的 FID 分数，相较于先前 SOTA 模型提升 29.75%。
• 在 CIFAR-10 上，模型实现 120.5 的 FID 分数，相较于先前最先进（SOTA）的基于 SNN 的生成模型降低 44.88%。
• 视觉样本显示，Spiking-Diffusion 生成的图像比 FSVAE 更为清晰，边界更分明，表明其对语义与结构细节的捕捉能力更强。