Spiking-Diffusion: Vector Quantized Discrete Diffusion Model with
  Spiking Neural Networks
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}.
Mingxuan Liu, Jie Gan, Rui Wen, Tao Li, Yongli Chen, Hong Chen
# Spiking-Diffusion: Vector Quantized Discrete Diffusion Model with Spiking Neural Networks

###### 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 [https://github.com/Arktis2022/Spiking-Diffusion](https://github.com/Arktis2022/Spiking-Diffusion).

Mingxuan Liu\({}^{\star}\)  Jie Gan\({}^{\dagger}\)  Rui Wen\({}^{\star}\)  Tao Li\({}^{\star}\)  Yongli Chen\({}^{\dagger}\)  Hong Chen\({}^{\star}\)\({}^{\star}\) Tsinghua University

\({}^{\dagger}\) Beijing Smartchip Microelectronics Technology Co., Ltd Spiking neural networks, Diffusion model, Image generation, Bionic Learning

## 1 Introduction

Spiking neural networks (SNNs) are the third generation neural networks with biological plausibility that encode and transmit information in form of spikes by mimicking the dynamics of neurons in the brain. Compared with Artificial neural networks (ANNs), the event-driven nature of SNNs allows for significant reduction in energy consumption when running on neuromorphic chips [1].

SNNs with deep learning techniques have shown promising results on simple tasks such as image classification, image segmentation, and optical flow estimation[2]. However, the scope of their utility in complex tasks, particularly image generation, is still confined. Spiking-GAN [3] is the first fully SNN-based GAN that uses time-to-first-spike (TTFS) encoding to generate MNIST images through adversarial training. Despite its significant progress, the quality of the generated images still falls short of ANN-based GAN. Another model, fully spiking variational autoencoder (FSVAE) [4], is a VAE model constructed entirely from SNN layers. This model produces comparable results to those of ANN-based VAEs, but its performance is still constrained by the inherent limitations of VAEs, namely the weak generative capacity due to the overly simplified decoder and compressed latent space [5]. Thus, more exploration and development are needed to harness the full potential of SNNs in image generation tasks.

In this paper, we propose a vector quantized discrete diffusion model with spiking neural networks (called Spiking-Diffusion), which is the first diffusion model using only SNN layers. The pipeline of proposed Spiking-Diffusion contains two stages. First, an image is transformed into a discrete matrix through proposed vector quantized spiking variational autoencoder (VQ-SVAE). However, creating VQ-VAEs [6] in SNNs brings two challenges. The one is converting spike sequences into dense features suitable for storage in the codebook, because storing spike sequences directly within the codebook will consume too much memory. To overcome this challenge, we combine spike frequency rate (SFR) and postsynaptic potential (PSP) to model the spike sequences and introduce the learnable operator \(k\) to optimize the weights of PSP and SFR. The other challenge is losslessly converting embedded features to spiking input for SNN decoder, because Poisson coding incurs information loss. To address this challenge, we design an adaptive spike generator (ASG) before the decoder and train it by using the same dictionary learning algorithm as traditional VQ-VAE. In the second stage, we utilize a spiking diffusion image decoder (SDID) to fit the prior discrete latent codes. In the forward process, we select a Markov transition matrix with an absorbing state [7] to gradually add masks to the discrete
matrix, and adopt an SNN denoising network SDID to recover the masks and achieve inverse process parameterization. Experiments on MNIST [8], FMNIST [9], KMNIST [10], Letters [11], and Cifar10 [12] show that Spiking-Diffusion outperforms the SOTA SNN-based generative model.

## 2 Methodology

### SNNs Learning Algorithms

In this study, we adopt he SNNs learning algorithms in [13]. Spike neurons in SNNs emulates the behavior of biological neurons by generating discrete pulses, or spikes, in response to input stimuli. The LIF neuron model [14] is used, which is described by the following dynamic equation:

