QKFormer: Hierarchical Spiking Transformer using Q-K Attention

Authors: chenlin  zhou
han  zhang
zhaokun  zhou
liutao  yu
liwei  huang
xiaopeng  fan
li  yuan
zhengyu  ma
huihui  zhou
yonghong  tian
pengcheng  laboratory
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 04:33:09

Abstract: Spiking Transformers, which integrate Spiking Neural Networks (SNNs) with Transformer architectures, have attracted significant attention due to their potential for low energy consumption and high performance. However, there remains a substantial gap in performance between SNNs and Artificial Neural Networks (ANNs). To narrow this gap, we have developed QKFormer, a direct training spiking transformer with the following features: i) _Linear complexity and high energy efficiency_, the novel spike-form Q-K attention module efficiently models the token or channel attention through binary vectors and enables the construction of larger models. ii) _Multi-scale spiking representation_, achieved by a hierarchical structure with the different numbers of tokens across blocks. iii) _Spiking Patch Embedding with Deformed Shortcut (SPEDS)_, enhances spiking information transmission and integration, thus improving overall performance. It is shown that QKFormer achieves significantly superior performance over existing state-of-the-art SNN models on various mainstream datasets. Notably, with comparable size to Spikformer (66.34 M, 74.81\%), QKFormer (64.96 M) achieves a groundbreaking top-1 accuracy of **85.65\%** on ImageNet-1k, substantially outperforming Spikformer by **10.84\%**. To our best knowledge, this is the first time that directly training SNNs have exceeded 85\% accuracy on ImageNet-1K.

