PaperlayerView exact revision on arXiv

Source-linked AI summary

Spike-driven Transformer

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

arXiv:2307.01694v1cs.NEcs.CV

TL;DR

SNNs offer energy-efficient spike-driven computation, but their limited task accuracy and existing hybrid Transformer designs leave a gap between efficiency and performance. The paper introduces a Spike-driven Transformer with SDSA and membrane-potential residual connections, reporting sparse-addition computation, 87.2× lower SDSA energy, and 77.1% ImageNet-1K top-1 accuracy. Its energy analysis assumes 32-bit floating-point implementation in 45nm technology.

  • Problem

    SNN applications are limited by poor task accuracy, while existing spiking Transformers retain MAC-heavy vanilla Transformer operations and are difficult to use on neuromorphic chips.

  • Method

    The paper introduces a fully spike-driven Transformer using SDSA with mask and sparse addition operations plus residual connections between membrane potentials.

  • Results

    77.1% top-1 accuracy on ImageNet-1K is reported as state of the art in the SNN field, while SDSA energy is 87.2× lower than vanilla self-attention.

  • Takeaways & Limitations

    The proposed architecture uses only sparse addition across the Transformer and is reported effective on static and neuromorphic datasets.

  • Takeaways & Limitations

    The energy-efficiency analysis assumes 32-bit floating-point implementation in 45nm technology and uses EMAC = 4.6pJ and EAC = 0.9pJ.

Abstract

from arXiv · show

Spiking Neural Networks (SNNs) provide an energy-efficient deep learning option due to their unique spike-based event-driven (i.e., spike-driven) paradigm. In this paper, we incorporate the spike-driven paradigm into Transformer by the proposed Spike-driven Transformer with four unique properties: 1) Event-driven, no calculation is triggered when the input of Transformer is zero; 2) Binary spike communication, all matrix multiplications associated with the spike matrix can be transformed into sparse additions; 3) Self-attention with linear complexity at both token and channel dimensions; 4) The operations between spike-form Query, Key, and Value are mask and addition. Together, there are only sparse addition operations in the Spike-driven Transformer. To this end, we design a novel Spike-Driven Self-Attention (SDSA), which exploits only mask and addition operations without any multiplication, and thus having up to $87.2\times$ lower computation energy than vanilla self-attention. Especially in SDSA, the matrix multiplication between Query, Key, and Value is designed as the mask operation. In addition, we rearrange all residual connections in the vanilla Transformer before the activation functions to ensure that all neurons transmit binary spike signals. It is shown that the Spike-driven Transformer can achieve 77.1\% top-1 accuracy on ImageNet-1K, which is the state-of-the-art result in the SNN field. The source code is available at https://github.com/BICLab/Spike-Driven-Transformer.

1 Introduction

The paper addresses limited SNN task accuracy and the energy costs of hybrid spiking Transformers by introducing a fully spike-driven Transformer. Its SDSA module and membrane-potential residual connections enable sparse-addition computation while achieving strong reported accuracy.

  • Motivation: SNNs offer event-driven, binary-spike computation, but their applications are limited by poor task accuracy.Event-driven computation is triggered sparsely, and binary communication permits synaptic accumulate operations.
  • Motivation: Existing spiking Transformers retain MAC-heavy operations such as dot products, softmax, and scaling, resulting in hybrid computing.These designs convert Q, K, and V to spikes but still perform matrix multiplications and can produce multi-bit outputs requiring normalization.
  • Method: SDSA replaces Query-Key matrix multiplication with a Hadamard mask and uses column summation plus a spiking neuron layer instead of softmax and scaling.The resulting attention vector is binary and masks channels in the Value tensor; SDSA is also characterized as linear in token and channel dimensions.
  • Method: 87.2× lower energy cost is reported for SDSA, including Query, Key, and Value generation, than for vanilla self-attention.The operator uses mask and sparse addition operations rather than multiplication-heavy attention computation.
  • Method: The redesigned residual connections are placed before spike neuron layers so spiking neurons communicate through binary spikes.Residual connections between membrane potentials allow spike-weight multiplication to become addition instead of producing multi-bit integer spikes.
  • Results: 77.1% top-1 accuracy on ImageNet-1K is reported as the state-of-the-art result in the SNN field.Experiments are also reported to outperform or match state-of-the-art SNNs on static and neuromorphic datasets.

2 Related Works