\[H[t]=V[t-1]+\frac{1}{\tau}(X[t]-(V[t-1]-V_{reset})) \tag{1}\]

\[S[t]=\Theta(H[t]-V_{th}) \tag{2}\]

\[V[t]=H[t](1-S[t])+V_{reset}S[t] \tag{3}\]

where \(\tau\) represents the membrane time constant, \(X[t]\) denotes the synaptic input current at time step \(t\), \(H[t]\) represents the membrane potential of a neuron after charging but before firing a spike, and \(S[t]\) denotes the spike generated by a LIF neuron when its membrane potential exceeds the discharge threshold \(V_{th}\). \(\Theta(v)\) is the Heaviside step function, which is equal to 1 when \(v\geq 0\) and 0 otherwise. \(V[t]\) stands for the membrane potential after a spike event, which is reset to the \(V_{reset}\) using a hard reset [15]. If no spike is generated, it is equal to \(H[t]\).

### Vq-Svae

Although diffusion models [16] can generate high-quality images, they need to predict the noise added in the forward step. However, SNNs are difficult to fit and process continuous analog signals due to their discrete spike encoding [17]. In order to solve the problem, we propose a VQ-SVAE to learn a discrete latent representation of images.

As shown in Fig. 1, Given a 2D image \(x\in\mathbb{R}^{1\times H\times W}\), we first convert it into sequences \(x_{1:T_{x}}\in\mathbb{R}^{T_{x}\times 1\times H\times W}\) using direct input encoding [18], and then the \(x_{1:T_{x}}\) is encoded as a spike sequence \(z_{1:T_{x}}^{E}\in\ \{0,1\}^{T_{x}\times C\times H^{\prime}\times W^{\prime}}\) after the SNN encoder. In order to prevent the codebook from consuming too much memory, we utilize spike firing rate (SFR) and postsynaptic potential (PSP) to model the spike sequence. SFR has been proven in neurobiology to have the ability to represent information. For example, auditory nerves use SFR to encode steady-state vowels [19].

\[\mathrm{SFR}(z_{1:T_{x}}^{E})=\frac{1}{T_{s}}\sum_{t=1}^{T_{s}}z_{t}^{E} \tag{4}\]

PSP simulates the response of postsynaptic neurons to the action potential (AP) sequence. Similarly, studies have shown that PSP is related to experience-dependent plasticity in the mammalian nervous system [20]. We use the following formula to update the PSP [21]:

\[\mathrm{PSP}(z_{\leq t_{s}}^{E})=(1-\frac{1}{\tau_{syn}})\times\mathrm{PSP} (z_{\leq t_{s-1}}^{E})+\frac{1}{\tau_{syn}}\times z_{t}^{E} \tag{5}\]

where \(\tau_{syn}\) is the synaptic time constant and \(\mathrm{PSP}(z_{\leq 0}^{E})\) is set to 0. We use a trainable operator \(k\) to allocate the weights of the PSP and SFR, so the features used for quantized encoding \(z^{E}\in\ \mathbb{R}^{C\times H^{\prime}\times W^{\prime}}\) can be obtained from the following formula:

\[z^{E}=k\times\mathrm{SFR}(z_{1:T_{x}}^{E})+(1-k)\times\mathrm{PSP}(z_{\leq T_ {x}}^{E}) \tag{6}\]

After obtaining \(z^{E}\), we construct a codebook \(\mathbb{Z}\in\mathbb{R}^{K\times C}\), with \(z_{k}\in\mathbb{Z}\) and \(z_{i:j}^{E}\in z^{E}\). The quantized encoding of \(z^{E}\), denoted as a discrete matrix \(h_{0}\), can then be derived through nearest neighbor search. According to \(h_{0}\), each spatial feature \(z_{i,j}^{E}\) is mapped to its nearest codebook entry \(z_{k}\) by indexing in \(\mathbb{Z}\), thereby obtaining the embedding feature \(z^{Q}\):