Main paper: 
 introduction: 
 Regarded as the third generation of neural networks [1] , the brain-inspired Spiking Neural Networks (SNNs) are potential competitors to Artificial Neural Networks (ANNs) due to their high biological plausibility and high energy efficiency attributed to their event-driven properties [2] . Transformer, originally designed for natural language processing [3] , has flourished in various computer vision tasks, including image classification [4, 5] , object detection [6, 7, 8] and semantic segmentation [9, 10] . Spiking Transformers (Transformer-based SNNs) [11, 12, 13, 14, 15] , which integrate spiking neural networks with transformer architecture, have attracted significant attention. This innovative combination provides great potential to develop advanced AI algorithms with high performance and low energy consumption.As the architecture of the transformers is essential to the model's performance [4, 5, 16, 8, 10] , designing new architectures for transformer-based SNNs is quite challenging in terms of space requirements for the following reasons [11, 13, 15] . i). Spiking Self Attention (SSA) [11] , the core module of spiking transformers, encodes Query, Key, and Value with sparse spikes. However, the Figure 1 : Illustration of Q-K attention with the two versions of Q-K token attention (QKTA) and Q-K channel attention (QKCA). The inputs are binary spikes and there are only sparse additions and mask operations in Q-K attention. As a spike-driven module, Q-K attention efficiently models the token or channel attention through spike-form binary vectors, performing linear complexity to #tokens (or #channels) and high energy efficiency. Spiking Neuron (SN) in this work adopts the Leaky-Integrate-and-Fire (LIF) model, which is shown in Appendix. 7.1.computational complexity (especially space complexity) of SSA scales quadratically to the number of tokens (#tokens), and is the main obstacle to explore architecture that incorporate multi-level features.ii). SNNs process data across the time domain, necessitating a high level of computational and memory resources. This combination leads to considerable consumption of computational resources, making the training process highly demanding in terms of both memory and processing power.To address these issues, we propose QKFormer with three innovations. i) Q-K attention with linear complexity and high energy efficiency. ii) A hierarchical architecture with decreasing number of tokens across blocks. iii) A novel patch embedding with deformed shortcut module. The linear complexity of Q-K attention is originated from the binary spike-form vector attention. This design lower the energy consumption and the space requirement. The hierarchical architecture starts from small patches and gradually merges neighboring patches in deeper spiking transformer layers with gradually decreasing #tokens, which enables multi-level spiking feature representation and benefits the model performance. The patch embedding with deformed shortcut facilitates spiking information transmission and integration. These merits make QKFormer achieve state-of-the-art performance in the SNN domain, in contrast to the previous transformer-based SNNs with spiking feature maps of a single resolution. Our main contributions are as follows:1) We develop a novel spike-form Q-K attention mechanism, tailor-made for the spatio-temporal spiking patterns of SNNs, which can easily model the importance of token or channel dimensions with binary values. The Q-K attention has linear complexity to #tokens (or #channels) and only adopts two spike-form components: Query (Q) and Key (K).2) We design a versatile and powerful Spiking Patch Embedding with Deformed Shortcut (SPEDS) module, which enhances spiking information transmission and integration thus improving the performance of spiking transformers significantly.3) We build a direct-training hierarchical spiking transformer with different number of tokens across blocks, incorporating Q-K attention and SPEDS, named QKFormer. This marks the effective exploration of hierarchical spiking representation in Transformer-based SNNs.4) Extensive experiments show that the proposed model outperforms the state-of-the-art (SOTA) SNNs on various static and neuromorphic datasets. Notably, QKFormer has achieved a significant milestone, surpassing 85% top-1 accuracy on ImageNet with 4 time steps using the direct training approach for the first time.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: See the Abstract and Section.1.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related work: 
 Learning Methods of Spiking Neural Networks. At present, there are mainly two ways to obtain trained SNNs. One involves converting pre-trained ANNs to SNNs (ANN2SNN) [17, 18, 19, 20, 21, 22, 23] , replacing the ReLU activation function in ANN with spiking neurons. However, This converted SNN suffers from long converting time steps and constraints on the original ANN design.Another method is to directly train SNNs [24] , using surrogate gradient [25, 26, 27, 28] to address the non-differentiability of spike excitation function during backpropagation. The direct training method has received more attention due to its low latency and supporting flexible architectural exploration.Direct Trained SNN Models. [28] proposed the Spike-Element-Wise block, which further addressed gradient explosion and gradient vanishing problems, and prolonged the directly trained SNNs beyond a depth of 100 layers with 69.26% accuracy on ImageNet-1k. Spikformer [11] designed a novel spikeform self-attention named Spiking Self Attention (SSA), using sparse spike-form Query, Key, and Value without softmax operation, which was used to construct the Spikformer. Spikformer achieved 74.81% accuracy on ImageNet-1k with 4 time steps, showing the great potential of transformer-based SNNs for the first time. Spikingformer [12] modified Spikformer with a pre-activation shortcut, which can avoid the floating-point multiplications in synaptic computing and has a lower firing rate. Spikingformer achieved 75.85% accuracy on ImageNet-1k. [13] designed a novel Spike-Driven Self-Attention (SDSA), which used only masks and addition operations without any multiplication, thus significantly reducing the computation energy compared to the vanilla self-attention. In addition, the proposed Spike-driven Transformer based on SDSA has achieved 77.07% on ImageNet-1k. However, all of these SNN models above remain a large performance gap compared with ANN.3 MethodQuestion: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: See the Appendix. 7.9.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
