[論文レビュー] 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 は、離散的潜在表現のためのベクトル量子化スパiking変分オートエノード（VQ-SVAE）と、生成的モデリングのためのスパiking拡散画像デコーダー（SDID）を組み合わせた、最初の完全なスパikingニューラルネットワーク（SNN）ベースの拡散モデルである。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}.

研究の動機と目的

• スパikingニューラルネットワーク（SNN）を、分類を越えた複雑な画像生成タスクに応用する際のギャップを埋めるため。
• イベント駆動のスパイクダイナミクスのエネルギー効率を活用した、完全にSNNベースの拡散モデルを構築するため。
• SNNベースのVQ-VAEにおける課題、特にメモリ効率の良いスパイク系列符号化と、適応的スパイク生成による情報損失のない特徴再構築を克服するため。
• 完全にSNNレイヤーで実装された離散的拡散を用いて、高精細な画像生成を実現するため。

提案手法

• スパイク発火頻度（SFR）とシナプス後ポテンシャル（PSP）の両方を用いて、メモリ効率の良い表現で画像を離散的潜在コードに符号化する、ベクトル量子化スパキング変分オートエノード（VQ-SVAE）を提案する。
• 最適な特徴モデリングを実現するため、コードブックにおけるSFRとPSPのバランスを取るために学習可能な重み演算子 $k$ を導入する。
• 埋め込み特徴をスパイクトレインに再構築するための適応的スパイクジェネレータ（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 提案されたSpiking-Diffusionモデルは、既存のSNNベースの生成モデルと比較して、画像再構築および生成品質において優れているか？

主な発見

• Spiking-Diffusion は、MNIST で FID 37.50 を達成し、以前のSOTAモデル比で58.60%のFID低減を実現した。
• FMNIST では FID 91.98 を達成し、前回のSOTAモデル比で18.75%の改善を示した。
• KMNIST では FID スコアが 59.23 で、前回の最良SNNベースのアプローチ比で64.51%の低減を達成した。
• Letters データセットでは、Spiking-Diffusion が FID 67.41 を達成し、前回のSOTAモデル比で29.75%の改善を示した。
• CIFAR-10 では、モデルが FID 120.5 を達成し、前回のSOTA SNNベースの生成モデル比で44.88%の低減を実現した。
• 視覚的サンプルでは、Spiking-Diffusion が FSVAE よりもシャープで明確な境界を持つ画像を生成しており、意味的および構造的詳細をよりよく捉えていることが示された。