[论文解读] Spikeformer: A Novel Architecture for Training High-Performance Low-Latency Spiking Neural Network

Yudong Li, Yunlin Lei+1 more|arXiv (Cornell University)|Nov 19, 2022

Advanced Memory and Neural Computing被引用 11

一句话总结

该论文提出Spikeformer，一种新型基于Transformer的脉冲神经网络架构，在脉冲神经网络（DVS-Gesture，DVS-CIFAR10）和静态图像（ImageNet）数据集上均实现了最先进性能，且仿真时间步数极少。通过引入卷积标记化模块（CT）以解决数据饥渴和训练不稳定性问题，并采用时空注意力（STA）机制联合建模空间与时间依赖关系，Spikeformer在ImageNet上比其人工神经网络（ANN）基线模型高出2.2%，在DVS-Gesture上高出3.1%，证明其在低延迟、高性能脉冲神经网络训练方面的优越性。

ABSTRACT

Spiking neural networks (SNNs) have made great progress on both performance and efficiency over the last few years,but their unique working pattern makes it hard to train a high-performance low-latency SNN.Thus the development of SNNs still lags behind traditional artificial neural networks (ANNs).To compensate this gap,many extraordinary works have been proposed.Nevertheless,these works are mainly based on the same kind of network structure (i.e.CNN) and their performance is worse than their ANN counterparts,which limits the applications of SNNs.To this end,we propose a novel Transformer-based SNN,termed "Spikeformer",which outperforms its ANN counterpart on both static dataset and neuromorphic dataset and may be an alternative architecture to CNN for training high-performance SNNs.First,to deal with the problem of "data hungry" and the unstable training period exhibited in the vanilla model,we design the Convolutional Tokenizer (CT) module,which improves the accuracy of the original model on DVS-Gesture by more than 16%.Besides,in order to better incorporate the attention mechanism inside Transformer and the spatio-temporal information inherent to SNN,we adopt spatio-temporal attention (STA) instead of spatial-wise or temporal-wise attention.With our proposed method,we achieve competitive or state-of-the-art (SOTA) SNN performance on DVS-CIFAR10,DVS-Gesture,and ImageNet datasets with the least simulation time steps (i.e.low latency).Remarkably,our Spikeformer outperforms other SNNs on ImageNet by a large margin (i.e.more than 5%) and even outperforms its ANN counterpart by 3.1% and 2.2% on DVS-Gesture and ImageNet respectively,indicating that Spikeformer is a promising architecture for training large-scale SNNs and may be more suitable for SNNs compared to CNN.We believe that this work shall keep the development of SNNs in step with ANNs as much as possible.Code will be available.

研究动机与目标

• 解决现有基于代理梯度或转换方法训练的脉冲神经网络（SNN）在性能与延迟方面的局限性。
• 克服将视觉Transformer（ViT）适配至SNN时出现的数据饥渴与训练不稳定性问题。
• 设计一种新型SNN架构，使其在ImageNet等大规模数据集上的表现优于基于CNN的SNN，并与或超越人工神经网络（ANN）性能。
• 实现SNN的直接端到端训练，仅需极少仿真时间步，同时保持高准确率与低能耗。

提出的方法

• 引入卷积标记化（CT）模块以替代ViT中的标准图像块化主干网络，提升SNN的数据效率与训练稳定性。
• 设计一种时空注意力（STA）机制，联合建模脉冲序列中的空间与时间依赖关系，区别于以往仅使用空间或时间注意力的方法。
• 采用代理梯度反向传播进行SNN的端到端训练，实现低延迟，避免依赖人工神经网络到脉冲神经网络的转换过程。
• 采用“200个标记”规则进行网络架构设计，确保足够长的序列长度以实现稳定训练，同时保持高效性。
• 通过消融实验优化网络深度与时间步数，验证其在不同配置下的鲁棒性。
• 在初始设计中采用轻量化、非残差结构，以优先突出Transformer模块与CT模块的影响。

实验结果

研究问题

• RQ1 基于Transformer的SNN架构是否能在脉冲神经网络与静态图像数据集上均实现竞争性或最先进性能，且具备低延迟？
• RQ2 当将视觉Transformer应用于SNN时，如何缓解其数据饥渴与训练不稳定性问题？
• RQ3 SNN中联合使用时空注意力是否能带来优于仅使用空间或时间注意力的表征学习效果？
• RQ4 直接训练的SNN是否能在ImageNet等大规模基准测试中超越其人工神经网络（ANN）基线？
• RQ5 哪些架构设计原则（如标记化方式、序列长度）对SNN的稳定与高性能训练至关重要？

主要发现

• 在DVS-Gesture数据集上，Spikeformer仅用4个时间步即达到98.96%的top-1准确率，比之前最先进方法高出1.04%，且延迟相同。
• 在DVS-CIFAR10数据集上，Spikeformer在4个时间步下达到85.21%的准确率，尽管时间步数与参数量更少，仍比之前最先进方法高出16.6%。
• 在ImageNet数据集上，Spikeformer在4个时间步与38.75M参数下达到78.31%的top-1准确率，比基于代理梯度或转换方法训练的最先进SNN高出5%以上。
• Spikeformer在ImageNet上比其ANN基线高出2.2%，在DVS-Gesture上高出3.1%，首次实现直接训练的SNN在这些基准上超越其ANN等效模型。
• 与原始ViT结构的SNN相比，卷积标记化（CT）模块使DVS-Gesture的准确率提升超过16%，显著稳定了训练过程并实现了收敛。
• 消融实验验证了CT模块的关键作用，以及Spikeformer在不同深度与时间步配置下的鲁棒性。