preliminary: 
 Vanilla Self Attention. Vanilla self-attention (VSA) [3] in transformers has three floating-point key components: query (Q F ), key (K F ), value (V F ) which are calculated by learnable linear matrics and input X. The calculation of VSA can be formulated as follows:EQUATIONEQUATIONwhere F denotes the floating-point form. Both floating-point matrix multiplication and softmax operation which contains exponent calculation and division, do not align with the properties of SNNs.Spiking Self Attention. Spikformer [11] demonstrated a novel spike-form self-attention named Spiking Self Attention (SSA), using sparse spike-form Q, K, V without softmax operation and floating-point matrix multiplication. The calculation process of SSA is formulated as follows:EQUATIONEQUATIONwhere Q, K, V ∈ R T ×N ×D , the spike-form Q, K, V are computed by learnable linear layers. s is a scaling factor. SN means spiking neuron layer. The calculation of SSA avoids floating-point multiplication, meeting the property of SNNs.An overview of Q-K attention is shown in Figure 1 . Both VSA and SSA use three key components and have O(N 2 d) or O(N d 2 ) computational complexity, while our proposed Q-K Attention which has linear complexity and only uses two spike-form components: Q and K, which are produced through learnable linear matrics.EQUATION)where X is the input spiking map. According to the detailed calculation mechanism of Q, K, Q-K Attention can be divided into Q-K Token Attention (QKTA) and Q-K Channel Attention (QKCA).Q-K Token Attention. We here assume T = 1 and single head attention for mathematical description.After obtaining spike-form Q, K ∈ R T ×N ×D , both Q and K can be formed as a spike matrix N × D (N is the token number, D is the channel number). QKTA process can be formulated as follows:EQUATIONTable 1 : Computational complexity comparison. N is the token number, D is the channel number.Methods VSA [3] SSA [11] SDSA [13] QKTA QKCAO(N 2 D) O(N 2 D) O(N D) O(D) O(N ) Space complexity O(N 2 + N D) O(N 2 + N D) O(N D) O(N ) O(D)where A t is the N * 1 token attention vector, which models the binary importance of different tokens.A t is a spike-form vector, which is obtained by addition operations (row summation) of Q spike matrix and a following spiking neuron. ⊗ is the Hadamard product between spike tensors, which is equivalent to the mask operation. We apply the spike-form token attention vector A t to the K spike matrix through the column mask operation (token mask), to obtain the output X ′ of QKTA.Q-K Channel Attention. The calculation process of Q-K channel attention is similar to the previous Q-K token attention, and can be formulated as :A c = SN( N j=0 Q i,j ), X ′ = A c ⊗ K, where A c is the 1 * D channel attention vector, which models the binary importance of different channels. A t is a spike-form vector, which is obtained by addition operations (column summation) of Q spike matrix and a following spiking neuron. Then, the output X ′ of Q-K Channel Attention is obtained by the row mask operation (channel mask) between A t and K.EQUATIONAs shown in Formula.7, a post-linear layer is also required after obtaining X ′ of Q-K Token or Channel Attention. In addition, the channel dimension is D/h in the multi-head Q-K attention, where h is the head number. In this work, the spiking neuron uses the LIF model [28] . Same with [11] , time step T is an independent dimension for the spiking neuron layer. In other layers, it is merged with the batch size. We exploit QKTA in our experiments by default.Linear Computational Complexity of Q-K Attention. As shown in Table 1 , the time complexity of Q-K attention varies based on the implementation approach. Specifically, when utilizing spike-form broadcasted element-wise multiplication, ⊗, the time complexity can reach up to O(N * D). When applying mask operation, the time complexity of Q-K attention is only O(N ) or O(D). The space complexity of Q-K attention with the whole process is O(N * D) at most, which is caused by the self-storage consumption Q and K matrix. In terms of the space complexity of attention operation, Q-K attention only requires an extra 1 * D or N * 1 space to store the attention vector with the space complexity of O(N ) or O(D). The linear complexity of Q-K attention makes it possible to successfully explore the large-scale hierarchical architecture SNN model.High Energy Efficiency of Q-K Attention. As a spike-driven attention module, the linear multiplication is transformed into sparse addition. Mask operation can be implemented on neuromorphic chips through addressing algorithms [29] or AND logic operations [30] with negligible power consumption. Compared with SSA, Q-K attention is much more energy-efficient, which comes from the following reasons: i) Q-K attention only adopts two spike-form components for spike [0, 1] operation without the V input and thus has less synaptic computing. ii) Q-K attention has much fewer spiking matrix operations due to its linear complexity of O(N ) or O(D). iii) Q-K attention discards the scale operation of SSA, which leads to reduced power consumption further.In VSA [3] , assume that q i q i ∈ R 1×d , Q ∈ R m×d and k i k i ∈ R 1×d , K ∈ R m×d are independent random variables with a mean of 0 and a variance of 1, then each element in the product of QK T has mean 0 and variance d. The variance magnitude of QK T grows with the embedding dimension d, which can result in gradient vanishing issues after softmax operation. Therefore, The product of matrices Q and K in VSA [3] is scaled by a factor 1 √ d in Eq. 2 to normalize the product to variance 1. Though the softmax function is not adopted due to its non-spike operations (division, exponential operation) in SNNs, SSA-based [11] SNNs will suffer obvious performance degradation even cannot converge without scaling because the variance of QK T V output is too large (Assuming that all the spiking elements are independent random variables and subject to Bernoulli Distribution). However, Q-K attention can discard scaling operations thus reducing power consumption because the variance of Q-K attention is much smaller than SSA (e.g. the max theoretical variance of Q-K token attention is only about 1 / 200 of SSA). The detailed analysis can be found in the Appendix.7.2 and Section.4.3.As the computational complexity (especially space complexity) of SSA is quadratic to #tokens, previous direct training spiking transformers are all limited to straight-through structures. Combining SSA with hierarchical architecture directly will lead to memory explosion easily when training spiking transformers. To overcome these issues, we proposed a hierarchical spiking transformer based on Q-K attention, named QKFormer, which constructs hierarchical spiking feature maps with linear computational complexity to #tokens or #channels. ). To produce a hierarchical spiking representation, the number of tokens is reduced in SPEDS-2 and SPEDS-3 as the network goes deeper. Both SPEDS-2 and SPEDS-3 reduce the number of tokens by a patch size of 2 × 2 (2× downsampling of resolution), and the number of channels is transformed into 2C and 4C, respectively. We denote the SPEDS-2 and the following QKFormer blocks as "Stage 2", which reduces the number of tokens ( H 8 × W 8 ). While SPEDS-3 and the following Spikformer or QKormer blocks are referred to as "Stage 3" with ( H 16 × W 16 ) tokens. The number of spiking transformer blocks (QKFormer or Spikformer) in each stage are N 1 , N 2 , and N 3 , respectively. These stages jointly produce a hierarchical spike-form representation.D N K  1  N A 1  N A D A  1 D A  1 D N X  D N X   D N X   (a) Q-K Token Attention (QKTA) (b) Q-K Channel Attention (QKCA) SPEDS-1 Q-K Attention SMLP Stage 1: 푇 × 퐻 4 × 푊 4 × 퐶 SPEDS-2 Q-K Attention SMLP Stage 2: 푇 × 퐻 8 × 푊 8 × 2퐶 SPEDS-3 SSA SMLP Stage 3: 푇 × 퐻 16 × 푊 16 × 4퐶 × 푁 × 푁 × 푁 QKFormer Blocks QKFormer Blocks Spikformer Blocks Input: 푇 × 퐻 × 푊 × 푛Mixed Spiking Attention Integration. In the former stage of a hierarchical architecture model, the number of channels is small while the number of tokens is large. In the latter stage, the channel number is large while the token number is small. Thus it leads to suboptimal performance when we only use a single type of Q-K attention in a hierarchical architecture model. Therefore, we use mixed spiking attention integration in QKFormer. QKTA is conducted in the former stage in hierarchical architecture, and we could choose QKCA or SSA in the latter stage. In the subsequent experiments, we use SSA in the last stage of QKFormer and QKTA in the former stages by default.QKFormer Blocks. Similar to the standard transformer encoder block, a QKFormer block contains a Q-K Attention module (QKTA or QKCA) and a Spiking MLP (SMLP) block, which can be formulated as follows:EQUATIONEQUATIONAt last, a fully connected layer is used as the classifier behind the last block.Residual shortcuts in SNNs [28] can implement identity mapping, which reduces information loss (facilitates information transmission and integration) in spike communication, thus ensuring the network can be well-behaved in a depth-insensitive way. Previous spiking transformers [11, 12, 13] use the residual shortcuts to achieve identity mapping, mainly focusing on the spiking attention block and spiking MLP block, and lacking identity mapping in patch embedding across the downsampling block. The input and output of a spiking patch embedding block in QKFormer have different channel and token numbers. To realize residual learning in spiking patch embedding, we can perform a lightweight linear projection W d in the shortcut connections to match the channel and token numbers, thus realizing the identity mapping cross downsampling blocks in spiking patch embedding. Given the input spiking map X, the process of patch embedding can be formulated as follows: In this work, the deformed linear projection W d is set as a lightweight convolutional layer with 1 × 1 kernel and stride > 1, to meet the channel and token numbers of the patch embedding block.EQUATIONThe function F involved in this work is set as {Conv2D-BN-MaxPooling-SN-Conv2D-BN-SN} or {Conv2D-BN-SN-Conv2D-BN-MaxPooling-SN}, while more layers or more variants are possible.There are mainly two types of residual shortcuts in deep SNNs. Formula.10 shows the patch embedding in the way of activation-before-addition [28, 11] . The other way of the patch embedding with the pre-activation residual shortcut [31, 12, 13 ] can be formulated as follows:EQUATIONwhere the function G correspondingly could be set as {Conv2D-BN-MaxPooling-SN-Conv2D-BN} or {Conv2D-BN-SN-Conv2D-BN-MaxPooling}. The intuitive representation of SPEDS is shown in Appendix 7.4.In this work, the spiking patch embedding of stage 2 or stage 3 in QKFormer can be formulated as Formula.10. The spiking patch embedding in stage 1 uses an extra {Conv2D-BN-SN} for spiking encoding in front of the block (Formula.10) to transform the non-spike input data into spikes. [33] and distributed across 8 Nvidia V100 GPUs. We trained QKFormer for 200 epochs. In addition, following DeiT [32] , data augmentation techniques including RandAugment [34] , random erasing [35] , and stochastic depth [36] are employed in this study. The number of blocks in the three stages is set as {1, 2, 7} respectively.Main Results on ImageNet. The experimental results demonstrate the superior performance of our proposed QKFormer, surpassing previous works' performance by a large margin (Table 2 ). QKFormer (64.96 M) achieves 85.65% top-1 accuracy and 97.74% top-5 accuracy on ImageNet. To begin with, we compare our model with the baseline spiking transformer (i.e., Spikformer [11] ). Our QKFormer models have slightly fewer parameters but much higher performance. For example, our QKFormer (64.96 M, 85.65%) significantly outperforms Spikformer (66.34 M, 74.81%) by 10.84%. In addition, compared with SDSA, our Q-K attention has lower computational complexity (shown in Table 1 ) and our QKFormer has much higher performance than S-Transformer (built by SDSA) [13] . In detail, QKFormer outperforms S-Transformer by 7.55%, 7.47%, and 8.58% respectively on three models with comparable #parameters. Finally, Our QKFormer outperforms the SOTA ANN-to-SNN model MST [23] 3 . For DVS128-Gesture dataset, our model with 1.50 M parameters achieves 98.6% accuracy using 16 time steps and 98.3% accuracy using 10 time steps. For CIFAR10-DVS dataset, our model achieves 84.0% accuracy with only 1.50 M parameters using 16 time steps. Our proposed QKFormer significantly outperforms Spikformer by 3.1% while reducing 1.07 M parameters. In addition, our model with 10 time steps achieves 83.8% accuracy, which outperforms Spikformer by 4.9% and outperforms the SOTA model (Spikingformer) by 3.9%.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer Justification: See Appendix. 7.5.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
