Quantized Spike-driven Transformer

TL;DR

This paper proposes a quantized spike-driven Transformer baseline (QSD-Transformer), which achieves reduced resource demands by utilizing a low bit-width parameter and introduces an information-enhanced LIF to rectify the information distribution in Q-SDSA.

Abstract

Spiking neural networks are emerging as a promising energy-efficient alternative to traditional artificial neural networks 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. 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.

Figures (5)

• Figure 1: Overview of the QSD-Transformer. (a) Proposed quantized spike-driven self-attention (Q-SDSA) module, where the membrane potential is rectified and then sent to the information-enhanced LIF (IE-LIF) neuron. (b) Proposed IE-LIF spiking neuron model, which utilizes the multi-bit spike during training while the binary spike during inference. (c) Proposed fine-grained distillation scheme.
• Figure 2: (a) Accuracy of quantizing different modules in the SD-Transformer v2 and its same ANN Transformer. (b) The distribution of the attention module (blue shadow), and the probability density function curve of normal distribution (red line). Experiments are conducted on ImageNet, and 3 layers in SD-Transformer v2 and its same ANN Transformer are selected for illustration.
• Figure 3: (a) Comparative results of accuracy, power, and parameters on ImageNet. (b) Comparison of information distribution in Q-SDSA before and after using the proposed IE-LIF and FGD scheme.
• Figure 4: Visualization of results on COCO dataset. Our QSD-Transformer excels in the target detection task.
• Figure 5: Visualization of results on ADE20K dataset. Our QSD-Transformer excels in the segmentation task.

Theorems & Definitions (2)

• proof
• proof