Quantized Spike-driven Transformer

Authors: Xuerui Qiu, Malu Zhang, Jieyuan Zhang, Wenjie Wei, Honglin Cao, Junsheng Guo, Rui-Jie Zhu, Yimeng Shan, Yang Yang, Haizhou Li
Primary area: foundation or frontier models, including LLMs
Keywords: Spiking Neural Network+Spike-driven+Quantized Spiking Transformer+ Neuromorphic Computing

Abstract:
Spiking neural networks (SNNs) are emerging as a promising energy-efficient alternative to traditional artificial neural networks (ANNs) due to their spike-driven paradigm.
However, recent research in the SNN domain has mainly focused on enhancing accuracy by designing large-scale Transformer structures, which typically rely on substantial computational resources, limiting their deployment on resource-constrained devices.
To overcome this challenge, we propose a quantized spike-driven Transformer baseline (QSD-Transformer), which achieves reduced resource demands by utilizing a low bit-width parameter. 
Regrettably, the QSD-Transformer often suffers from severe performance degradation.
In this paper, we first conduct empirical analysis and find that the bimodal distribution of quantized spike-driven self-attention (Q-SDSA) leads to spike information distortion (SID) during quantization, causing significant performance degradation. To mitigate this issue, we take inspiration from mutual information entropy and propose a bi-level optimization strategy to rectify the information distribution in Q-SDSA.
Specifically, at the lower level, we introduce an information-enhanced LIF to rectify the information distribution in Q-SDSA.
At the upper level, we propose a fine-grained distillation scheme for the QSD-Transformer to align the distribution in Q-SDSA with that in the counterpart ANN.
By integrating the bi-level optimization strategy, the QSD-Transformer can attain enhanced energy efficiency without sacrificing its high-performance advantage.
We validate the QSD-Transformer on various visual tasks, and experimental results indicate that our method achieves state-of-the-art results in the SNN domain.
For instance, when compared to the prior SNN benchmark on ImageNet, the QSD-Transformer achieves 80.3\% top-1 accuracy, accompanied by significant reductions of 6.0$\times$ and 8.1$\times$ in power consumption and model size, respectively. Code is available at https://github.com/bollossom/QSD-Transformer.

Decision: accept (poster)
Reviewer scores: 5, 6, 8, 6, 5

Reviews and discussion:
<thread>
creation date:
1730103661228

thread_type:
official_review

<thread_object>
creator: reviewer_ynnk

<content>
summary:
The author proposed Quantized Spike-Driven Transformer(QSD-Transformer) to tackle with the spike information distortion (SID) challenge resulted from quantized spike-driven self-attention (Q-SDSA). The author addressed the problems through two levels: 1) Information-Enhanced LIF(IE-LIF) neuron to rectify the information distribution in Q-SDSA at the lower level. 2) A fine- grained distillation scheme for the QSD-Transformer to align the distribution in Q-SDSA with that in the counterpart ANN at the upper level. QSD-Transformer achieved promising results on multiple computer vision tasks.

strengths:
1. Extensive experiments and ablation studies on Image Classification, Object Detection and Semantic Segmentation.
2. The proposed Information-Enhanced LIF(IE-LIF) neuron is effective to rectify the information distribution in Q-SDSA through Information Theory.
3. Clear writing and methodology.

weaknesses:
1. The comparison between ANN2SNN and Direct Training methods is limited. Currently, MST is no longer the SOTA method on ANN2SNN method. SpikeZIP-TF (ICML 2024) [1] and ECMT (ACM MM 2024) achieve better performance on Image Classification tasks. The performance of  SpikeZIP-TF and ECMT on ImageNet surpass QSD-Transformer by a large margin. In addition, ANN2SNN methods has advantage on saving computational resources compared with Direct Training methods. It is recommended that the author should conduct a more comprehensive comparison between those two methods.
2. The method proposed by the author is somewhat cumbersome, did the training time provided by the authors in appendix include the training time of FGD?
3.  It is recommended that the author should extend the method to NLP tasks to verify the transferability of QSD-Transformer.

[1] Kang You*, Zekai Xu* et al. SpikeZIP-TF: Conversion is All You Need for Transformer-based SNN. International Conference on Machine Learning 2024
[2] Zihan Huang, Xinyu Shi et al. Towards High-performance Spiking Transformers from ANN to SNN Conversion. ACM Multimedia 2024

questions:
1. I wonder whether you quantized the membrane potential or not. If you didn't quantize the membrane potential, it seems hard to implement your method on hardware.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730438330310

thread_type:
official_review

<thread_object>
creator: reviewer_uosk

<content>
summary:
This paper proposes a quantized spike-driven transformer and increases the accuracy of the SNN by proposing an information-enhanced LIF model and fine-grained distillation from the lower level and upper level respectively. Experiments show that these technologies reduce energy consumption significantly while increasing the accuracy of SNN.

strengths:
1. The paper is well-written and easy to follow.
2. The IE-LIF neuron combines the conversion algorithm and training algorithm, which is novel.
3. The experimental results are significant.

weaknesses:
1. The training of IE-LIF neurons does not utilize temporal information, which is not suitable for temporal benchmarks.
2. The reason why the IE-LIF neuron and fine-grained distillation that reduces the energy consumption is not provided.

questions:
1. Why do the IE-LIF neuron and fine-grained distillation that reduce the energy consumption? Do these technologies reduce the number of synaptic operations?
2. Why the energy reduction in the COCO2017 dataset is not significant?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730601890293

thread_type:
official_review

<thread_object>
creator: reviewer_ifne

