PaperlayerView exact revision on arXiv

Source-linked AI summary

Spike-driven Transformer V2: Meta Spiking Neural Network Architecture Inspiring the Design of Next-generation Neuromorphic Chips

Man Yao, Jiakui Hu, Tianxiang Hu, Yifan Xu, Zhaokun Zhou, Yonghong Tian, Bo Xu, Guoqi Li

arXiv:2404.03663v1cs.NEcs.CV

TL;DR

Transformer-based SNNs had not clearly surpassed CNN-based SNNs while preserving their low-power spike-driven paradigm. Meta-SpikeFormer generalizes spike-driven Transformer design across architecture, attention, and shortcuts, achieving stronger reported performance and multi-task versatility. Its results also motivate Transformer-oriented neuromorphic chip design.

  • Problem

    Transformer-based SNNs had not clearly combined sparse-addition computation with advantages over CNN-based SNNs, limiting evidence for their broader value.

  • Method

    Meta-SpikeFormer extends spike-driven Transformer into a meta architecture with Conv-based and Transformer-based blocks, sparse-addition SDSA, and redesigned shortcuts.

  • Results

    80.0% top-1 accuracy on ImageNet-1K exceeds the current SOTA SNN baseline by 3.7% with 55M versus 66M parameters, while the backbone achieves SOTA results across tested vision tasks.

  • Takeaways & Limitations

    The results support Meta-SpikeFormer as a universal vision-backbone candidate and provide design inspiration for future Transformer-based neuromorphic chips.

  • Takeaways & Limitations

    Architectural descriptions simplify the temporal setting to T = 1, although static images are repeated across T timesteps in SNN inputs.

Abstract

from arXiv · show

Neuromorphic computing, which exploits Spiking Neural Networks (SNNs) on neuromorphic chips, is a promising energy-efficient alternative to traditional AI. CNN-based SNNs are the current mainstream of neuromorphic computing. By contrast, no neuromorphic chips are designed especially for Transformer-based SNNs, which have just emerged, and their performance is only on par with CNN-based SNNs, offering no distinct advantage. In this work, we propose a general Transformer-based SNN architecture, termed as ``Meta-SpikeFormer", whose goals are: 1) Lower-power, supports the spike-driven paradigm that there is only sparse addition in the network; 2) Versatility, handles various vision tasks; 3) High-performance, shows overwhelming performance advantages over CNN-based SNNs; 4) Meta-architecture, provides inspiration for future next-generation Transformer-based neuromorphic chip designs. Specifically, we extend the Spike-driven Transformer in \citet{yao2023spike} into a meta architecture, and explore the impact of structure, spike-driven self-attention, and skip connection on its performance. On ImageNet-1K, Meta-SpikeFormer achieves 80.0\% top-1 accuracy (55M), surpassing the current state-of-the-art (SOTA) SNN baselines (66M) by 3.7\%. This is the first direct training SNN backbone that can simultaneously supports classification, detection, and segmentation, obtaining SOTA results in SNNs. Finally, we discuss the inspiration of the meta SNN architecture for neuromorphic chip design. Source code and models are available at \url{https://github.com/BICLab/Spike-Driven-Transformer-V2}.

1 INTRODUCTION

Meta-SpikeFormer extends spike-driven Transformers into a general SNN architecture targeting sparse-addition computation, broad vision-task support, and stronger performance than CNN-based SNNs.

  • Meta-SpikeFormer investigates Transformer-based SNN design across network structure, skip connections, and fully addition-based Spike-Driven Self-Attention.
  • The architecture aligns Spike-driven Transformer with a general vision Transformer and expands four spiking encoding layers into four Conv-based SNN blocks.
  • The authors evaluate the architecture across classification, event-based action recognition, detection, and semantic segmentation datasets.
  • 80% ImageNet-1K accuracy exceeds the current SOTA SNN baseline by 3.7% with 17% fewer parameters, using 55M rather than 66M parameters.
  • Meta-SpikeFormer is the first direct-training SNN backbone reported to handle image classification, object detection, and semantic segmentation concurrently.
  • The work positions its meta architecture and component analysis as inspiration for future Transformer-based neuromorphic chip designs.

2 RELATED WORK

Related work spans direct-training and converted SNN backbones, Vision Transformer developments, and neuromorphic hardware supporting spike-driven computation.

  • SNN training commonly uses ANN-to-SNN conversion or direct training, because the spike function is non-differentiable for ordinary backpropagation.
  • Conv-based SNNs extend residual CNN designs to deeper networks, but current CNN-based SNNs have not demonstrated generality across vision tasks.
  • Vision Transformer research has explored architecture design, self-attention enhancement, training optimization, and efficient Transformer variants.
  • Neuromorphic chips are non-von Neumann systems commonly characterized by collocated processing and memory and spike-driven computing.
  • Existing neuromorphic designs primarily support spike-driven convolutional and MLP operators, while Transformer-based SNN hardware remains less established.

3 SPIKE-DRIVEN TRANSFORMER V2: META-SPIKEFORMER

