Spikformer: When Spiking Neural Network Meets Transformer

TL;DR

This work introduces Spiking Self Attention (SSA) and the Spiking Transformer (Spikformer) to merge the energy efficiency of spiking neural networks with Transformer-style self-attention. By operating on spike-form Q/K/V without softmax, SSA delivers sparse, multiplication-free attention suited to SNNs, enabling direct training on ImageNet with competitive accuracy. The Spikformer architecture combines Spiking Patch Splitting, SSA-based encoder blocks, and a classifier, achieving state-of-the-art results among directly trained SNNs on ImageNet and strong performance on neuromorphic datasets. Ablation studies validate SSA’s efficiency and effectiveness, while reducing energy consumption compared to traditional attention mechanisms. Overall, the paper demonstrates the feasibility and advantages of incorporating self-attention into SNNs, opening avenues for transformer-based spike models in vision tasks.

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.

Figures (6)

• Figure 1: Illustration of vanilla self-attention (VSA) and our Spiking Self Attention (SSA). A red spike indicates a value of 1 at that location. The blue dashed boxes provide examples of matrix dot product operation. For convenience, we choose one of the heads of SSA, where $N$ is the number of input patches and $d$ is the feature dimension of one head. $\rm{FLOPs}$ is the floating point operations and $\rm{SOPs}$ is the theoretical synaptic operations. The theoretical energy consumption to perform one calculation between Query, Key and Value in one time step is obtained from $8$-encoder-blocks 512-embedding-dimension Spikformer on ImageNet test set according to kundu2021hirehu2018residual. More details about the calculation of theoretical SOP and energy consumption are included in appendix. \ref{['sec:energy']}. (a) In VSA, $Q_{\mathcal{F}},K_{\mathcal{F}},V_{\mathcal{F}}$ are float-point forms. After the dot-product of $Q_{\mathcal{F}}$ and $K_{\mathcal{F}}$, the softmax function regularizes negative values in the attention map to positive values. (b) In SSA, all value in attention map is non-negative and the computation is sparse using spike-form $Q, K, V$ ($5.5\times 10^6$ VS. $77 \times 10^6$ in VSA). Therefore, the computation in SSA consumes less energy compared with VSA ($354.2\mu \rm{J}$). In addition, the SSA is decomposable (the calculation order of $Q,K$ and $V$ is changeable).
• Figure 2: The overview of Spiking Transformer (Spikformer), which consists of a spiking patch splitting module (SPS), a Spikformer encoder and a Linear classification head. We empircally find that the layer normalization (LN) does not apply to SNNs, so we use batch normalization (BN) instead.
• Figure 3: Attention map examples of SSA. The black region is 0.
• Figure 4: Fire rate of Query, Key and Value of blocks in Spikformer-$8$-$512$ on ImageNet test set.
• Figure 5: Training loss, testing loss and test accuracy on ImageNet.