analyses on q-k attention: 
 Attention Visualization. In this part, we visualize the Q-K token attention (Stage 1 and Stage 2 of the QKFormer model) on ImageNet. As shown in Figure 3 (a), A t is the N * 1 token attention vector, and X ′ is the output of the attention process, which is obtained by the mask operation between matrix K and attention vector A t . Specifically, the longitudinal axis denotes the channel index of one head, while the horizontal axis denotes the token index. The #tokens in stage 1 and stage 2 are 56 2 and 28 2 , respectively. To facilitate visualization, we choose a continuous segment with a length of 100 extracted from the whole token vector. The visualization shows Q-K attention can lead to high sparsity of spikes.Memory Consumption. In this experiment, we compare the memory consumption between QKTA (Formula.6) and SSA (Formula.4) under different token numbers. We calculate the memory consumption of a QKTA and an SSA on a GPU by forwarding the input tensor (T, B, C, N ). To facilitate the statistics of the impact of #tokens N on memory consumption, the #channels C is set to 256, and the time step T and batch size B are set to 1. The experiment result is shown in Figure 3(b) . With the increment of #Tokens, SSA consumes much more GPU memory than QKTA, of which the complexity is linear to #Tokens. For example, SSA consumes about 10× GPU memory than QKTA when We can find SSA has a much larger variance and expectation than QKTA on the whole. For example, the maximum theoretical variance of QKTA is 16, but the maximum theoretical variance of SSA is over 3000. This is the main reason that Q-K attention can discard scaling operations thus reducing power consumption, but SSA can not.√ N = 50.SPEDS Module. In this experiment, We replaced the Spiking Patch Splitting (SPS) module in Spikformer with Spiking Patch Embedding with Deformed Shortcut (SPEDS) module, while other conditions remain unchanged. The results (Table 5 ) show that the SPEDS module is essential to QKFormer on both static and neuromorphic datasets. In addition, the addition of SPEDS to Spikformer leads to great gains in CIFAR100 (+2.05%) and CIFAR10-DVS (+1.30%), which further verified the effectiveness of SPEDS. 6 ) show that using a single type of Q-K attention (QKTA or QKCA only) in a hierarchical architecture model leads to suboptimal performance. In particular, the performance decline in QKTA is more obvious. While the mixed spiking attention solutions, such as QKFormer(QKTA + QKCA), QKFormer(QKTA + SSA), and QKFormer(QKCA + SSA) can achieve comparable performance to QKFormer(SSA) while requiring fewer parameters and much fewer memory resources (Figure 3 (b)). Consequently, the mixed spiking attention solutions are wellsuited for larger architectures and more challenging scenarios when considering both computational efficiency and performance. 7 . In this block, we replaced the Activation-Before-Addition (ABA) [28, 11] residual connection of QKFormer with the Pre-activation (PA) [31, 37] way, and the performance slightly improved. In addition, we replaced the LIF spiking neuron with Integrateand-Fire (IF) and Parametric-Leaky-Integrateand-Fire (PLIF) [38] , which led to slight performance degradation. The accuracy regarding different simulation time steps of QKFormer is shown in the last column. When we increase the simulation time steps, the performance of QKFormer can be further improved. Specifically, QKFormer achieves 81.30 % accuracy on CIFAR 100 when T=6. 
