SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Man Yao $^{1,2*}$ , Jiakui Hu $^{3,1*}$ , Zhaokun Zhou $^{3,2*}$ , Li Yuan $^{3,2}$ , Yonghong Tian $^{3,2}$ , Bo Xu $^{1}$ , Guoqi Li $^{1\dagger}$
$^{1}$ Institute of Automation, Chinese Academy of Sciences, Beijing, China  
 $^{2}$ Peng Cheng Laboratory, Shenzhen, Guangzhou, China  
 $^{3}$ Peking University, Beijing, China
Spiking Neural Networks (SNNs) provide an energy-efficient deep learning option due to their unique spike-based event-driven (i.e., spike-driven) paradigm. In this paper, we incorporate the spike-driven paradigm into Transformer by the proposed Spike-driven Transformer with four unique properties: i) Event-driven, no calculation is triggered when the input of Transformer is zero; ii) Binary spike communication, all matrix multiplications associated with the spike matrix can be transformed into sparse additions; iii) Self-attention with linear complexity at both token and channel dimensions; iv) The operations between spike-form Query, Key, and Value are mask and addition. Together, there are only sparse addition operations in the Spike-driven Transformer. To this end, we design a novel Spike-Driven Self-Attention (SDSA), which exploits only mask and addition operations without any multiplication, and thus having up to  $87.2 \times$  lower computation energy than vanilla self-attention. Especially in SDSA, the matrix multiplication between Query, Key, and Value is designed as the mask operation. In addition, we rearrange all residual connections in the vanilla Transformer before the activation functions to ensure that all neurons transmit binary spike signals. It is shown that the Spike-driven Transformer can achieve  $77.1\%$  top-1 accuracy on ImageNet-1K, which is the state-of-the-art result in the SNN field. All source code and models are available at https://github.com/BICLab/Spike-Driven-Transformer.
One of the most crucial computational characteristics of bio-inspired Spiking Neural Networks (SNNs) [1] is spike-based event-driven (spike-driven): i) When a computation is event-driven, it is triggered sparsely as events (spike with address information) occur; ii) If only binary spikes (0 or 1) are employed for communication between spiking neurons, the network's operations are synaptic ACumulate (AC). When implementing SNNs on neuromorphic chips, such as TrueNorth [2], Loihi [3], and Tianjic [4], only a small fraction of spiking neurons at any moment being active and the rest being idle. Thus, spike-driven neuromorphic computing that only performs sparse addition operations is regarded as a promising low-power alternative to traditional AI [5, 6, 7].
Although SNNs have apparent advantages in bio-plausibility and energy efficiency, their applications are limited by poor task accuracy. Transformers have shown high performance in various tasks for their self-attention [8, 9, 10]. Incorporating the effectiveness of Transformer with the high energy efficiency of SNNs is a natural and exciting idea. There has been some research in this direction, but all so far have relied on "hybrid computing". Namely, Multiply-and-ACcumulate (MAC) operations dominated
† Corresponding author, guoqi.li@ia.ac.cn
Figure 1: Comparison Vanilla Self-Attention (VSA) and our Spike-Driven Self-Attention (SDSA). (a) is a typical Vanilla Self-Attention (VSA) [8]. (b) are two equivalent versions of SDSA. The input of SDSA are binary spikes. In SDSA, there are only mask and sparse additions. Version 1: Spike  $Q$  and  $K$  first perform element-wise mask, i.e., Hadamard product; then column summation and spike neuron layer are adopted to obtain the binary attention vector; finally, the binary attention vector is applied to the spike  $V$  to mask some channels (features). Version 2: An equivalent version of Version 1 (see Section 3.3) reveals that SDSA is a unique type of linear attention (spiking neuron layer is the kernel function) whose time complexity is linear with both token and channel dimensions. Typically, performing self-attention in VSA and SDSA requires  $2N^2 D$  multiply-and-accumulate and  $0.02ND$  accumulate operations respectively, where  $N$  is the number of tokens,  $D$  is the channel dimensions,  $0.02$  is the non-zero ratio of the matrix after the mask of  $Q$  and  $K$ . Thus, the self-attention operator between the spike  $Q$ ,  $K$ , and  $V$  has almost no energy consumption.
by vanilla Transformer components and AC operations caused by spiking neurons both exist in the existing spiking Transformers. One popular approach is to replace some of the neurons in Transformer with spiking neurons to do a various of tasks [11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22], and keeping the MAC-required operations like dot-product, softmax, scale, etc.
Though hybrid computing helps reduce the accuracy loss brought on by adding spiking neurons to the Transformer, it can be challenging to benefit from SNN's low energy cost, especially given that current spiking Transformers are hardly usable on neuromorphic chips. To address this issue, we propose a novel Spike-driven Transformer that achieves the spike-driven nature of SNNs throughout the network while having great task performance. Two core modules of Transformer, Vanilla Self-Attention (VSA) and Multi-Layer Perceptron (MLP), are re-designed to have a spike-driven paradigm.
The three input matrices for VSA are Query  $(Q)$ , Key  $(K)$ , and Value  $(V)$ , (Fig. 1(a)).  $Q$  and  $K$  first perform similarity calculations to obtain the attention map, which includes three steps of matrix multiplication, scale and softmax. The attention map is then used to weight the  $V$  (another matrix multiplication). The typical spiking self-attention in the current spiking Transformers [20, 19] would convert  $Q$ ,  $K$ ,  $V$  into spike-form before performing two matrix multiplications similar to those in VSA. The distinction is that spike matrix multiplications can be converted into addition, and softmax is not necessary [20]. But these methods not only yield large integers in the output (thus requiring an additional scale multiplication for normalization to avoid gradient vanishing), but also fails to exploit the full energy-efficiency potential of the spike-driven paradigm combined with self-attention.
We propose Spike-Driven Self-Attention (SDSA) to address these issues, including two aspects (see SDSA Version 1 in Fig. 1(b)): i) Hadamard product replaces matrix multiplication; ii) matrix column-wise summation and spiking neuron layer take the role of softmax and scale. The former can be considered as not consuming energy because the Hadamard product between spikes is equivalent to the element-wise mask. The latter also consumes almost no energy since the matrix to be summed column-by-column is very sparse (typically, the ratio of non-zero elements is less than 0.02). We also observe that SDSA is a special kind of linear attention [23, 24], i.e., Version 2 of Fig. 1(b). In this view, the spiking neuron layer that converts  $Q$ ,  $K$ , and  $V$  into spike form is a kernel function.
Additionally, existing spiking Transformers [20, 12] usually follow the SEW-SNN residual design [25], whose shortcut is spike addition and thus outputs multi-bit (integer) spikes. This shortcut can satisfy event-driven, but introduces integer multiplication. We modified the residual connections throughout the Transformer architecture as shortcuts between membrane potentials [26, 27] to address this issue (Section 3.2). The proposed Spike-driven Transformer improves task accuracy on both static and neuromorphic event-based datasets. The main contributions of this paper are as follows:
- We propose a novel Spike-driven Transformer that only exploits sparse addition. This is the first time that the spike-driven paradigm has been incorporated into Transformer, and the proposed model is hardware-friendly to neuromorphic chips.
- We design a Spike-Driven Self-Attention (SDSA). The self-attention operator between spike Query, Key, Value is replaced by mask and sparse addition with essentially no energy consumption. SDSA is computationally linear in both tokens and channels. Overall, the energy cost of SDSA (including Query, Key, Value generation parts) is  $87.2 \times$  lower than its vanilla self-attention counterpart.
- We rearrange the residual connections so that all spiking neurons in Spike-driven Transformer communicate via binary spikes.
- Extensive experiments show that the proposed architecture can outperform or comparable to State-of-The-Art (SOTA) SNNs on both static and neuromorphic datasets. We achieved  $77.1\%$  accuracy on ImageNet-1K, which is the SOTA result in the SNN field.
Bio-inspired Spiking Neural Networks can profit from advanced deep learning and neuroscience knowledge concurrently [28, 5, 7]. Many biological mechanisms are leveraged to inspire SNN's neuron modeling [1, 29], learning rules [30, 31], etc. Existing studies have shown that SNNs are more suited for incorporating with brain mechanisms, e.g., long short-term memory [32, 33], attention [34, 27, 35], etc. Moreover, while keeping its own spike-driven benefits, SNNS have greatly improved its task accuracy by integrating deep learning technologies like network architecture [26, 25, 36], gradient backpropagation [37, 38], normalization [39, 40], etc. Our goal is to combine SNN and Transformer architectures. One way is to discretize Transformer into spike form through neuron equivalence [41, 42], i.e., ANN2SNN, but this requires a long simulation timestep and boosts the energy consumption. We employ the direct training method, using the first SNN layer as the spike encoding layer and applying surrogate gradient training [43].
Neuromorphic Chips. As opposed to the compute and memory separated processors used in ANNs, neuromorphic chips use non-von Neumann architectures, which are inspired by the structure and function of the brain [5, 7, 28]. Because of the choice that uses spiking neurons and synapses as basic units, neuromorphic chips [44, 45, 46, 2, 47, 3, 4] have unique features, such as highly parallel operation, collocated processing and memory, inherent scalability, and spike-driven computing, etc. Typical neuromorphic chips consume tens to hundreds of mWs [48]. Conv and MLP in neuromorphic chips are equivalent to a cheap addressing algorithm [49] since the sparse spike-driven computing, i.e., to find out which synapses and neurons need to be involved in the addition operation. Our Transformer design strictly follows the spike-driven paradigm, thus it is friendly to deploy on neuromorphic chips.
Efficient Transformers. Transformer and its variants have been widely used in numerous tasks, such as natural language processing [8, 50, 51] and computer vision [52, 53, 54]. However, deploying these models on mobile devices with limited resources remains challenging because of their inherent complexity [24, 55]. Typical optimization methods include convolution and self-attention mixing [55, 56], Transformer's own mechanism (token mechanism [57, 58], self-attention [59, 60], multi-head [61, 62] and so on) optimization, etc. An important direction for efficient Transformers is linear attention on tokens since the computation scale of self-attention is quadratic with token number. Removing the softmax in self-attention and re-arranging the computation order of Query, Key, and Value is the main way to achieve linear attention [60, 63, 23, 64, 65, 62, 66]. For the spiking Transformer, softmax cannot exist, thus spiking Transformer can be a kind of linear attention.
Figure 2: The overview of Spike-driven Transformer. We follow the network structure in [20], but make two new key designs. First, we propose a Spike-driven Self-Attention (SDSA) module, which consists of only mask and sparse addition operations (Fig. 1(b)). Second, we redesign the shortcuts in the whole network, involving position embedding, self-attention, and MLP parts. As indicated by the red line, the shortcut is constructed before the spike neuron layer. That is, we establish residual connections between membrane potentials to make sure that the values in the spike matrix are all binary, which allows the multiplication of the spike and weight matrices to be converted into addition operations. By contrast, previous works [20, 19] build shortcut between spike tensors in different layers, resulting in the output of spike neurons as multi-bit (integer) spikes.
3 Spike-driven Transformer
We propose a Spike-driven Transformer, which incorporates Transformer into the spike-driven paradigm with only sparse addition. We first briefly introduce the spike neuron layer, then introduce the overview and components of the Spike-driven Transformer one by one.
The spiking neuron model is simplified from the biological neuron model [67, 29]. Leaky Integrate-and-Fire (LIF) spiking neuron [1], which have biological neuronal dynamics and are easy to simulate on a computer, is uniformly adopted in this work. The dynamics of the LIF layer [37] is governed by
$$
U [ t ] = H [ t - 1 ] + X [ t ], \tag {1}
$$
$$
S [ t ] = \operatorname {H e a} \left(U [ t ] - u _ {t h}\right), \tag {2}
$$
$$
H [ t ] = V _ {\text {r e s e t}} S [ t ] + (\beta U [ t ]) \left(\mathbf {1} - S [ t ]\right), \tag {3}
$$
where  $t$  denote the timestep,  $U[t]$  means the membrane potential which is produced by coupling the spatial input information  $X[t]$  and the temporal input  $H[t - 1]$ , where  $X[t]$  can be obtained through operators such as Conv, MLP, and self-attention. When membrane potential exceeds the threshold  $u_{th}$ , the neuron will fire a spike, otherwise it will not. Thus, the spatial output tensor  $S[t]$  contains only 1 or 0.  $\mathrm{Hea}(\cdot)$  is a Heaviside step function that satisfies  $\mathrm{Hea}(x) = 1$  when  $x \geq 0$ , otherwise  $\mathrm{Hea}(x) = 0$ .  $H[t]$  indicates the temporal output, where  $V_{reset}$  denotes the reset potential which is set after activating the output spiking.  $\beta < 1$  is the decay factor, if the spiking neuron does not fire, the membrane potential  $U[t]$  will decay to  $H[t]$ .
Fig. 2 shows the overview of Spike-driven Transformer that includes four parts: Spiking Patch Splitting (SPS), SDSA, MLP, and a linear classification head. For the SPS part, we follow the design in [20]. Given a 2D image sequence  $I \in \mathbb{R}^{T \times C \times H \times W}$ , the Patch Splitting Module (PSM), i.e., the first four Conv layers, linearly projects and splits it into a sequence of  $N$  flattened spike patches  $s$  with  $D$  dimensional channel, where  $T$  (images are repeated  $T$  times in the static dataset as input),  $C$ ,  $H$ , and  $W$  denote timestep, channel, height and width of the 2D image sequence. Another Conv layer
is then used to generate Relative Position Embedding (RPE). Together, the SPS part is written as:
$$
u = \operatorname {P S M} (I), \quad I \in \mathbb {R} ^ {T \times C \times H \times W}, u \in \mathbb {R} ^ {T \times N \times D} \tag {4}
$$
$$
s = \mathcal {S N} (u), \quad s \in \mathbb {R} ^ {T \times N \times D} \tag {5}
$$
$$
\mathrm {R P E} = \mathrm {B N} (\operatorname {C o n v 2 d} (s)), \quad \mathrm {R P E} \in \mathbb {R} ^ {T \times N \times D} \tag {6}
$$
$$
U _ {0} = u + \mathrm {R P E}, \quad U _ {0} \in \mathbb {R} ^ {T \times N \times D} \tag {7}
$$
where  $u$  and  $U_{0}$  are the output membrane potential tensor of PSM and SPS respectively,  $\mathcal{SN}(\cdot)$  denote the spike neuron layer. Note, in Eq 6, before executing  $Conv2d(\cdot)$ ,  $s \in \mathbb{R}^{T \times N \times D}$  will be transposed into  $s \in \mathbb{R}^{T \times C \times H \times W}$ . We pass the  $U_{0}$  to the  $L$ -block Spike-driven Transformer encoder, which consists of a SDSA and a MLP block. Residual connections are applied to membrane potentials in both SDSA and MLP block. SDSA provides an efficient approach to model the local-global information of images utilizing spike  $Q$ ,  $K$ , and  $V$  without scale and softmax (see Sec. 3.3). A Global Average-Pooling (GAP) is utilized on the processed feature from spike-driven encoder and outputs the  $D$ -dimension channel which will be sent to the fully-connected-layer Classification Head (CH) to output the prediction  $Y$ . The three parts SDSA, MLP and CH can be written as follows:
$$
S _ {0} = \mathcal {S N} \left(U _ {0}\right), \quad S _ {0} \in \mathbb {R} ^ {T \times N \times D} \tag {8}
$$
$$
U _ {l} ^ {\prime} = \operatorname {S D S A} \left(S _ {l - 1}\right) + U _ {l - 1}, \quad U _ {l} ^ {\prime} \in \mathbb {R} ^ {T \times N \times D}, l = 1 \dots L \tag {9}
$$
$$
S _ {l} ^ {\prime} = \mathcal {S N} \left(U _ {l} ^ {\prime}\right), \quad S _ {l} ^ {\prime} \in \mathbb {R} ^ {T \times N \times D}, l = 1 \dots L \tag {10}
$$
$$
S _ {l} = \mathcal {S N} \left(\operatorname {M L P} \left(S _ {l} ^ {\prime}\right) + U _ {l} ^ {\prime}\right), \quad S _ {l} \in \mathbb {R} ^ {T \times N \times D}, l = 1 \dots L \tag {11}
$$
$$
Y = \mathrm {C H} (\mathrm {G A P} (S _ {L})), \tag {12}
$$
where  $U_{l}^{\prime}$  and  $S_{l}^{\prime}$  represent membrane potential and spike tensor output by SDSA at  $l$ -th layer.
3.2 Membrane Shortcut in Spike-driven Transformer
Residual connection [68, 69] is a crucial basic operation in Transformer architecture. There are three shortcut techniques in existing Conv-based SNNs [27]. Vanilla Res-SNN [39], similar to vanilla Res-CNN [68], performs a shortcut between membrane potential and spike. Spike-Element-Wise (SEW) Res-SNN [25] employs a shortcut to connect the output spikes in different layers. Membrane Shortcut (MS) Res-SNN [26], creating a shortcut between membrane potential of spiking neurons in various layers. There is no uniformly standard shortcut in the current SNN community, and SEW shortcut is adopted by existing spiking Transformers [20, 12]. As shown in Eq. 7, Eq. 9 and Eq. 11, we leverage the membrane shortcut in the proposed Spike-driven Transformer for four reasons:
- Spike-driven refers to the ability to transform matrix multiplication between weight and spike tensors into sparse additions. Only binary spikes can support the spike-driven function. However, the values in the spike tensors are multi-bit (integer) spikes, as the SEW shortcut builds the addition between binary spikes. By contrast, as shown in Eq. 8, Eq. 10, Eq. 11,  $S\mathcal{N}$  is followed by the MS shortcut, which ensures that there are always only binary spike signals in the spike tensor.
- High performance. The task accuracy of MS-Res-SNN is higher than that of SEW-Res-SNN [26, 27, 25], also in Transformer-based SNN (see Table 5 in this work).
- Bio-plausibility. MS shortcut can be understood as an approach to optimize the membrane potential distribution. This is consistent with other neuroscience-inspired methods to optimize the internal dynamics of SNNs, such as complex spiking neuron design [70], attention mechanism [27], long short-term memory [33], recurrent connection [71], information maximization [72], etc.
- Dynamical isometry. MS-Res-SNN has been proven [26] to satisfy dynamical isometry theory [73], which is a theoretical explanation of well-behaved deep neural networks [74, 75].
3.3 Spike-driven Self-Attention
Vanilla Self-Attention (VSA). Given a float-point input feature sequence  $X \in \mathbb{R}^{N \times D}$ , float-point Query  $(Q)$ , Key  $(K)$ , and Value  $(V)$  in  $\mathbb{R}^{N \times D}$  are calculated by three learnable linear matrices, respectively. Standard scaled dot-product self-attention (Fig. 2(a)) is computed as [52]:
$$
\operatorname {V S A} (Q, K, V) = \operatorname {s o f t m a x} \left(\frac {Q K ^ {\mathrm {T}}}{\sqrt {d}}\right) V, \tag {13}
$$
where  $d = D / H$  is the feature dimension of one head and  $H$  is the head number,  $\sqrt{d}$  is the scale factor. The time complexity of VSA is  $O(N^2 D + N^2 D)$ .
Spike-Driven Self-Attention (SDSA) Version 1. As shown in Fig. 2(b) left part, given a spike input feature sequence  $S \in \mathbb{R}^{T \times N \times D}$ , float-point  $Q$ ,  $K$ , and  $V$  in  $\mathbb{R}^{T \times N \times D}$  are calculated by three learnable linear matrices, respectively. Note, the linear operation here is only addition, because the input  $S$  is a spike tensor. A spike neuron layer  $\mathcal{SN}(\cdot)$  follows, converting  $Q$ ,  $K$ ,  $V$  into spike tensor  $Q_{S}$ ,  $K_{S}$ , and  $V_{S}$ . SDSA Version 1 (SDSA-V1) is presented as:
$$
\operatorname {S D S A} (Q, K, V) = g \left(Q _ {S}, K _ {S}\right) \otimes V _ {S} = \mathcal {S N} \left(\operatorname {S U M} _ {\mathrm {c}} \left(Q _ {S} \otimes K _ {S}\right)\right) \otimes V _ {S}, \tag {14}
$$
where  $\otimes$  is the Hadamard product,  $g(\cdot)$  is used to compute the attention map,  $\mathrm{SUM}_{\mathrm{c}}(\cdot)$  represents the sum of each column. The outputs of both  $g(\cdot)$  and  $\mathrm{SUM}_{\mathrm{c}}(\cdot)$  are  $D$ -dimensional row vectors. The Hadamard product between spike tensors is equivalent to the mask operation.
Discussion on SDSA. Since the Hadamard product among  $Q_S$ ,  $K_S$ , and  $V_S$  in  $\mathbb{R}^{N\times D}$  (we here assume  $T = 1$  for mathematical understanding) can be exchanged, Eq. 14 can also be written as:
$$
\operatorname {S D S A} (Q, K, V) = Q _ {S} \otimes g \left(K _ {S}, V _ {S}\right) = Q _ {S} \otimes \mathcal {S N} \left(\operatorname {S U M} _ {\mathrm {c}} \left(K _ {S} \otimes V _ {S}\right)\right). \tag {15}
$$
Note, Eq.14 and Eq.15 are equivalent functionally. In this view, Eq.15 is a linear attention [23, 63] whose computational complexity is linear in token number  $N$  because  $K_{S}$  and  $V_{S}$  can participate in calculate first. This is thanks to the softmax operation in the VSA is dropped here. The function of softmax needs to be replaced by the kernel function. Specific to our SDSA,  $\mathcal{SN}(\cdot)$  is the kernel function. Further, we can assume a special case [62],  $H = D$ , i.e., the number of channels per head is one. After the self-attention operation is performed on the  $H$  heads respectively, the outputs are concatenated together. Specifically,
$$
\operatorname {S D S A} \left(Q ^ {i}, K ^ {i}, V ^ {i}\right) = \mathcal {S N} \left(Q ^ {i}\right) g \left(K ^ {i}, V ^ {i}\right) = \mathcal {S N} \left(Q ^ {i}\right) \mathcal {S N} \left(\mathcal {S N} \left(Q ^ {i}\right) ^ {\mathrm {T}} \odot \mathcal {S N} \left(V ^ {i}\right)\right), \tag {16}
$$
where  $Q^i, K^i, V^i$  in  $\mathbb{R}^{N \times 1}$  are the  $i$ -th vectors in  $Q, K, V$  respectively,  $\odot$  is the dot product operation. The output of  $g(K^i, V^i)$  is a scalar, 0 or 1. Since the operation between  $\mathcal{SN}(Q^i)$  and  $g(K^i, V^i)$  is a mask, the whole SDSA only needs to be calculated  $H = D$  times for  $g(K^i, V^i)$ . The computational complexity of SDSA is  $O(0 + ND)$ , which is linear with both  $N$  and  $D$  (see Fig. 1(b) right part). Vectors  $K^i$  and  $V^i$  are very sparse, typically less than 0.01 (Table 2). Together, the whole SDSA only needs about 0.02ND times of addition, and its energy consumption is negligible.
Interestingly, Eq. 15 actually converts the soft vanilla self-attention to hard self-attention, where the attention scores in soft and hard attention are continuous- and binary-valued, respectively [76]. Thus, the practice of the spike-driven paradigm in this work leverages binary self-attention scores to directly mask unimportant channels in the sparse spike Value tensor. Although this introduces a slight loss of accuracy (Table 5),  $\mathrm{SDSA}(\cdot)$  consumes almost no energy.
4 Theoretical Energy Consumption Analysis
Three key computational modules in deep learning are Conv, MLP, and self-attention. In this Section, We discuss how the spike-driven paradigm achieves high energy efficiency on these operators.
Spike-driven in Conv and MLP. Spike-driven combines two properties, event-driven and binary spike-based communication. The former means that no computation is triggered when the input is zero. The binary restriction in the latter indicates that there are only additions. In summary, in spike-driven Conv and MLP, matrix multiplication is transformed into sparse addition, which is implemented as addressable addition in neuromorphic chips [49].