[Paper Review] Fully Spiking Variational Autoencoder

Hiromichi Kamata, Yusuke Mukuta+1 more|arXiv (Cornell University)|Sep 26, 2021

Advanced Memory and Neural Computing4 citations

TL;DR

This paper proposes the Fully Spiking Variational Autoencoder (FSVAE), the first end-to-end spiking neural network (SNN) VAE where all components—including the latent space—are implemented using SNNs. By introducing autoregressive Bernoulli spike sampling, the method enables stochastic sampling from a latent Bernoulli process without floating-point operations, allowing high-quality image generation on neuromorphic hardware with superior energy efficiency and speed compared to equivalent ANNs.

ABSTRACT

Spiking neural networks (SNNs) can be run on neuromorphic devices with ultra-high speed and ultra-low energy consumption because of their binary and event-driven nature. Therefore, SNNs are expected to have various applications, including as generative models being running on edge devices to create high-quality images. In this study, we build a variational autoencoder (VAE) with SNN to enable image generation. VAE is known for its stability among generative models; recently, its quality advanced. In vanilla VAE, the latent space is represented as a normal distribution, and floating-point calculations are required in sampling. However, this is not possible in SNNs because all features must be binary time series data. Therefore, we constructed the latent space with an autoregressive SNN model, and randomly selected samples from its output to sample the latent variables. This allows the latent variables to follow the Bernoulli process and allows variational learning. Thus, we build the Fully Spiking Variational Autoencoder where all modules are constructed with SNN. To the best of our knowledge, we are the first to build a VAE only with SNN layers. We experimented with several datasets, and confirmed that it can generate images with the same or better quality compared to conventional ANNs. The code is available at https://github.com/kamata1729/FullySpikingVAE

Motivation & Objective

• To enable high-quality image generation using spiking neural networks (SNNs) on energy-constrained edge devices.
• To overcome the challenge of representing continuous latent variables in SNNs, which require binary, event-driven spike trains.
• To develop a fully SNN-based variational autoencoder (VAE) that maintains generative quality comparable or superior to equivalent artificial neural network (ANN) VAEs.
• To enable efficient, low-power inference on neuromorphic hardware by eliminating floating-point operations in sampling and latent space modeling.

Proposed method

• Proposes autoregressive Bernoulli spike sampling: uses autoregressive SNNs to model prior and posterior distributions over latent spike trains as Bernoulli processes.
• Employs SNN-based encoder and decoder networks to process input images and reconstruct them using binary spike trains.
• Uses a random number generator on neuromorphic hardware to sample from the Bernoulli-distributed latent spike trains during inference.
• Applies a posterior shift (PSP) to the MMD kernel to stabilize training and improve disentanglement in the latent space.
• Models the latent space as a sequence of binary spike events, avoiding floating-point operations and enabling compatibility with neuromorphic devices.
• Trains the model using either Kullback-Leibler (KLD) or Maximum Mean Discrepancy (MMD) as the regularizer in the VAE loss function.

Experimental results

Research questions

• RQ1 Can a fully spiking VAE be constructed using only SNN layers, without any hybrid ANN components?
• RQ2 Can autoregressive Bernoulli spike sampling effectively model the latent space in a way compatible with SNNs and enable high-quality image generation?
• RQ3 Does the proposed FSVAE achieve image generation quality comparable or superior to equivalent ANN-based VAEs on standard benchmarks?
• RQ4 How does the computational cost and inference speed of FSVAE compare to ANN VAEs, especially when deployed on neuromorphic hardware?
• RQ5 Can discrete spike-based latent variables prevent posterior collapse, leading to more meaningful and disentangled representations?

Key findings

• FSVAE outperforms equivalent ANN VAEs in inception score across all datasets (MNIST, FashionMNIST, CIFAR10, CelebA), indicating higher-quality image generation.
• On MNIST and FashionMNIST, FSVAE achieves lower Fréchet Inception Distance (FID) than the ANN VAE, indicating better distributional similarity to real images.
• For CIFAR10, FSVAE achieves superior performance in all metrics, including FID and inception score, demonstrating strong generative quality.
• On CelebA, FSVAE achieves the lowest FID among all configurations when using MMD loss with PSP applied to the kernel, confirming improved latent space quality.
• The computational cost of FSVAE is significantly reduced in multiplications (5.6×10⁸ vs. 7.4×10⁹ in ANN), and SNN inference is expected to be ~100× faster on neuromorphic hardware.
• The optimal timestep for latent representation was found to be 16, and the best FID was achieved with k=20 for the channel multiplier, indicating robustness to hyperparameter choice.