QKFormer: Hierarchical Spiking Transformer using Q-K Attention

Summary

The authors introduce a novel spiking transformer, QKFormer, which incorporates several innovative features: a spike-form Q-K attention mechanism with linear complexity and enhanced energy efficiency, a hierarchical structure that facilitates multi-scale spiking representation, and a spiking patch embedding with a deformed shortcut designed to optimize spiking information transmission and integration. Notably, this model achieves a remarkable top-1 accuracy of 85.65% on the ImageNet-1k dataset, marking the first instance where directly trained spiking neural networks have surpassed the 85% accuracy threshold on this benchmark.

Strengths

1) The authors' introduction of a novel spike-form Q-K attention module with linear complexity represents a significant advancement in mitigating the quadratic computational complexity of spiking self-attention, a core component of traditional spiking transformers. This innovation demonstrates a substantial degree of originality and contributes to the optimization of spiking neural network architectures. 2) The authors' design of a powerful spiking patch embedding with deformed shortcut (SPEDS) is a notable achievement, as it effectively enhances spiking information transmission and integration. This enhancement contributes to the overall performance of the proposed QKFormer model and underscores the authors' creative approach to improving SNN performance. 3) The proposed QKFormer model, leveraging the innovative Q-K attention and SPEDS, represents a substantial leap forward in SNN performance. Its ability to outperform state-of-the-art SNNs on various static and neuromorphic datasets, particularly the ImageNet-1K dataset, where it achieves over 85% accuracy for the first time with directly trained SNNs, underscores its significance and potential impact on real-world applications of SNNs. This achievement not only demonstrates the model's quality but also highlights its potential to revolutionize the field of SNN research.

Weaknesses

1) The authors' assertion that SPEDS can facilitate spiking information transmission and integration lacks sufficient elaboration on the underlying mechanisms. 2) Clarifying the rationale behind the assertion that the low variance of Q-K attention negates the necessity for scaling operations would significantly enhance the paper's rigor and credibility. 3) To improve the overall writing quality and grammar of the paper, several revisions are recommended. Firstly, in the caption of Figure 1, the phrase "The input is" should be revised to "The inputs are" to ensure consistency with the plural form of "inputs." Secondly, the expression in line 246 appears unusual and requires refinement to enhance clarity. Lastly, the abbreviation "LF" should be corrected to "IF" in line 316, assuming it is a typographical error and should stand for "if." Additionally, a comprehensive review and refinement of the entire paper's language, grammar, and structure would further elevate its academic style and readability.

Questions

1) Could the summation operation in the QK attention mechanism potentially result in an excessively high firing rate within the attention vector? 2) Could you elaborate on the relationship between Figure 3a and the original input image? 3) As noted in the weaknesses section.

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors briefly mention limitations in the Appendix, but discussing them in the main text would enhance clarity. Another limitation is that the model was only evaluated on image classification, limiting its generalizability.

Summary

This model introduces a spike-form Q-K attention mechanism that efficiently models the importance of token or channel dimensions using binary values, significantly reducing computational complexity .The model is evaluated on the ImageNet-1K dataset, achieving an impressive top-1 accuracy of 85.65%, surpassing existing benchmarks for spiking transformers.

Strengths

1. The proposed spike-form Q-K attention mechanism provides an efficient way to handle token or channel dimensions with binary spikes. 2. The model achieves over 85% accuracy on ImageNet-1K, which is a very impressive result.

Weaknesses