Related work spans bio-inspired SNNs, neuromorphic hardware, and efficient Transformers. The paper connects spike-driven computation with linear-attention ideas for reducing Transformer complexity.

  • Spiking Neural Networks: SNN research combines biological mechanisms with deep-learning technologies, including neuron models, learning rules, attention, and network architectures.These efforts seek to retain spike-driven benefits while improving task accuracy.
  • Neuromorphic Chips: Neuromorphic chips use brain-inspired non-von Neumann architectures with parallel operation, collocated processing and memory, scalability, and spike-driven computing.Spiking neurons and synapses serve as basic computational units.
  • Efficient Transformers: Efficient Transformer research targets quadratic token scaling through convolution-attention mixing, architectural optimization, and linear attention.Linear attention commonly removes softmax and rearranges Query, Key, and Value computation; spiking Transformers inherently omit softmax.
  • Positioning: The Spike-driven Transformer combines an SDSA module with redesigned shortcuts across position embedding, self-attention, and MLP components.The shortcuts are constructed before spike neuron layers to keep spike matrices binary and enable addition operations.

3 Spike-driven Transformer

The Spike-driven Transformer redesigns Transformer components around binary spikes, sparse additions, and membrane-potential residual connections. Its SDSA uses mask-like operations and linear complexity while retaining task performance.

  • 3 Spike-driven Transformer: LIF neurons produce binary spike tensors by firing when membrane potential exceeds a threshold and otherwise emitting zero.The membrane potential couples spatial input with temporal state and decays when the neuron does not fire.
  • 3 Spike-driven Transformer: The architecture comprises Spiking Patch Splitting, SDSA, MLP, and a linear classification head, with membrane-potential residual connections in SDSA and MLP blocks.SPS converts image sequences into flattened spike patches, and global average pooling feeds the classification head.
  • 3.2 Membrane Shortcut in Spike-driven Transformer: The membrane shortcut preserves binary spike signals, whereas spike-addition shortcuts can produce multi-bit integer spikes and introduce integer multiplication.The paper also reports higher task accuracy for MS-Res-SNN than SEW-Res-SNN in Transformer-based SNNs.
  • 3.3 Spike-driven Self-Attention: SDSA replaces conventional Query-Key and Value matrix multiplications with Hadamard products, column-wise summation, spiking, and mask operations.Hadamard products between spike tensors act as element-wise masks, while sparse column sums have negligible energy cost.
  • 3.3 Spike-driven Self-Attention: SDSA is linear in token number N because Key and Value can be combined first, avoiding vanilla self-attention’s softmax operation.The spiking neuron layer serves as the kernel function in this linear-attention interpretation.
  • 3.3 Spike-driven Self-Attention: SDSA has complexity O(0 + ND) and typically requires about 0.02ND additions because Key and Value vectors are highly sparse.Binary attention scores mask unimportant channels in the sparse Value tensor, with a slight reported accuracy loss.

4 Theoretical Energy Consumption Analysis

The energy analysis explains how spike-driven computation reduces operations in convolution, MLP, and self-attention. It estimates SNN energy using spike firing rate and simulation timestep relative to ANN computation.

  • Spike-driven in Conv and MLP: Spike-driven Conv and MLP transform matrix multiplication into sparse, addressable additions by combining event-driven inputs with binary spike communication.No computation is triggered when the input is zero.
  • Spike-driven in Self-attention: In spike-driven self-attention, element-wise multiplication between binary spike tensors becomes a mask operation rather than direct matrix multiplication.The mask operation is described as requiring no energy because spike elements are restricted to 0 or 1.
  • Energy Consumption Comparison: SNN energy under a fixed architecture can be estimated from ANN FLOPs together with spike firing rate R and simulation timestep T.The comparison evaluates Conv, self-attention, and MLP modules at the same scale in vanilla and Spike-driven Transformers.

5 Experiments

