Spike-driven Transformer

Summary

This paper proposes a Spike-driven Transformer that incorporates the spike-driven paradigm of SNNs into the Transformer and is hardware-friendly for neuromorphic chips. The authors use SDSA to transform the multiplication operation between Query, Key, and Value into a combination of mask and sparse addition operations and modify the self-attention structure to achieve linear complexity. As a result, this approach can achieve up to 87.2× lower energy consumption compared to the traditional self-attention method. Furthermore, membrane shortcut is introduced to ensure all spiking neurons communicate via binary spikes. The paper provides insights into implementing each module on neuromorphic chips and achieves competitive results on multiple datasets while reducing energy consumption compared to previous works.

Strengths

The authors introduce the Spike-Driven Self-Attention (SDSA) module as a replacement for the self-attention module in the current spiking Transformer. The proposed module exhibits superior performance while significantly reducing energy consumption, resulting in an architecture that is better suited for deployment on neuromorphic chips.

Weaknesses

This paper explores the feasibility of implementing Conv, MLP, and Self-attention models driven by spikes on neuromorphic chips. However, it does not delve into the implementation details of membrane shortcuts. This approach involves the direct transmission of membrane potential signals between spiking neurons and may contradict the spike-driven approach discussed in the original text. The effectiveness of a spikformer-based framework with only SDSA, without MS, remains unclear as most experiments in the literature combine both techniques. This weakens the argument for the performance of the spike-driven Transformer proposed in this paper. Although an ablation study was conducted on MS and SDSA in Table 5, it was limited to the CIFAR10/100 dataset and Spiking Transformer-2-512. The results suggest that the use of MS may be crucial to achieve state-of-the-art performance.

Questions

How is the membrane shortcut implemented on a neuromorphic chip? Is it necessary to use it instead of SEW shortcut? The article explains that membrane shortcuts are used to optimize the distribution of membrane potential, but for spiking neurons, wouldn't it be more appropriate to perform such an optimization in the time domain?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

No

Summary

This paper proposed a Spike-Driven Self-Attention (SDSA) module, which uses Hadamard product, column-wise summation, and spiking neuron layer to replace the matrix multiplication and softmax operation. Experiments on static and neuromorphic image classification demonstrate competitive performance and energy efficiency.

Strengths

1.. The authors proposed a novel form of linear self-attention module, SDSA, which leverages the features of the spikes and increases computational and memory efficiency. 2. The authors show extensive results from different models and different datasets. The performance of accuracy and energy efficiency is very strong. 3. SDSA reduces the computational complexity with a slight accuracy drop, and the authors show some attention maps to validate the effectiveness of the SDSA.

Weaknesses

1. There is still a noticeable drop between the spike-driven transformer and the original transformer. Can the authors show some results of the model with more time steps and report the limitations of the model? 2. Can the authors show the performance and energy consumption as the number of time steps decreases?

Questions

1. The ablation study shows the membrane shortcut brings significant accuracy improvement. I slightly doubt the importance of the SDSA. Can the authors try the combination of the membrane shortcut and some attention-free transformers like [1][2] to check if they can also achieve similar performance? [1] J Lee-Thorp et al., “FNet: Mixing Tokens with Fourier Transforms”, NAACL 2022 [2] W Yu et al, “MetaFormer Is Actually What You Need for Vision”, CVPR 2022

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The paper has few limitations.

Summary

This paper proposes an improved spike transformer by replacing Spike-Element-Wise shortcut in an existing spike transformer (ref [20]) with Membrane Shortcut from spike ResNet (ref [26]) .

Strengths

Match SOTA ImageNet top-1 precision achieved by ResNet [26] with slightly lower estimated energy consumption.

Weaknesses

Combining known techniques from two papers without modification is considered an incremental improvement. It is not clear why Transformer is a better choice than ConvNet in the context of Spike Neural Network. In ViT papers, the motivation for the Transformer is its model capacity can scale easily to handle large datasets and it has much better parallelism on TPU/GPU than ConvNet. None of these has been indicated as the goal for SNN. As shown in table-2, the ResNet SNN model in ref[27] is equally competitive as the proposed model. The notation is a bit confusing in a few equations. The s in Eq (5) is binary. This should be distinguished from floating-point numbers in Eq (4). But the same symbol R is used for both. The same observation holds for Eq (8)(10)(11).

Questions

The outputs of LIF spiking neurons at all timesteps are used as the inputs to the Transformer. This may be the tested method in [20], but why is this a better idea than using the output at the last timestep? As shown in table-2, the power is directly proportional to the timestep count. So this seems a critical parameter to optimize over. In Eq (12), S_L has the shape TxNxD. So outputs at all timesteps are used in classification head? In Eq (6), the input s has no 2d spatial dimensions. What are the 2d dimensions in Conv2d()?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

4 excellent

Limitations

None

Summary

The authors propose a variant of transformer networks with spiking neurons based on the LIF neuron. The submission reformulates the self-attention to use sparse addition and masking, and modifies the residual connections to transmit information in the domain of membrane potentials. They achieve a state of the art result on Imagenet. They present an energy analysis.

Strengths

It increases the amount of binary spike-based computations in the transformer. The experimental results are strong. A small ablation study is performed.

Weaknesses

The explanation for equations 15 and 16 can be improved The explanation of power estimates can be improved.

Questions

Can the authors provide an additional explanation for equations 15 and 16?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

not much discussed

why is the right hand side in eq (16) equal to the right hand side in eq (14) ?

Thank you for writing the explanation, however this still raises doubts: eq 16, rhs is: $SN(Q^i) \cdot SN ( SN(K^i)^\top \odot SN(V^i) ) $ eq 14 rhs is: $ SN ( SN(K^i)^\top \odot SN(Q^i) ) \cdot SN(V^i) $ the hadamard product can be exchanged: $Q_s \otimes K_s \otimes V_s = (Q_s \otimes K_s ) \otimes V_s = Q_s \otimes (K_s \otimes V_s) $, but not when there are thresholding non-linearities applied in between. As far as the reviewer understands, SN is a thresholding operation resulting in 1 or 0. So something needs to be proven to show that $SN ( SUM_c Q_s \otimes K_s ) \otimes SN( V_s) = SN(Q_s) \otimes SN ( SUM_c K_s \otimes V_s) $ This could be due to the fact that the summings are all non-negative, but even then the threshold should be at 1 or below. if the threshold is above 1, then this equality might not hold anymore. Can you please clarify ?

Thank you for addressing our questions

Spike Transformer: The authors have addressed our questions with more experimental results. We have raised our score to 7.