Combining Aggregated Attention and Transformer Architecture for Accurate and Efficient Performance of Spiking Neural Networks

TL;DR

SAFormer addresses the challenge of combining low-power Spiking Neural Networks with Transformer performance by introducing Spike Aggregated Self-Attention (SASA), which computes attention using only $Q$ and $K$ derived from inputs and down-sampled aggregation $AG(\cdot)$, omitting the $V$ matrix and achieving $O(nD)$ time and space. A Depthwise Convolution (DWC) module enhances feature diversity in SASA. The architecture combines Spike Patch Splitting, LIF neurons, and encoder blocks with residual paths to deliver energy-efficient high accuracy. Evaluations on static datasets (e.g., CIFAR-10/100) and neuromorphic datasets (CIFAR-10-DVS, DVS128-Gesture) show SAFormer attains competitive accuracy with substantial energy savings, outperforming several baselines. This work highlights SAFormer as a viable route to edge-friendly neuromorphic learning with Transformer-inspired capacity.

Abstract

Spiking Neural Networks have attracted significant attention in recent years due to their distinctive low-power characteristics. Meanwhile, Transformer models, known for their powerful self-attention mechanisms and parallel processing capabilities, have demonstrated exceptional performance across various domains, including natural language processing and computer vision. Despite the significant advantages of both SNNs and Transformers, directly combining the low-power benefits of SNNs with the high performance of Transformers remains challenging. Specifically, while the sparse computing mode of SNNs contributes to reduced energy consumption, traditional attention mechanisms depend on dense matrix computations and complex softmax operations. This reliance poses significant challenges for effective execution in low-power scenarios. Given the tremendous success of Transformers in deep learning, it is a necessary step to explore the integration of SNNs and Transformers to harness the strengths of both. In this paper, we propose a novel model architecture, Spike Aggregation Transformer (SAFormer), that integrates the low-power characteristics of SNNs with the high-performance advantages of Transformer models. The core contribution of SAFormer lies in the design of the Spike Aggregated Self-Attention (SASA) mechanism, which significantly simplifies the computation process by calculating attention weights using only the spike matrices query and key, thereby effectively reducing energy consumption. Additionally, we introduce a Depthwise Convolution Module (DWC) to enhance the feature extraction capabilities, further improving overall accuracy. We evaluated and demonstrated that SAFormer outperforms state-of-the-art SNNs in both accuracy and energy consumption, highlighting its significant advantages in low-power and high-performance computing.

Figures (6)

• Figure 1: The overall architecture of our proposed method integrates a Spike Patch Splitting (SPS) module, multiple SAFormer encoder blocks, and a classification layer. SPS divides the image into multiple patches and then extracts multi-scale features from each patch. Next, stacking multiple Encoder Blocks facilitates efficient feature encoding, ultimately generating the prediction output through the classification layer. We have reconfigured the downsampling module and residual connection strategy to optimize the network structure, ensuring the flow of information in binary sparse matrix form throughout the neural network. The Encoder Blocks serve as core building units, combining the SASA mechanism with a multi-layer perceptron (MLP) module.
• Figure 2: Performance comparison of different methods on the CIFAR-10 and CIFAR-100 datasets. Both error rates and energy efficiency of our model relative to each method are shown to highlight the significance of our model. The averages and standard deviations are calculated over three independent runs.
• Figure 3: The performance comparison on the CIFAR10-DVS and DVS128-Gesture datasets underscores the advantages of our model in terms of both accuracy rates and energy efficiency relative to other methods. The averages and standard deviations are calculated over three independent runs.
• Figure 4: We randomly selected six images (left column) from the CIFAR-10 dataset and visualized the attention maps generated by SDSA (middle column) and our SASA (right colume). The color intensity of the attention maps indicates the degree of attention assigned to each region, with darker colors representing higher attention.
• Figure 5: The performance of five different deep learning models over time on the standard image classification benchmark datasets CIFAR-10 and CIFAR-100.