Spiking Transformer with Spatial-Temporal Spiking Self-Attention

Abstract

Spiking Neural Networks are celebrated for energy efficiency and biological plausibility. Building on Spiking Self-Attention (SSA), Spiking Transformers are extensively studied due to their exceptional performance. However, SSA focuses solely on spatial dimension at each time step, overlooking the crucial features across temporal dimension. To address this, we propose the Spatial-Temporal Spiking Self-Attention (STSSA), a spike-driven mechanism that leverages both spatial and temporal information with negligible additional computational overhead. Specifically, we extract the Representative Spiking Temporal Tokens (RSTT) and apply temporal window masking to the RSTT. These tokens are inserted between the Query and Key to integrate temporal features. Furthermore, we design a Multi-dimensional Learnable Scaling Factor (MLSF) to adapt to STSSA. Our results consistently demonstrate that STSSA outperforms SSA across extensive experiments on sequential, neuromorphic, and static datasets. Notably, STSSA achieves performance improvements of 5.7% and 2.9% over SSA on Sequential CIFAR-100 and CIFAR-10DVS, respectively. STSSA provides a powerful alternative within the family of Spiking Self-Attention mechanisms.