Meta-SpikeFormer combines a general meta-Transformer block with Conv-based and Transformer-based SNN blocks, spiking neurons, sparse-addition SDSA, and membrane shortcuts.

  • 3.1 THE CONCEPT OF META TRANSFORMER ARCHITECTURE IN ANNS: A meta Transformer block alternates token mixing with a channel MLP, allowing different token mixers while preserving a general architecture.
  • 3.1 THE CONCEPT OF META TRANSFORMER ARCHITECTURE IN ANNS: The model embeds images as token sequences and feeds them through repeated meta Transformer blocks.
  • 3.2 SPIKING NEURON LAYER: Spiking neurons integrate spatial and temporal inputs into membrane potentials, emit binary spikes after threshold crossing, and reset or decay their state.
  • 3.3 META-SPIKEFORMER: For static datasets, images are repeated across T timesteps, while architectural descriptions assume T = 1 for simplicity.
  • 3.3 META-SPIKEFORMER: The macro architecture uses early Conv-based SNN stages followed by pyramid-shaped Transformer-based stages, replacing four encoding layers with Conv-based SNN blocks.
  • 3.3 META-SPIKEFORMER: Conv-based blocks use separable convolution for token mixing and replace a 1 × 1 channel MLP with a 3 × 3 ChannelConv.
  • 3.3 META-SPIKEFORMER: SDSA uses spiking Query, Key, and Value tensors without softmax or scale, with computational complexity linear in token number N.
  • 3.3 META-SPIKEFORMER: SDSA-1/2 use Hadamard products, SDSA-3/4 use matrix multiplication, and all four operators reduce computation to sparse addition.

4 EXPERIMENTS

Meta-SpikeFormer is evaluated across classification, event-based recognition, detection, and segmentation, combining strong accuracy with sparse-addition operation and architectural trade-offs in power, parameters, and accuracy.

  • Image Classification: 79.7% accuracy with 55M parameters surpasses MS-Res-SNN at 75.3% and Spike-driven Transformer at 76.3%.The comparison reports 55M versus 77M and 66M parameters, respectively.
  • Downstream Vision Tasks: Meta-SpikeFormer supports classification, object detection, and semantic segmentation as a directly trained SNN backbone, with SOTA results across tested datasets.COCO evaluation uses Mask R-CNN and direct training with a spike-driven Yolov5 variant; segmentation is evaluated on ADE20K and VOC2012.
  • Semantic Segmentation: 32.3% MIoU on ADE20K with 16.5M parameters requires 22.1mJ, compared with 32.9% and 147.1mJ for ResNet-18.The comparison also reports 35.7% MIoU and 152.7mJ for PVT-Tiny.
  • Ablation Studies: Removing SpeConv reduces power by 29.5% while losing 0.3% accuracy, whereas replacing channel convolution with MLP lowers accuracy by up to 2%.These ablations indicate different sensitivity to token-mixing and channel-processing choices in Conv-based SNN blocks.
  • Ablation Studies: RepConv improves accuracy and reduces parameters but increases energy, while SDSA-3 has the highest computational complexity and accuracy among tested SDSA operators.Shortcut choice has little effect on power, and MS provides the highest accuracy among the tested shortcuts.

5 DISCUSSION AND CONCLUSION

Meta-SpikeFormer frames Transformer-based SNNs as a co-design template for versatile vision processing and future neuromorphic chips. Its discussion emphasizes hybrid Conv–ViT modeling, spike-driven self-attention, and configurable meta-blocks.

  • Chip-design inspiration: Conv+ViT combines progressive local and global modeling, with convolution modeling features and Transformers capturing long-range dependencies.The paper reports this hybrid design benefits SNN performance and versatility.
  • Chip-design inspiration: SDSA provides long-distance dependency modeling for Transformer-based SNN blocks, addressing a capability current neuromorphic chips lack.
  • Chip-design inspiration: Meta Conv-based and Transformer-based blocks allow targeted optimization of internal designs for accuracy, parameters, and power.
  • Broader significance: Neuromorphic chip design is presented as algorithm-hardware co-design, with algorithmic features potentially inspiring algorithm-driven hardware design.
  • Conclusion: The paper reports 80% ImageNet-1K accuracy, 3.7% above the prior SOTA result with 17% fewer parameters, alongside SOTA results across tested classification, detection, and segmentation datasets.

A.2 SPIKE-DRIVEN SELF-ATTENTION (SDSA)

