[论文解读] Fully Spiking Variational Autoencoder

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

Advanced Memory and Neural Computing被引用 4

一句话总结

本文提出全脉冲变分自编码器（FSVAE），这是首个端到端的脉冲神经网络（SNN）变分自编码器（VAE），其中所有组件——包括潜在空间——均使用SNN实现。通过引入自回归伯努利脉冲采样，该方法可在无需浮点运算的情况下，从潜在伯努利过程中实现随机采样，从而在神经形态硬件上实现高质量图像生成，且相比等效的人工神经网络（ANN）具有更高的能效和速度。

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

研究动机与目标

• 在能效受限的边缘设备上，利用脉冲神经网络（SNN）实现高质量图像生成。
• 解决在SNN中表示连续潜在变量的挑战，因为SNN需要使用二值化、事件驱动的脉冲序列。
• 开发一种完全基于SNN的变分自编码器（VAE），其生成质量与等效的人工神经网络（ANN）VAE相当或更优。
• 通过在采样和潜在空间建模中消除浮点运算，实现在神经形态硬件上的高效、低功耗推理。

提出的方法

• 提出自回归伯努利脉冲采样：利用自回归SNN对潜在脉冲序列的先验和后验分布进行建模，表示为伯努利过程。
• 采用基于SNN的编码器和解码器网络，通过二进制脉冲序列处理输入图像并实现重建。
• 在推理过程中，利用神经形态硬件上的随机数生成器，从伯努利分布的潜在脉冲序列中进行采样。
• 对MMD核应用后验偏移（PSP），以稳定训练并改善潜在空间的解耦性。
• 将潜在空间建模为一系列二进制脉冲事件，避免使用浮点运算，从而实现与神经形态设备的兼容性。
• 使用Kullback-Leibler（KLD）或最大均值差异（MMD）作为VAE损失函数中的正则化项进行模型训练。

实验结果

研究问题

• RQ1 能否仅使用SNN层构建一个完全的脉冲VAE，而无需任何混合ANN组件？
• RQ2 自回归伯努利脉冲采样是否能有效建模SNN兼容的潜在空间，并实现高质量图像生成？
• RQ3 所提出的FSVAE在标准基准测试中，其图像生成质量是否与等效的基于ANN的VAE相当或更优？
• RQ4 FSVAE的计算成本和推理速度与ANN VAE相比如何，尤其是在部署于神经形态硬件时？
• RQ5 基于离散脉冲的潜在变量是否能防止后验坍缩，从而生成更具意义且解耦的表征？

主要发现

• 在所有数据集（MNIST、FashionMNIST、CIFAR10、CelebA）上，FSVAE的Inception Score均优于等效的ANN VAE，表明其图像生成质量更高。
• 在MNIST和FashionMNIST上，FSVAE的Fréchet Inception Distance（FID）低于ANN VAE，表明其生成分布与真实图像更接近。
• 在CIFAR10上，FSVAE在所有指标（包括FID和Inception Score）上均表现更优，证明其具有强大的生成质量。
• 在CelebA上，当使用MMD损失并应用PSP于核函数时，FSVAE在所有配置中实现了最低的FID，证实其潜在空间质量得到提升。
• FSVAE的计算成本显著降低（乘法运算量为5.6×10⁸，而ANN为7.4×10⁹），且在神经形态硬件上SNN推理预计可快约100倍。
• 潜在表征的最佳时间步长为16，通道乘数k=20时FID最佳，表明对超参数选择具有鲁棒性。