conclusion: 
 In this work, we design a novel spike-form Q-K attention considering the properties of SNNs, which can easily model the importance of token or channel dimensions through binary vectors. Q-K attention has linear complexity to #tokens (or #channels) and only adopts two spike-form components: Query (Q) and Key (K). We design a versatile and powerful Spiking Patch Embedding with Deformed Shortcut (SPEDS), which enhances spiking information transmission and integration, thus improving the performance of spiking transformers. In addition, we develop a hierarchical spiking transformer based on the proposed Q-K attention and SPEDS in a direct training way, named QKFormer, which marks the effective exploration of hierarchical spiking representation in Transformer-based SNNs.Extensive experiments show that the proposed model achieves SOTA performance on both static and neuromorphic datasets. Note that QKFormer achieved top-1 accuracy beyond 85% on ImageNet-1k with 4 time steps using the direct training way for the first time. With its powerful performance, we aim for our investigations to instill optimism in the application of SNNs.Limitation. Currently, our model is limited to image / DVS classification tasks. We will extend this work to more tasks, such as segmentation, detection, and language tasks, to test the generalizability in the further. In addition, we will explore efficient and high-performance network architectures with fewer time steps based on Q-K attention and other efficient modules, to further reduce the training consumption. 
spiking neuron model: 
 Spiking neuron is the fundamental unit of SNNs, we choose the Leaky Integrate-and-Fire (LIF) model as the spiking neuron in our work. The dynamics of a LIF neuron can be formulated as follows:EQUATIONEQUATIONEQUATIONwhere τ is the membrane time constant, and XEQUATIONEQUATIONEQUATION)EQUATIONwhere d = D/h is the feature dimension of a head in the multi-head Q-K attention and N is the token number.Mathematical Characteristics of SSA. Similar to the above analysis process, assume that all elementsQ i,j [t], K j,i [t], V i,j [t] in SSA are independent random variables and subject to Bernoulli distribution B(f Q ), B(f K ), B(f V ), respectively. f Q , f K and f V are the average firing rate of Q, Kand V, respectively. We can calculate the expectation and variance of SSA as follows.EQUATIONEQUATIONFigure .4 visualize the variance and expectation of QKTA (Formula.15 and 16) and SSA (Formula.19 and 20). N is set as 196 and d is set as 64, respectively. We can find SSA has a much larger variance and expectation than QKTA on the whole. For example, the maximum theoretical variance of QKTA is 16, but the maximum theoretical variance of SSA is over 3000. This is the main reason that Q-K attention can discard scaling operations thus reducing power consumption, but SSA can not.The Complexity of Attention Mechanisms. The computational complexity of SSA:Q, K ∈ [0, 1] N ×D . The attention map (Q × K T ∈ Z N ×N) is obtained by matrix multiplication of matrix [0, 1] N ×D and matrix [0, 1] D×N , which thus need O(N 2 D) computation. The computational complexity of SDSA:Q, K ∈ [0, 1] N ×D . The attention map (Q ⊗ K ∈ [0, 1] N ×D) is obtained by the Hadamard product of matrix [0, 1] N ×D and matrix [0, 1] N ×D , which thus need O(N D) computation. The computational complexity of Q-K Attention: Our attention vector(A t ∈ [0, 1] N ×1 ) is computed by A t = SN ( D i=0 Q i,j), which depends on the row or column accumulation of theQ matrix (Q ∈ [0, 1] N ×D ), thus only needs O(N ) or O(D) computation.PLIF for Scaling. Q-K attention can discard scaling operations to reduce power consumption On these datasets used in this article's experiments because the variance of Q-K attention is much smaller than SSA (e.g. the max theoretical variance of Q-K token attention is only about 1 / 200 of SSA). For generality, we can also replace the LIF after attention calculation with PLIF (LIF with trainable parameters) allowing for adaptively controlling the fire rate of that spiking neuron, which can be seen as a learnable scaling. It can be expressed asA t = PLIF( D i=0 Q i,j). The results show that this modification brings a 0.2% performance improvement on CIFAR 100 (Acc = 81.17%, the firing rate of A t is 0.2952 in stage1 and 0.4008 in stage 2), while increasing the training time to 1.3 times. Datasets. We evaluate QKFormer on static image classification and neuromorphic classification. The former includes ImageNet-1K [39] , CIFAR10/100 [40] . The latter contains CIFAR10-DVS [41] and DVS128 Gesture [42] . Training Details. In our experiments, we use 8 NVIDIA Tesla V100 SXM2 32GB GPUs when training models on ImageNet, while 1 GPU is used to train other datasets (CIFAR10, CIFAR100, DVS128 Gesture, CIFAR10-DVS). In direct training SNN models with surrogate function,ImageNetEQUATIONwe select the Sigmoid function σ(x) as the surrogate function with α = 4 during the backpropagation of direct training in all experiments.Experimental Details on CIFAR and Neuromorphic Classification. We evaluate our QKFormer on small-scale datasets, including CIFAR10, CIFAR100 [40] and temporal neuromorphic datasets (CIFAR10-DVS and DVS128 Gesture [42] ). The detailed results on the four small-scale datasets are presented in Table 8 . The homogeneity of convolution allows the following BN and linear scaling transformation to be equivalently fused into the convolutional layer with an added bias when deployment [45, 46, 47, 48] . Therefore, when calculating the theoretical energy consumption, the consumption of BN layers could be ignored. We calculate the number of Synaptic Operations (SOPs) of spike before calculating theoretical energy consumption for QKFormer.EQUATIONwhere l is a block/layer in QKFormer, f r is the firing rate of the block/layer and T is the simulation time step of spiking neuron. FLOPs l refers to floating point operations of block/layer l, which is the number of multiply-and-accumulate (MAC) operations. And SOP l is the number of spike-based accumulate (AC) operations. Refer to previous works [49, 31, 50, 11, 12, 51, 52] . we assume that the MAC and AC operations are implemented on the 45nm hardware [50] , where E M AC = 4.6pJ and E AC = 0.9pJ. The theoretical energy consumption of QKFormer can be calculated as follows:EQUATIONEq.23 shows the energy consumption of QKFormer. FLOP 1 Conv is the first layer encoding the non-spike input into spike-form. Then the SOPs of N SNN Conv layers, M QKTA layers, and Z SSA layers are added together and multiplied by E AC . For ANNs, the theoretical energy consumption can be calculated:EQUATION7.7 Supplementary for Memory Consumption in Experiment 4.3Table 9 shows the detailed values of Figure 4 in the main body of this paper (Experiment 4.3). We compare the memory consumption between QKTA (Formula.6) and SSA (Formula.4) under different token numbers, which is calculated on a GPU by forwarding the input tensor (T, B, C, N ). To facilitate the statistics of the impact of #tokens N on memory consumption, the #channels C are set to 256, and the time step T and batch size B are set to 1. The experiment result is shown in Figure 3(b) . With the increment of #Tokens, SSA consumes much more GPU memory than QKTA, of which the complexity is linear to #Tokens. For example, SSA consumes about 10× GPU memory than QKTA when √ N = 50.We organized the experiment to test the training and inference time of QKFormer compared with Spikformer. We carried out this experiment on ImageNet with an input size of 224*224. This Prospect. The human brain has powerful intelligence that runs with low power consumption, so how to develop novel artificial intelligence algorithms to achieve high performance with the low power consumption of the human brain level is one of the AI's ultimate goals. SNN is an attractive potential way to achieve it. QKFormer directly trained on ImageNet-1K has a groundbreaking leap forward with 10.84% accuracy improvement compared to the previous SNN model while maintaining the energy-efficient feature, which marks an important step towards this goal. Combined with pre-training in the future, the performance of QKFormer is expected to be further improved.Reproducibility. The experimental results in this paper are reproducible. All experiments are implemented based on Pytorch [53] , SpikingJelly [54] and Timm [55] . We explain the details of model training and configuration in the main text and Appendix. Our codes and models of QKFormer will be available on GitHub after review.NeurIPS Paper Checklist 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: This research is in every respect with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: This research belongs to foundational research and is not tied to particular applications.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This paper poses no such risks Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: See Appendix. 7.5.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix:  