\[h_{0}=Q(z^{E},\mathbb{Z})=\mathrm{argmin}_{k}||z_{i,j}^{E}-z_{k}||_{2} \tag{7}\]

\[z_{Q}=\mathrm{Index}(\mathbb{Z},h_{0}) \tag{8}\]

However, the embedding feature \(z^{Q}\) incorporates the SFR and PSP features of spike sequences. Consequently, it cannot be Poisson encoded [21] as input to the SNN decoder without loss of information. To solve the problem, we use a SNN layer to construct a trainable adaptive spike generator (ASG) that learns to restore \(z^{Q}\) to the original spike sequence \(z_{1:T_{x}}^{Q}\). At this stage, all parameters in the network can be trained end-to-end with the following loss function:

\[\mathcal{L}_{\mathrm{VQ-SVAE}}=||x-\hat{x}||_{2}+||z^{E}-z^{Q}||_{2}^{2} \tag{9}\] \[+(\sum_{t=1}^{T}||\mathrm{PSP}(sg[z_{\leq t_{s}}^{E}])-\mathrm{PSP }(z_{\leq t_{s}}^{Q})||^{2}\] \[+\beta\sum_{t=1}^{T}||\mathrm{PSP}(sg[z_{\leq t_{s}}^{Q}])- \mathrm{PSP}(z_{\leq t_{s}}^{E})||^{2})\]

In the formula, \(sg[\cdot]\) denotes the stop gradient operation. Notably, compared with VQ-VAE [6], VQ-SVAE imposes an additional constraint on training the ASG, which corresponds to the third term above. Moreover, the straight-through estimator propagates the gradients from the spike sequences \(z_{1:T_{x}}^{Q}\) to \(z_{1:T_{x}}^{E}\) during backpropagation, rather than propagating the gradients of the embedded features \(z^{Q}\). Since spike sequences exhibit sparsity, directly computing the mean squared error (MSE) loss between sequences is not the optimal distance metric, therefore, we set the third term in Eq. (9) to maximum mean discrepancy (MMD), whose kernel function is PSP.
### Absorbing State Diffusion

The training dataset is encoded into discrete matrices \(h_{0}\) by the VQ-SVAE to train the spiking diffusion image decoder (SDID). For the forward diffusion process \(q(h_{t_{d}}|h_{t_{d}-1})\), we use a transition matrix with an absorbing state [7] to convert all the elements in the discrete matrix into a mask after a fixed number of \(T_{d}\) time steps. In detail, for each discrete random variable element \(h_{t_{d}}(i,j)\) in \(h_{t_{d}}\) with \(K\) categories, denoting the one-hot version of \(h\) with the row vector \(\mathbf{h}\), we can define the forward process as:

\[q(\mathbf{h}_{t_{d}}|\mathbf{h}_{t_{d}-1})=\mathrm{Cat}(\mathbf{h}_{t_{d}};\mathbf{p}=\mathbf{h}_{t _{d}-1}\mathbf{Q}_{t_{d}}) \tag{10}\]

where \(\mathrm{Cat}(\mathbf{h};\mathbf{p})\) is a categorical distribution parameterized by \(\mathbf{p}\), and \([\mathbf{Q}_{t_{d}}]_{ij}=q(h_{t_{d}}=j|h_{t_{d}-1}=i)\) is the transition matrix of the forward Markov chain. We adopt absorbing state diffusion which is compatible with SNNs. Its transition matrix \(Q_{t_{d}}\in\mathbb{R}^{(K+1)\times(K+1)}\) can be expressed as follows:

\[Q_{t_{d}}=\begin{bmatrix}1-\gamma_{t_{d}}&0&\cdots&\gamma_{t_{d}}\\ 0&1-\gamma_{t_{d}}&\cdots&\gamma_{t_{d}}\\ \vdots&\vdots&\ddots&\vdots\\ 0&0&\cdots&1\end{bmatrix} \tag{11}\]

