SGM-Transformer: Rethinking Gradient Information Loss and Compensation in Spiking Neural Networks

摘要

Spiking Neural Network (SNN), as a next-generation neural network technology, use binary spike signals as carriers of information. They offer advantages such as low energy consumption, low computational complexity, and high information transmission rates. However, deep SNNs suffer from the gradient vanishing problem due to issues like the non-differentiability of the step function and neuron dormancy. To address this gradient problem, we first propose the MCLIF neuron, which optimizes the backpropagation mechanism and compensates for gradient information from both temporal and spatial dimensions. Furthermore, we design a spiking attention mechanism tailored to the temporal characteristics of SNNs. By introducing the QK Memory to embed temporal features, we make full use of information from different time steps. Additionally, we propose a gradient correction module to enhance the model's representational power from both temporal and spatial dimensions. The proposed SGM-Transformer achieves state-of-the-art (SOTA) performance in image classification tasks such as CIFAR10, CIFAR100, and CIFAR10-DVS, and also excels in industrial defect classification scenarios. The code will be made available after the paper is accepted.