[논문 리뷰] Spiking-Diffusion: Vector Quantized Discrete Diffusion Model with Spiking Neural Networks

Mingxuan Liu, Gan, Jie외 4명|arXiv (Cornell University)|2023. 08. 20.

Advanced Memory and Neural Computing인용 수 4

한 줄 요약

스피킹-디퓨전은 완전히 스파iking 신경망(SNN) 기반의 디퓨전 모델로, 이산 잠재 표현을 위한 벡터 양자화 스파이킹 변동형 오토에코더(VQ-SVAE)와 생성 모델링을 위한 스파이킹 디퓨전 이미지 디코더(SDID)를 조합한다. 이 모델은 MNIST에서 FID 점수 37.50, FMNIST에서 91.98, KMNIST에서 59.23, Letters에서 67.41, CIFAR-10에서 120.5를 기록하며, 기존 SNN 기반 생성 모델 대비 최대 58.60%의 FID 감소를 달성하여 최신 기준 성능을 확립한다.

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 레이어로 구성된 이산 디퓨전을 통해 고해상도 이미지 생성을 가능하게 하기 위해.

제안 방법

• 스파이크 발화 빈도(SFR)와 postsynaptic potential(PSP)를 동시에 사용하여 이미지를 이산 잠재 코드로 메모리 효율적인 표현으로 인코딩하는 벡터 양자화 스파이킹 변동형 오토에코더(VQ-SVAE)를 제안한다.
• 코드북 내에서 SFR와 PSP의 균형을 맞추기 위해 학습 가능한 가중치 연산자 $k$를 도입하여 최적의 특징 모델링을 달성한다.
• 임bed된 특징을 스파이크 트레인으로 재구성하기 위해 정보 손실을 최소화하는 적응형 스파이크 생성기(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 SNN 기반 VQ-VAE 프레임워크에서 스파이크 시퀀스를 과도한 메모리 소비나 정보 손실 없이 효과적으로 인코딩하고 디코딩할 수 있는가?
• RQ3 이산 잠재 공간에서의 흡수 상태 마르코프 체인과 SNN 기반 디노이저를 조합하면 고품질의 이미지 생성을 가능하게 하는가?
• RQ4 제안된 스피킹-디퓨전 모델이 기존의 SNN 기반 생성 모델 대비 이미지 복원 및 생성 품질에서 뛰어나게 성능을 발휘하는가?

주요 결과

• 스피킹-디퓨전은 MNIST에서 FID 점수 37.50을 기록하며, 이는 이전 최신 기준 SNN 기반 모델 대비 58.60% 감소한 결과이다.
• FMNIST에서는 FID 점수 91.98을 기록하며, 이는 이전 최신 기준 모델 대비 18.75% 향상된 성능이다.
• KMNIST에서는 FID 점수 59.23을 기록하며, 이는 이전 최고의 SNN 기반 접근 대비 64.51% 감소한 결과이다.
• Letters 데이터셋에서는 FID 점수 67.41을 기록하며, 이는 이전 최신 기준 모델 대비 29.75% 향상된 성능이다.
• CIFAR-10에서는 FID 점수 120.5를 기록하며, 이는 이전 최신 기준 SNN 기반 생성 모델 대비 44.88% 감소한 결과이다.
• 시각적 샘플 분석 결과, 스피킹-디퓨전은 FSVAE보다 더 선명한 이미지와 뚜렷한 윤곽선을 생성하여 의미적이고 구조적인 세부 정보를 더 잘 포착하고 있음을 확인할 수 있다.