\[\gamma_{t_{d}}=\frac{1}{T_{d}-t_{d}+1} \tag{12}\]

The Eq. (11) indicates that at each timestep, each regular token has probability \(\gamma_{t_{d}}\) of being replaced by the [MASK] token, and probability \(1-\gamma_{t_{d}}\) of remaining unchanged, but the [MASK] token remains unchanged. In our work, we set the [MASK] token as \(K\), as a result, \(\mathbf{Q}_{t_{d}}\in\mathbb{R}^{(K+1)\times(K+1)}\) and each token possesses \(K+1\) states. Therefore, the transition matrix with the absorbing state is non-zero only on the diagonal and the last column, so the sparsity of transition matrix will help to reduce the computational cost of the forward process.

The reverse process is defined as a Markov chain parameterized by \(\theta\), which gradually denoises from the data distribution \(p_{\theta}(\mathbf{h}_{0:T_{d}})=p(\mathbf{h}_{T_{d}})\prod_{d=1}^{T_{d}}p_{\theta}( \mathbf{h}_{t_{d}-1}|\mathbf{h}_{t_{d}})\) and is optimized by the evidence lower bound (ELBO). With \(t_{d}^{\mathrm{th}}\) term the loss can be written as:

\[\mathcal{L}_{t_{d}}=D_{\mathrm{KL}}(q(\mathbf{h}_{t_{d}-1}|\mathbf{h}_{0})||p_{\theta} (\mathbf{h}_{t_{d}-1}|\mathbf{h}_{t_{d}})) \tag{13}\]

To reduce the randomness of training, SDID \(S(h_{t_{d}},t_{d})\) is employed to predict \(p_{\theta}(\mathbf{h}_{0}|\mathbf{h}_{t_{d}})\) rather than learn \(p_{\theta}(\mathbf{h}_{t_{d}-1}|\mathbf{h}_{t_{d}})\) directly. As a result, variational bound reduces to:

\[\mathbb{E}_{q(\mathbf{h}_{0})}\Bigg{[}\sum_{t_{d}=1}^{T_{d}}\frac{1}{t}\mathbb{E} _{q(\mathbf{h}_{t_{d}}|\mathbf{h}_{0})}\Big{[}\sum_{\mathbf{h}_{t_{d}}(i,j)=m}\log p_{ \theta}(\mathbf{h}_{0}(i,j)|\mathbf{h}_{t_{d}})\Big{]}\Bigg{]} \tag{14}\]

Figure 1: The training process of Spiking-Diffusion consists of two stages. Step 1: Compress the images into discrete variables through VQ-SVAE. Step 2: The spiking diffusion image decoder (SDID) models the discrete latent space by reversing the forward diffusion process, which gradually adds masks in the discrete matrix through a fixed Markov chain. Step 3: During the test process, the SDID lifts the masks through an autoregressive process to obtain a discrete matrix with the target distribution.
## 3 Experiments

### Datasets

We test our model on five datasets: The MNIST [8] consists of 60,000 handwritten digit images divided into 10 classes (digits 0 to 9); FMNIST [9] contains 60,000 images of 10 different classes of clothing; KMNIST [10] consists of 60,000 images of Japanese characters, with each image belonging to one of 10 classes; Letters [11] provides 145,000 images of 26 Latin letters; for Cifar10[12], 50,000 images are used for training and 10,000 images for evaluation.

### Implementation details