SDSA replaces vanilla self-attention’s dense operations with spike-based operators and develops four variants with linear token complexity. The variants trade accuracy, parameters, and energy through different tensor operations and threshold designs.

  • SDSA construction: Spike inputs are transformed into binary QS, KS, and VS using learnable 3 × 3 re-parameterization convolutions.The subsequent formulas assume T = 1 for convenience.
  • SDSA variants: SDSA-1 replaces matrix multiplication among QS, KS, and VS with Hadamard products, whose column sums and spiking neurons replace softmax and scale.The Hadamard product acts as a spike mask.
  • Complexity: SDSA-1 has O(ND) complexity, matching hydra attention, because it performs D vector dot products with N additions each.
  • SDSA variants: SDSA-2 obtains binary attention scores using QS alone, removing KS from the attention computation.Compared with SDSA-1, this reduces parameters and energy while losing 0.4 percentage points in accuracy.
  • SDSA variants: SDSA-3 is the default operator and is linear in N, while SDSA-4 makes the spiking threshold learnable to avoid explicit scale multiplication.SDSA-3 and SDSA-4 have nearly identical performance; SDSA-4 uses 0.1mJ less energy.

A.3 DISCUSSION ABOUT SDSA OPERATORS

All SDSA variants preserve spike-driven computation while lowering attention complexity relative to vanilla self-attention. Their implementation choices differ in whether spike interactions use masking or sparse-addition-compatible multiplication.

  • Energy implications: SDSA operators consume only sparse addition overall, and their energy is much lower than vanilla self-attention.
  • Operator comparison: Performance differences among SDSA-1/2/3/4 remain below 1.2%, supporting a meta-architecture view that leaves the token mixer unspecified.
  • Energy evaluation: Table 6 reports FLOPs for VSA and SDSA modules and converts them to energy using EMAC = 4.6pJ and EAC = 0.9pJ.RC and bR denote sums of spike firing rates for various spike matrices.

B THEORETICAL ENERGY EVALUATION

The paper estimates SNN energy by combining operation counts, AC energy, timesteps, and spike firing rates. It distinguishes theoretical qualitative comparisons from more refined layer-wise estimates and applies this framework to SDSA.

  • Energy model: Spike-driven Conv and MLP convert weight–spike matrix multiplication into sparse addition implemented as addressable addition on neuromorphic chips.
  • Evaluation scope: The paper notes that theoretical energy estimates facilitate qualitative comparisons and do not incorporate specific hardware implementation details.
  • Energy model: SNN energy is estimated as FLOPs × EAC × network spiking firing rate, whereas ANN energy uses FLOPs × EMAC.The paper uses EMAC = 4.6pJ and EAC = 0.9pJ for 45nm technology.
  • Layer-wise evaluation: A refined estimate counts each layer’s firing rate because network structure changes the additions triggered by a spike.
  • Layer-wise evaluation: For spike-driven Conv and MLP, power depends on EAC, timestep T, layer firing rate, and the corresponding ANN FLOPs.RC and RM represent layer spiking firing rates.
  • SDSA evaluation: SDSA energy includes the RepConv cost and a separate operator cost; sparse Q, K, and V make the SDSA function itself relatively inexpensive.SDSA-1 saves about 0.6mJ compared with SDSA-3 in the reported comparison.

C.1 IMAGENET-1K EXPERIMENTS

The ImageNet-1K experiments evaluate three Meta-SpikeFormer scales using the paper’s specified training hyperparameters and configurations.

  • Three Meta-SpikeFormer scales are evaluated on the ImageNet-1K classification benchmark.
  • The ImageNet-1K models are trained with the hyperparameters listed in Table 8.

C.2 COCO EXPERIMENTS

The COCO experiments evaluate Meta-SpikeFormer for object detection using both a Mask R-CNN fine-tuning setup and a direct-training spike-driven Yolov5 setup.

  • Meta-SpikeFormer is first used as a backbone for Mask R-CNN and fine-tuned for 24 epochs.The setup uses AdamW with an initial learning rate of 1 × 10^-4 and images resized and cropped to 1333 × 800.
  • Direct training combines Meta-SpikeFormer with a lightweight Yolov5 detector re-implemented in a spike-driven manner.
  • Better COCO performance than EMS-Res-SNN is obtained using a parameter count close to that model.

C.3 ADE20K EXPERIMENTS

The ADE20K experiments use Meta-SpikeFormer as a backbone with a spike-driven Semantic FPN and initialize the backbone from ImageNet-1K checkpoints.

  • Meta-SpikeFormer is employed as the backbone with a Semantic FPN re-implemented in a spike-driven manner.
  • Training runs for 160K iterations with batch size 20 and T = 1.ImageNet-1K-trained checkpoints initialize the backbone, while Xavier initializes newly added SNN layers.

C.4 ADDITIONAL RESULTS ON VOC2012 SEGMENTATION

Additional VOC2012 segmentation experiments test Meta-SpikeFormer on a 21-category benchmark using direct training at T = 1 and ImageNet-1K initialization.

  • VOC2012 contains 1460 training images, 1456 validation images, and 21 categories.
  • Meta-SpikeFormer is trained for 80k iterations at T = 1 on the VOC2012 segmentation task.ImageNet-1K-trained checkpoints initialize the backbone and Xavier initializes newly added SNN layers.
  • Table 10 reports semantic-segmentation performance on VOC2012.
  • Table 11 reports layer spiking firing rates for a Meta-SpikeFormer model with T = 4, 31.3M parameters, and SDSA-3 on ImageNet-1K.
Loading 2404.03663v1…