Experiments evaluate the Spike-driven Transformer on static and neuromorphic datasets, ImageNet accuracy and energy, spike firing rates, and component ablations.

  • Datasets: The experiments cover ImageNet, CIFAR-10/100, CIFAR10-DVS, and DVS128 Gesture, including both static and neuromorphic datasets.Gesture and CIFAR10-DVS require converting event streams into frame sequences before processing.
  • ImageNet accuracy: 77.1% top-1 accuracy is achieved on ImageNet with Spike-driven Transformer-8-768 at 288×288 inference resolution.At 224×224, the same model obtains 76.32% with 66.34M parameters.
  • ImageNet accuracy: 2.0% higher accuracy is obtained by Spike-driven Transformer-8-384 than SpikFormer-8-384 under the same parameter setting.
  • ImageNet accuracy: 74.11% accuracy is reported for Spike-driven Transformer-6-512, compared with 67.04% for SEW-Res-SNN-34 and 69.15% for MS-Res-SNN-34.The corresponding parameter counts are 23.27M, 21.79M, and 21.80M.
  • Energy consumption: 4.50mJ is the reported power for Spike-driven Transformer-8-512, versus 4.89mJ for SEW-Res-SNN-50 and 5.11mJ for MS-Res-SNN-34.The corresponding accuracies are 74.57%, 67.78%, and 69.42%.
  • Energy consumption: 9.3× higher energy is required by ANN-ViT than the spike-driven counterpart under the same 8-512 structure.The reported powers are 41.77mJ for ANN-ViT and 4.50mJ for the spike-driven model; with T = 1, the advantage reaches 36.7×.
  • Ablation study: SDSA causes a slight performance loss, while membrane shortcuts produce significant accuracy improvements in the ablation study.The study also reports that SDSA masks background information and reduces the spike firing rate of the processed Value tensor.

6 Conclusion

The paper proposes a Spike-driven Transformer built around SDSA and rearranged residual connections, using sparse additions and linear-complexity self-attention across token and channel dimensions.

  • The Spike-driven Transformer uses only sparse addition operations by combining SDSA with rearranged residual connections.Vanilla self-attention’s matrix multiplication, softmax, and scale are replaced with mask, addition, and a spike neuron layer.
  • SDSA replaces matrix multiplication, softmax, and scale with mask, addition, and a spike neuron layer.
  • SDSA has linear complexity in both token and channel dimensions.
  • Experiments on static-image and neuromorphic datasets verify the method’s effectiveness and efficiency.

S1 Energy Consumption Analysis Details

The energy analysis estimates Spike-driven Transformer consumption from vanilla Transformer FLOPs, timesteps, and spike firing rates under a stated hardware-energy model.

  • Spike-driven Transformer energy is estimated using vanilla Transformer FLOPs together with timesteps T and spike firing rates R.The analysis treats T and R as known quantities after evaluating the vanilla Transformer’s FLOPs.
  • The MLP FLOPs equation uses input dimension im and output dimension om.
  • The spike firing rate is the proportion of non-zero elements in a spike tensor.R averages the firing rates of QS, KS, and VS, while bR sums those of QS and KS.
  • Under the 45nm 32-bit floating-point model, EMAC = 4.6pJ and EAC = 0.9pJ.
  • SNNs are theoretically more energy efficient when EAC × T × R < EMAC.With EAC × T usually constant, smaller R yields lower theoretical energy cost.

S2 Experiment Details

The experiments cover static-image and neuromorphic classification, using repeated image inputs for static data and temporal event streams for neuromorphic data.

  • The evaluation uses ImageNet-1K, CIFAR-10/100, CIFAR10-DVS, and DVS128 Gesture.The first two categories are static image datasets, while the latter two are neuromorphic classification datasets.
  • ImageNet-1K contains 1.28 million training images, 50k test images, and 1,000 categories.
  • CIFAR10-DVS contains 10,000 event-based samples, while DVS128 Gesture records event-based gesture data at 128 × 128 spatial resolution.
  • Static images are repeatedly input at each timestep, and more timesteps improve accuracy while increasing training, hardware, and inference-energy requirements.
  • Neuromorphic datasets exploit spatio-temporal dynamics through event streams encoding brightness-change time, location, and polarity.
  • The small-dataset experiments use T = 4 for static CIFAR-10/100 and T = 16 for neuromorphic CIFAR10-DVS and Gesture.Training lasts 200 epochs, with dataset-specific batch sizes.

S3 Attention Map

SDSA computes a binary attention map from spike-form Query and Key tensors, then masks Value channels; the resulting firing-rate maps visualize this modulation.

  • Given spike input features, learnable linear matrices produce Q, K, and V, which a spike neuron layer converts into QS, KS, and VS.
  • SDSA computes ˆVS by applying a Hadamard-product mask derived from QS and KS to VS.
  • The attention map is a binary vector, so SDSA performs hard attention by masking unimportant Value channels.
  • The attention-map visualization averages VS and ˆVS across T timesteps and H heads and plots their spike firing rates.
  • Figure S1 maps each patch’s firing rate on a 14 × 14 grid, with redder colors indicating higher rates and bluer colors lower rates.The comparison is between the original VS attention map and the SDSA-regulated ˆVS map.
Loading 2307.01694v1…