We make a comprehensive comparison between Spiking-Diffusion and F5VAE [4], which is the state-of-the-art SNN-based generative model generating images with quality equal to or better than ANN-based models. We test the F5VAE and Spiking-Diffusion on all datasets with the same experimental setup. The LIF neuron model's surrogate gradient function is \(g(x)=\frac{1}{\pi}\mathrm{arctan}(\frac{1}{\pi}\alpha x)+\frac{1}{2}\), and its derivative is \(g^{\prime}(x)=\frac{\alpha}{2(1+\frac{\alpha}{2}\alpha x)^{2}}\), where \(\alpha\) represents the slope parameter. For all neurons, \(\alpha=2\), \(V_{reset}=0\), and \(V_{th}=1\). The Optimizer is AdamW with \(\beta=(0.9,0.999)\) and weight decay 0.001. The training lasts for 100 epochs. In addition, the original F5VAE used the STBP-tdBN [22] framework to train SNNs, while we adopt SpikingJelly [13] in our work.

### Evaluation Metrics and Results

**Evaluation Metrics.** We evaluate the reconstruction ability of the model by comparing the input and output images using two metrics: mean squared error (MSE) loss and structural similarity (SSIM) loss. To assess the quality of sampled images, we use two commonly used scores: Kernel Inception Distance (KID) and Frechet inception distance (FID).

**Results.** Table 1 shows that our Spiking-Diffusion model outperforms F5VAE in reconstruction and generation with the SpikingJelly framework. Spiking-Diffusion yields lower MSE and SSIM losses of reconstruction quality, and lower FID and KID scores of generated image quality. Figure 2 shows the examples of generated images on the four datasets, from which we find that with the same SNN encoder and decoder structures, the images generated by F5VAE are blurrier, while the images generated by our Spiking-Diffusion are clearer with obvious boundaries. The is because that only the Bernoulli distribution is used in F5VAE to approximate the true posterior distribution, which cannot fully capture the rich semantic information in the image. In contrast, the Spiking-Diffusion can approximate the complex posterior distribution through nearest neighbor search and generate images by seeing context.

## 4 Conclusion

In this work, we present Spiking-Diffusion, the first implementation of diffusion models in SNNs. Our approach is based on VQ-DDM, which has two stages. First, we train the VQ-SVAE model to achieve image reconstruction, by using both PSP and SFR to model spiking features and constructing a codebook for image discretization. Secondly, we construct a discrete diffusion model in the discrete feature domain of images, and using absorbing diffusion state in Spiking-Diffusion is proven to generate higher quality images. The experimental results demonstrate that Spiking-Diffusion currently outperforms other SNN-based generative models. Future work will focus on how to train larger scale SNN generative models.

\begin{table}
\begin{tabular}{l l l l l l} \hline \hline Dataset & Model & MSE \(\downarrow\) & SSIM \(\downarrow\) & KID \(\downarrow\) & FID \(\downarrow\) \\ \hline \multirow{2}{*}{MNIST} & F5VAE & 0.023 & 0.219 & 0.054 & 90.57 \\  & **Ours** & **0.006** & **0.077** & **0.018** & **37.50** \\ \hline \multirow{2}{*}{FMNIST} & F5VAE & 0.023 & 0.378 & 0.070 & 113.2 \\  & **Ours** & **0.011** & **0.233** & **0.055** & **91.98** \\ \hline \multirow{2}{*}{KMNIST} & F5VAE & 0.054 & 0.452 & 0.272 & 166.9 \\  & **Ours** & **0.014** & **0.151** & **0.068** & **59.23** \\ \hline \multirow{2}{*}{Letters} & F5VAE & 0.015 & 0.180 & 0.140 & 95.96 \\  & **Ours** & **0.005** & **0.078** & **0.075** & **67.41** \\ \hline \multirow{2}{*}{Cifar10} & F5VAE & 0.024 & 0.689 & 0.181 & 218.6 \\  & **Ours** & **0.009** & **0.399** & **0.121** & **120.5** \\ \hline \hline \end{tabular}
\end{table}
Table 1: Performance comparisons of Spiking-Diffusion(Ours) and F5VAE on SpikingJelly [13].

Figure 2: Generated images of F5VAE and our Spiking-Diffusion (From top to down: MNIST, FMNIST, KMNIST, and Letters).