<content>
summary:
To tackle the issues of high-parameter spike-based Transformer in resource-constrained applications and the low performance of directly quantized spike-based Transformers, this paper introduces a Quantized Spike-driven Transformer. It uses a bi-level optimization strategy, including an Information-Enhanced LIF neuron and fine-grained distillation, to counteract quantization-induced performance degradation. The comparative and ablation experiments demonstrate the effectiveness of the proposed methods.

strengths:
1. Based on Information entropy, this paper proposes a bi-level optimization strategy, which mitigates quantization-induced performance drops in baseline quantized spiking Transformers. 

2. The proposed IE-LIF spike neuron is hardware-friendly and converts to a binary-spiking LIF neuron during inference to maintain spike-driven nature.

3. Experimental results on ImageNet and a large number of vision tasks (detection, segmentation, and transfer learning) show that the method is effective and energy-efficient on various spike-based transformers.

weaknesses:
1. The authors describe the implementation method of surrogate gradients for binary spikes in the appendix. However, the proposed IE-LIF in the main text is multi-bit. Could the authors explain how the aforementioned surrogate gradients are deployed in the proposed neurons?

2. Fast-SNN [1] converts quantized ANNs to SNNs and is a competitive ANN-to-SNN method. Like this paper, it aims to reduce quantization error and employs quantization techniques to enhance energy efficiency. The basic idea of Fast-SNN is to reduce both quantization error and accumulating error, achieving excellent performance on many vision tasks (classification, detection, and segmentation). Could you include comparative experimental results with Fast-SNN?

3. There are some typos in the text, such as the equation (1) on line 164  and on line 292, it seems you intended to reference Fig 1(b). 

[1] Hu, Yangfan, et al. "Fast-SNN: fast spiking neural network by converting quantized ANN." IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 45, no. 12, pp. 14546-14562, Dec. 2023, doi: 10.1109/TPAMI.2023.3275769.

questions:
See Weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730603142641

thread_type:
official_review

<thread_object>
creator: reviewer_pktr

<content>
summary:
This paper presents QSD-Transformer, a quantized spike-driven transformer that addresses the challenge of implementing efficient spiking neural networks (SNNs) while maintaining performance. The work shows three key points: (1) a lightweight quantized baseline for spike-driven transformers, (2) a bi-level optimization strategy to address spike information distortion (SID), and (3) information-enhanced LIF (IE-LIF) neurons with fine-grained distillation for improved performance.

strengths:
- First approach to quantizing spike-based transformers with SID analysis
- Solid theoretical analysis and proofs for the proposed methods
- Competitive accuracy on ImageNet (80.3%) with low energy
- Extensive experiments on multiple tasks (classification, object detection, segmentation)

weaknesses:
- Lack of comparison with the previous SNN and quantized ANN transformer models (Connected to question #1)
- Limited scalability in ImageNet experiments (Connected to question #2)
- Huge training overhead compared to the conventional spike-based transformer due to multi-bit spike and knowledge distillation (Connected to question #3)

questions:
1. To make the ImageNet results table solid, authors can add additional results such as QKFormer [1] and previous QAT-ViT models [2].
2. I just wondered why only small sizes of architecture (1.8M, 3.9M, 6.8M) are used for training the ImageNet dataset. Is there any scalability issue with this method?
3. This work uses multi-bit spikes during training and knowledge distillation with ANN architecture, which causes huge training overhead in training time and memory. Can you present any analysis of this overhead?
4. In the transfer learning section, the overall information is insufficient. Which bit-width did you use? and could you provide us the accuracy of CIFAR10/100, and CIFAR10-DVS without transfer learning?
5. Can the authors provide the firing rate information? Compared to the original Spike-driven Transformer-V2, how has the firing rate changed in the self-attention part?

[1] Zhou, Chenlin, et al. "QKFormer: Hierarchical Spiking Transformer using QK Attention." arXiv preprint arXiv:2403.16552 (2024).
[2] Li, Yanjing, et al. "Q-vit: Accurate and fully quantized low-bit vision transformer." Advances in neural information processing systems 35 (2022): 34451-34463.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731466652049

thread_type:
official_review

<thread_object>
creator: reviewer_xmwz

<content>
summary:
This paper presents a Quantized Spike-Driven Transformer, addressing the spike information distortion during quantization caused by the bimodal distribution of Quantized Spike-Driven Self-Attention (Q-SDSA). A bi-level optimization strategy is introduced, incorporating information-enhanced LIF and fine-grained distillation to rectify the distribution of Q-SDSA.

strengths:
1. This paper aims to address the problem of performance degradation of the Spike Transformer after quantization, attributing it to the spiking information distortion (SID) problem. The authors presents a loss function based on mutual information maximization to tackle the problem.
2. The authors conduct experiments across various tasks to demonstrate the effectiveness of the proposed method.
3. The paper is well organized and easy to follow.

weaknesses:
1. The primary reason for the improved quantization performance of the proposed method is the use of multi-bit spikes instead of traditional 0-1 spikes. Specifically, the implementation extends to 4 virtual timesteps, which inevitably reduces the training and inference efficiency. However, the manuscript does not provide a detailed explanation or analysis of this trade-off, which would be beneficial for understanding the overall impact on efficiency.
2. The empirical comparison can be done in a more thoroughly by comparing with other latest state-of-the-art methods.
3. Some content in the paper seems unnecessary, such as Appendix A, which does not contribute significantly to the main arguments or findings and could be omitted for conciseness.

questions:
Please see the weaknesses.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes the QSD-Transformer which is a quantized framework for spike-based Transformers by introducing a bi-level optimization strategy, incorporating information-enhanced LIF, and fine-grained distillation to rectify the attention distribution.  Most of the reviewers have positive comments on this work and part of them raised their score after the response. Thus, this paper can be accepted and please prepare the final version well.