1、The Q-K attention module discussed in the paper has already been introduced in Spike-Driven Transformer V2[1] (SDSA-2 in Figure 3). Additionally, the overall architecture of the QKFormer model closely resembles that in [2]. 2、 In Section 4.1, the authors incorrectly compare the full-precision results of their work with Binary Neural Networks. Upon reviewing the original texts, it was found that PokeBNN (20.7MB) [2] and MeliusNet59 (17.4MB) [3] report model sizes, not parameter counts (20.7M, 17.4M). Comparisons should use consistent metrics. 3、The baseline for the Q-K attention ablation study (QKCA + SSA) is not reasonable. There is no information on the model with only SSA or without the attention module. 4、The effectiveness of the SPEDS module with only Q-K attention and without SPEDS remains unclear. It appears that the architecture of Spikformer differs from QKFormer, with QKFormer’s architecture being more similar to that in Spike-Driven Transformer V2 [1] Man Yao, JiaKui Hu, Tianxiang Hu, Yifan Xu, Zhaokun Zhou, Yonghong Tian, XU Bo, and Guoqi Li. Spike-driven transformer v2: Meta spiking neural network architecture inspiring the design of next-generation neuromorphic chips. In The Twelfth International Conference on Learning Representations, 2023. [2] Zhang Y, Zhang Z, Lew L. Pokebnn: A binary pursuit of lightweight accuracy[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022: 12475-12485. [3] Bethge J, Bartz C, Yang H, et al. Meliusnet: An improved network architecture for binary neural networks[C]//Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. 2021: 1439-1448.

Questions

1. The paper suggests that Q-K attention proposed here is simpler compared to SDSA-3 in Spike-Driven Transformer V2, and ablation experiments show performance degradation with Q-K Attention. The QKFormer architecture closely resembles [1] yet achieves a 5.65% improvement on ImageNet. Could this improvement be attributed to superior data augmentation or training strategies rather than the effectiveness of the attention methods and model architecture? 2. Since the accuracy on ImageNet-1K significantly exceeds those of existing comparable models, it is challenging to evaluate the results without access to the corresponding code. Could you please make the code open source?

Rating

5

Confidence

5

Soundness

3

Presentation

3

Contribution

2

Limitations

N/A

Summary

The author has thoughtfully considered the attention structure of the existing Spiking Transformer (SSA) as well as issues present in other parts. Spiking neural networks are characterized by their high efficiency and energy-saving features. For this purpose, the author has proposed a more efficient attention mechanism. To better enhance performance, the author has also made modifications to the embedding part. From the experimental results, the new model achieved better performance on several typical datasets with lower energy consumption and fewer parameters. Additionally, the ablation experiments designed by the author demonstrated the rationality of the attention mechanism and embedding modifications. Overall, the method proposed in the article is highly feasible and logically coherent.

Strengths

1. The article analyzes the issue of computational redundancy in existing Spiking Transformer attention mechanisms. 2. The Q-K Attention offers an almost perfect solution, being simple in design (not a bad thing), yet reducing computational complexity while enhancing performance. 3. A better embedding method has been designed. 4. Ablation experiments have demonstrated the rationality of these approaches.

Weaknesses

1. Lack of an overview of pipeline. Some charts, such as Figure 3, seem to be of little significance. In contrast, SPEDS might require a more intuitive representation, given that you have made significant modifications to the previous embedding. 2. Could the ablation study include additional experiments? We have seen experiments on SPEDS and Q-K Attention using CIFAR-100. Could we conduct a similar set of experiments on neuromorphic datasets? Given the unique nature of events, we hope to see SPEDS achieve performance improvements on neuromorphic datasets as well.

Questions

I ampleased to see that such a simplified attention mechanism can improve performance while saving energy. But I still have the following questions: 1. The capability of SPEDS on event datasets needs to be demonstrated. 2. As shown Eq.6, row and column-wise summation is performed when computing Q-K attention. Attention mechanisms are designed to integrate global attention. Would Q-K attention cause the loss of features that do not align with the row and column directions? 3. After summation in Eq.6,, Q undergoes neuron processing, is this to ensure the model is spike-driven? Representations in SNNs are inherently sparse; would this lead to feature loss?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

See weaknesses and Questions

Summary

The authors proposed QKFormer, which pushes SNN to 85% accuracy in ImageNet, becoming a new sota and contributing to the SNN community.

Strengths

1. The authors propose a model accuracy of Sota, which is 10% higher than the previous spikformer, and the number of parameters is lower than spikformer, which is undoubtedly the leader of the SNN-Transformer series. 2. The description is clear and the comparison is complete. The working mechanism of Q-K Attention is clearly expressed. The visualization results are rich. Comparison experiments are complete, comparing the time complexity and space complexity of different attentions in SNN domain in recent years, comparing with the Transformer architecture in ANN domain, covering the commonly used static datasets CIFAR10, CIFAR100, ImageNet, and DVS datasets DVS128, CIFAR10-DVS.

Weaknesses

1. The current accuracy is high, but it still requires a large time step and a large amount of computation.

Questions

1. In Table2 the authors list a comparison of multiple metrics for different models, but there is a large time step for SNN, how does QKFormer's time during training or inference compare to other models? 2. The authors could add an explanation of Table1 to the supplemental material to visually analyze the complexity of other attention mechanisms. 3. Q-K token attention Sum each column and then use LIF on the column vectors, as shown in Fig. 1(a), 4 to 1,2 to 0. How is the threshold set?

Rating

8

Confidence

5

Soundness

4

Presentation

3

Contribution

4

Limitations

See Questions and Weaknesses.