[Paper Review] Spikformer: When Spiking Neural Network Meets Transformer

Zhaokun Zhou, Yuesheng Zhu+5 more|arXiv (Cornell University)|Sep 29, 2022

Advanced Memory and Neural ComputingEngineering103 citations

TL;DR

This paper introduces Spiking Self Attention (SSA) and a Spiking Transformer (Spikformer) that operates directly on spike-based data, achieving state-of-the-art results among directly trained SNNs on ImageNet and competitive results on neuromorphic datasets while remaining energy-efficient.

ABSTRACT

We consider two biologically plausible structures, the Spiking Neural Network (SNN) and the self-attention mechanism. The former offers an energy-efficient and event-driven paradigm for deep learning, while the latter has the ability to capture feature dependencies, enabling Transformer to achieve good performance. It is intuitively promising to explore the marriage between them. In this paper, we consider leveraging both self-attention capability and biological properties of SNNs, and propose a novel Spiking Self Attention (SSA) as well as a powerful framework, named Spiking Transformer (Spikformer). The SSA mechanism in Spikformer models the sparse visual feature by using spike-form Query, Key, and Value without softmax. Since its computation is sparse and avoids multiplication, SSA is efficient and has low computational energy consumption. It is shown that Spikformer with SSA can outperform the state-of-the-art SNNs-like frameworks in image classification on both neuromorphic and static datasets. Spikformer (66.3M parameters) with comparable size to SEW-ResNet-152 (60.2M,69.26%) can achieve 74.81% top1 accuracy on ImageNet using 4 time steps, which is the state-of-the-art in directly trained SNNs models.

Motivation & Objective

• Motivate the integration of self-attention with spiking neural networks (SNNs) to leverage both energy efficiency and global feature modeling.
• Design a self-attention variant (SSA) compatible with spike-based computation that avoids multiplications and softmax.
• Build a Spiking Transformer (Spikformer) built from SSA and verify its effectiveness on static and neuromorphic datasets.
• Demonstrate that directly trained Spikformer surpasses prior SNN architectures on ImageNet and neuromorphic benchmarks while offering energy efficiency.

Proposed method

• Propose Spiking Self Attention (SSA) that uses spike-form Query, Key, and Value without softmax to compute attention.
• Compute Q, K, V via learnable projections followed by spike neuron layers, producing sequences with 0/1 values.
• Apply a scaling factor to stabilize dot-product computations and perform the attention via spike operations (logical AND and addition).
• Embed patches using a Spiking Patch Splitting (SPS) module and add spike-form relative position embeddings (RPE) generated by a spike-based block.
• Stack Spikformer encoder blocks (SSA + MLP with residuals) and use global average pooling and a linear head for classification.
• Use batch normalization in place of layer normalization, consistent with SNN training, and train from scratch with AdamW.

Experimental results

Research questions

• RQ1 Can self-attention be effectively implemented in SNNs without softmax and multiplications?
• RQ2 Does a spike-form self-attention (SSA) enable competitive or superior performance for image classification on both static and neuromorphic datasets compared to existing SNNs?
• RQ3 What are the accuracy and energy efficiency trade-offs of Spikformer versus state-of-the-art directly trained SNNs on ImageNet and neuromorphic datasets?

Key findings

• Spikformer with SSA outperforms state-of-the-art SNN-like frameworks on both static and neuromorphic datasets.
• On ImageNet, Spikformer-8-512 achieves 73.38% top-1 with 4 time steps, and Spikformer-8-768 reaches 74.81% top-1, surpassing SEW-ResNet-152.
• Spikformer demonstrates competitive or superior accuracy on CIFAR-10/100 and neuromorphic datasets (DVS128 Gesture and CIFAR10-DVS) with significantly smaller models and lower energy estimates.
• SSA avoids softmax and multiplications, operating via spike-form Q/K/V with sparse computations, resulting in lower theoretical energy consumption (SOPs) while maintaining or improving accuracy.
• A four-block SPS and spike-form embedding enable effective patch processing, with robust performance across varying time steps and network sizes.