[Paper Review] 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 Computing4 citations

TL;DR

Spiking-Diffusion is the first fully spiking neural network (SNN)-based diffusion model, combining a vector-quantized spiking variational autoencoder (VQ-SVAE) for discrete latent representation with a spiking diffusion image decoder (SDID) for generative modeling. It achieves state-of-the-art FID scores of 37.50 (MNIST), 91.98 (FMNIST), 59.23 (KMNIST), 67.41 (Letters), and 120.5 (CIFAR-10), outperforming prior SNN-based generative models by up to 58.60% in FID reduction.

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}.

Motivation & Objective

• To bridge the gap in applying spiking neural networks (SNNs) to complex image generation tasks, particularly beyond classification.
• To develop a fully SNN-based diffusion model that leverages the energy efficiency of event-driven spiking dynamics for generative modeling.
• To overcome challenges in SNN-based VQ-VAEs, including memory-efficient spike sequence encoding and lossless feature reconstruction via adaptive spike generation.
• To enable high-fidelity image generation using discrete diffusion in the latent space, fully implemented with SNN layers.

Proposed method

• Proposes a vector-quantized spiking variational autoencoder (VQ-SVAE) that encodes images into discrete latent codes using both spike firing rate (SFR) and postsynaptic potential (PSP) for memory-efficient representation.
• Introduces a learnable weight operator $k$ to balance SFR and PSP in the codebook for optimal feature modeling.
• Designs an adaptive spike generator (ASG) to reconstruct embedded features as spike trains, minimizing information loss during decoding.
• Employs absorbing state Markov transition matrices to gradually mask discrete latent codes in the forward diffusion process.
• Constructs a spiking diffusion image decoder (SDID) using SNNs to denoise masked latent codes via an autoregressive process, learning the reverse denoising distribution.
• Uses a variational bound with $\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}})]]$ to stabilize training.

Experimental results

Research questions

• RQ1 Can a fully spiking neural network (SNN) architecture be used to implement a diffusion model for image generation?
• RQ2 How can spike sequences be effectively encoded and decoded in a VQ-VAE framework using SNNs without excessive memory or information loss?
• RQ3 Can absorbing state Markov chains in the discrete latent space enable high-quality image generation when combined with SNN-based denoisers?
• RQ4 Does the proposed Spiking-Diffusion model outperform existing SNN-based generative models in image reconstruction and generation quality?

Key findings

• Spiking-Diffusion achieves a Fréchet Inception Distance (FID) of 37.50 on MNIST, representing a 58.60% reduction compared to the previous state-of-the-art SNN-based model.
• On FMNIST, the model achieves an FID of 91.98, a 18.75% improvement over the prior SOTA model.
• For KMNIST, the FID score is 59.23, a 64.51% reduction compared to the previous best SNN-based approach.
• On the Letters dataset, Spiking-Diffusion achieves an FID of 67.41, a 29.75% improvement over the prior SOTA model.
• On CIFAR-10, the model achieves an FID of 120.5, a 44.88% reduction compared to the previous state-of-the-art SNN-based generative model.
• Visual samples show that Spiking-Diffusion generates sharper images with clearer boundaries than FSVAE, indicating better capture of semantic and structural details.