Decision: Spiking neural networks that more closely resemble real neurons have long been of interest as a more biologically grounded alternative to deep networks, but their performance have long lagged behind. In this paper, a spiking neural network transformer model is presented that achieves a large improvement above the state of the art in spiking neural networks. Specific design decisions also make this model very efficient along with being very accurate (such as a new efficient spike-based attention module that allows the construction of larger models, a multi-scale spiking architecture that models a hierarchy of representations, a new embedding approach). The paper is well written and describes the approach clearly. The authors perform a comprehensive analysis of the time and space complexity of recent spiking neural network architecture on the main datasets typically used in such approaches, effectively comparing their approach in a sound way to existing work. The author also perform informative ablation analysis. The analyses in the paper were appreciated by the reviewers, who also requested additional clarifications and analyses. All these modifications should be added to the paper for the camera ready version.

Reviews: review: paper_summary: paper_summary: The authors proposed QKFormer, which pushes SNN to 85% accuracy in ImageNet, becoming a new sota and contributing to the SNN community.
strength_weakness: strength_weakness: 1. The current accuracy is high, but it still requires a large time step and a large amount of computation.
questions: 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?
limitations: limitations: See Questions and Weaknesses.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: See weaknesses and Questions
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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.
limitations: 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.

Comments: comment: The author answered my confusion and therefore raised the score.
I have reviewed the rebuttal and have decided to maintain my positive rating.
I would like to thank the authors for addressing some of my concerns. I want to keep my score.
Thanks for your rebuttal. I keep my score.