Scaling Spike-driven Transformer with Efficient Spike Firing Approximation Training

Abstract

The ambition of brain-inspired Spiking Neural Networks (SNNs) is to become a low-power alternative to traditional Artificial Neural Networks (ANNs). This work addresses two major challenges in realizing this vision: the performance gap between SNNs and ANNs, and the high training costs of SNNs. We identify intrinsic flaws in spiking neurons caused by binary firing mechanisms and propose a Spike Firing Approximation (SFA) method using integer training and spike-driven inference. This optimizes the spike firing pattern of spiking neurons, enhancing efficient training, reducing power consumption, improving performance, enabling easier scaling, and better utilizing neuromorphic chips. We also develop an efficient spike-driven Transformer architecture and a spike-masked autoencoder to prevent performance degradation during SNN scaling. On ImageNet-1k, we achieve state-of-the-art top-1 accuracy of 78.5\%, 79.8\%, 84.0\%, and 86.2\% with models containing 10M, 19M, 83M, and 173M parameters, respectively. For instance, the 10M model outperforms the best existing SNN by 7.2\% on ImageNet, with training time acceleration and inference energy efficiency improved by 4.5$\times$ and 3.9$\times$, respectively. We validate the effectiveness and efficiency of the proposed method across various tasks, including object detection, semantic segmentation, and neuromorphic vision tasks. This work enables SNNs to match ANN performance while maintaining the low-power advantage, marking a significant step towards SNNs as a general visual backbone. Code is available at https://github.com/BICLab/Spike-Driven-Transformer-V3.

References in corpus (38)

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Is Space-Time Attention All You Need for Video Understanding?
• Spatio-Temporal Backpropagation for Training High-performance Spiking Neural Networks
• BEiT: BERT Pre-Training of Image Transformers
• MetaFormer Baselines for Vision
• Focal Self-attention for Local-Global Interactions in Vision Transformers
• Deep Residual Learning in Spiking Neural Networks
• Speck: A Smart event-based Vision Sensor with a low latency 327K Neuron Convolutional Neuronal Network Processing Pipeline
• Spikformer: When Spiking Neural Network Meets Transformer
• Temporal Efficient Training of Spiking Neural Network via Gradient Re-weighting
• Optimal ANN-SNN Conversion for High-accuracy and Ultra-low-latency Spiking Neural Networks
• MobileViTv3: Mobile-Friendly Vision Transformer with Simple and Effective Fusion of Local, Global and Input Features
• Designing BERT for Convolutional Networks: Sparse and Hierarchical Masked Modeling
• EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
• Spike-driven Transformer
• Towards Lossless ANN-SNN Conversion under Ultra-Low Latency with Dual-Phase Optimization
• Online Training Through Time for Spiking Neural Networks
• Rethinking Mobile Block for Efficient Attention-based Models
• How Mask Matters: Towards Theoretical Understandings of Masked Autoencoders
• Spike-driven Transformer V2: Meta Spiking Neural Network Architecture Inspiring the Design of Next-generation Neuromorphic Chips
• Deep Directly-Trained Spiking Neural Networks for Object Detection
• Spike Calibration: Fast and Accurate Conversion of Spiking Neural Network for Object Detection and Segmentation
• Parallel Spiking Neurons with High Efficiency and Ability to Learn Long-term Dependencies
• Event-Driven Learning for Spiking Neural Networks
• SEENN: Towards Temporal Spiking Early-Exit Neural Networks
• Accurate and efficient time-domain classification with adaptive spiking recurrent neural networks
• Spikformer V2: Join the High Accuracy Club on ImageNet with an SNN Ticket
• A Long Short-Term Memory for AI Applications in Spike-based Neuromorphic Hardware
• Masked Spiking Transformer
• Towards Memory- and Time-Efficient Backpropagation for Training Spiking Neural Networks
• Tensor Decomposition Based Attention Module for Spiking Neural Networks
• Memory-Efficient Reversible Spiking Neural Networks
• Inherent Redundancy in Spiking Neural Networks
• Gated Attention Coding for Training High-performance and Efficient Spiking Neural Networks
• High-Performance Temporal Reversible Spiking Neural Networks with $O(L)$ Training Memory and $O(1)$ Inference Cost
• Is Conventional SNN Really Efficient? A Perspective from Network Quantization
• MetaLA: Unified Optimal Linear Approximation to Softmax Attention Map