One-step Spiking Transformer with a Linear Complexity

1 Introduction

Spiking Neural Network (SNN) is the third generation neural network inspired by the behavior of biological neurons [Maass, 1997]. While Artificial Neural Network (ANN) utilizes computationally intensive floating-point Multiplyand-Accumulate (MAC) operations, SNN significantly reduces energy consumption by employing event-driven binary spike singles [Roy et al., 2019]. In practice, SNN is widely acknowledged as an energy-efficient alternative to ANN and has been extensively used in tasks such as image classification [Wu et al., 2019] and object detection [Kim et al., 2020]. On the other hand, Transformer [Vaswani et al., 2017], initially designed for natural language processing, has become a prominent technique owing to its excellent performance.

∗Corresponding author.

However, transformer's energy consumption is often substantial due to the high computational cost required for training and inference. For example, Switch Transformer [Fedus et al., 2022] consumes 179,000 kWh of power, while the energy consumption for training GPT-3 [Brown et al., 2020] is estimated to be 1,278,000 kWh [Patterson et al., 2021]. To leverage the high performance of transformer and low energy use of SNN, great effort has been dedicated to incorporating energy-efficient spiking calculations into transformer models, known as spiking transformers [Mueller et al., 2021; Yao et al., 2021; Zhou et al., 2023; Yao et al., 2023b; Wang et al., 2023; Yao et al., 2023a].

Although spiking transformers have demonstrated great performance across various tasks, there is still room for improvement in both performance and energy consumption. Most spiking transformers still require a level of energy consumption too high to be suitable for deployment on mainstream neuromorphic hardware. The two contributing factors that often lead to a dramatic increase in energy consumption are (1) lengthy timesteps and (2) complex MAC operations. Firstly, spiking transformers require additional T timesteps of temporal information (T > 3) [Meng et al., 2022]. Separate training is needed for every individual timestep, hence naturally resulting in T× spatio-temporal overhead compared to vanilla transformer. More importantly, deployment on neuromorphic hardware favors spiking transformers with fewer timesteps for simplicity and efficiency. Secondly, complex MAC operation [Zhou et al., 2023] is an integral part of the self-attention mechanism in existing spiking transformers. That leads to a relatively high complexity of O(L2d), where L and d denote the sequence length and feature dimension of input, respectively. Hereby a load of floating-point matrix dot product operations are involved, defeating the original purpose of SNN design, i.e., minimizing computations.

Aiming to address the above issues, in this study, we present a one-step binarized spiking transformer, reducing the timesteps from multiple to only one. More specifically, we introduce Time Domain Compression and Compensation (TDCC) to mitigate the spatio-temporal overhead. The "Compression" here is to compress T timesteps of information into a single one, thus substantially reducing the computational cost. The subsequent "Compensation" is to compensate for the information loss in the time domain, caused by "Compression", which is not lossless. Combining both, TDCC can fully leverage the spatio-temporal information and achieve pre-compression performance, but with a much lower cost. Furthermore, we introduce the Spiking Linear Transformation (SLT) to avoid the use of O(L2d) self-attention mechanism. Several variants of transformers, including Sparse Transformer [Child et al., 2019], Linear Transformers [Katharopoulos et al., 2020], and FLASH [Hua et al., 2022], endeavor to achieve the same purpose. They however still rely on floating-point dot product, not ideal for SNNs. The proposed SLT, however, has two parts well suited for binary spike signals. The first is SAF, Spiking AttentionFree, which eliminates floating-point dot product computation for self-attention. The second is SGFF, Spiking Gate Feed-Forward, which can convert the floating-point elementwise product to a binary logical AND (&) operation.

The contributions of this study can be summarized as:

1. One-step Spiking Transformer (OST) is proposed that can achieve SOTA performance with low energy use.
2. Time Domain Compression and Compensation (TDCC) is introduced, which reduces timestep to only 1, significantly lowering transformer's spatio-temporal overhead.
3. Spiking Linear Transformation (SLT) is introduced, greatly reducing floating-point MAC operations.

2 RelatedWorks

2.1 Vision Transformers (ViTs)

ViTs achieved remarkable performance in various computer vision tasks, such as image classification [Dosovitskiy et al., 2020] and object detection [Carion et al., 2020]. Despite their success, one of the main challenges faced by ViTs is the high computational complexity caused by the self-attention mechanism [Khan et al., 2022]. The self-attention mechanism enables ViTs to selectively focus on different parts of the input image, effectively capturing global information. However, the calculation of dot products between each pair of input features becomes inevitable, resulting in a quadratic computational complexity, relative to the input sequence length.

Several recent methods have been proposed to address the above challenge. Sparse Transformer [Child et al., 2019] used sparse attention mechanisms as a remedy, while the Synthesizer [Tay et al., 2021] introduced approximate attention. AFT [Zhai et al., 2021] utilized attention-free models for the same purpose. These methods still involve plenty of floatingpoint operations, not ideal for SNNs. Besides, they all need costly MAC operations, not helpful for energy saving.

2.2 Spiking Neural Networks (SNNs)

Training SNNs poses a significant challenge because their spike signals are non-differentiable, rendering stochastic gradient descent (SGD) ineffective. Two typical strategies are often used in practice for the training: (1) converting ANN to SNN (ANN-to-SNN) [Rueckauer et al., 2017] and (2) direct training [Neftci et al., 2019]. More specifically, the ANNto-SNN strategy converts an ANN pre-trained with SGD to an SNN for a given task and often achieves competitive accuracy. For example, the works in [Mueller et al., 2021]

and [Wang et al., 2023] can obtain a nearly equivalent accuracy compared to the vanilla ANN transformer models. However, ANN-to-SNN requires hundreds of timesteps to converge, causing significant inference latency and energy consumption. On the other hand, the direct training strategy uses surrogate gradients to replace discrete spikes with continuous functions, thereby bridging ANN transformers into spiking transformers across various tasks. Nevertheless, they still rely on multiple timesteps, which lead to higher energy consumption. Zhou et al. [Zhou et al., 2023] proposed a spiking self-attention transformer using four timesteps, incurring nearly four times higher energy consumption compared to one timestep. In addition, Yao et al. [Yao et al., 2023b] proposed attention spiking neural networks which do utilize one timestep. Despite significantly reduced energy consumption, the accuracy of their models suffers greatly with one timestep, unless the data are augmented. In comparison, this paper focuses on designing a spiking transformer with only one timestep while still maintaining high accuracy.

2.3 Leaky Integrate and Fire (LIF)

Spiking neurons transmit information in SNNs. They generate binary spike signals for transmission and are the fundamental part of SNNs. A key signal generation mechanism is LIF [Gerstner et al., 2014], which is also used in this study. We follow the conventions in [Zheng et al., 2021; Zhou et al., 2023], using SpikingJelly [Fang et al., 2023] to imply LIF. Their fire process is represented as Equation (1):

where uti is the membrane potential of the i-th neuron at the t-th timestep; τ is the decay constant; wij is the connection weight between the i-th and the j-th neurons; and ot-1 j is the spike signal generated by the j-th neuron at the (t - 1)-th timestep. When uti reaches the activation threshold, the spiking neuron fires. After the activation, the membrane potential is reset for the next cycle of accumulation.

3 Methodology

The proposed OST, the associated TDCC and SLT, and their sub-components are detailed in this section.

3.1 Overall Architecture

The overall architecture of OST is illustrated in Figure 1(a), where the proposed TDCC and SLT components are highlighted in different colors. The input of OST is a 2D image sequence I ∈ RT×C×H×W , where T , C, H, W denote timestep, channel, height, and width, respectively. Then the input is embedded. I is transformed into x ∈ RT×H×W×D . D represents the embedding dimension of OST. Note that OST uses Input Embedding (IE) and Position Embedding (PE) for embedding, following Spikformer. Subsequently, x is passed through the Compression module of TDCC, and is compressed from T timesteps to X0, meaning one timestep. After that, X0 is passed to the block of K encoders. Each encoder consists of three parts: Spiking Attention-free (SAF)

Y

(a) The architecture of OST

(b) Implementation details of TDCC and SLT

Figure 1: The overview of One-step Spiking Transformer (OST), which contains a Time Domain Compression and Compensation (TDCC) component and a Spiking Linear Transformation (SLT) component. Spiking matrices with Spike flags are represented in a white background. Other matrices are floating-point, represented in gray.

of SLT, Compensation of TDCC, and Spiking Gate Feedforward (SGFF) of SLT. The SAF is a linear transformer module that does not require a self-attention mechanism, while the Compensation module is to compensate for the information loss due to the compression. Moreover, the SGFF is a feedforward module, more suitable for processing spike signals. The input of each encoder is passed into SAF, Compensation, and SGFF in turn to obtain the output XK. We use skip connections before SAF and SGFF to avoid the vanishing gradient problem. Finally, XK is fed into a fully connected layer, Classification Head (CH), to produce the prediction Y . The transformation process of the sequence I is shown below:

where R1, R2, and R3 are RT×C×H×W , RT×N×D , and R1×N×D , respectively, and k = 1 · · ·K.

3.2 Time Domain Compression and Compensation (TDCC)

As shown in Figure 1(b), TDCC consists of two parts: Compression and Compensation.

Compression. It is to compress x in T timesteps toX0, e.g. one timestep, as indicated by Equation (8):

where xt is the image sequence at the t-th timestep. The rationale behind Compression is that x contains T timesteps after embedding. If x is directly passed to the encoder, the transformation information needs to be computed at each timestep, resulting in significant spatio-temporal overhead.

In fact, the human brain also involves multiple timesteps while processing dynamic visual information [Rao and Ballard, 1999]. However, because of its remarkable capacity for information integration, it efficiently extracts crucial visual features with low latency and energy consumption. In particular, temporal integration employed by visual neurons plays a vital role in the processing of visual information [Wolff et al., 2022]. Visual neurons receive inputs from the retina, as well as other regions of the visual cortex, and subsequently integrate and respond to these inputs. In this paper, we employ additive operations to create an analogy for such temporal integration, aiming to compress time-domain information. With the integration, each point in the matrix represents the intensity of the spike signals throughout all timesteps. However, compression naturally leads to information loss and subsequently performance deterioration (See Ablation Study, Section 4.3). Therefore, Compensation is introduced to regain more time-domain information.

Compensation. To accumulate time-domain information, we employ two LIF spiking neurons with different thresholds as illustrated in Equation (9):

where LIF1 and LIF2 are two LIF spiking neurons with initial threshold values of 0.5 and 1, respectively. Furthermore, inspired by dynamic thresholds in SNNs [Hao et al., 2020], the thresholds of LIF1 and LIF2 are set as learnable parameters. They update dynamically during training.

Neurons in the human brain are able to dynamically adjust their activation thresholds in response to changes in input information [Zhang and Linden, 2003]. We accumulate the spike signals at two different thresholds, to improve the diversity of time-domain information. To be more specifically, Compensation is to deal with possible information loss caused by compression. Its foundation is the rate coding of SNNs, converting information into spikes, making the number of spikes proportional to the information intensity [Kim et al., 2022]. In this way, the information loss due to compression can be compensated, thus performance improved.

3.3 Spiking Linear Transformation (SLT)

As shown in Figure 1(b), OST also has two parts: SAF, Spiking Attention Free, and SGFF, Spiking Gate FeedForward.

SAF. It is of linear complexity and is a variant of the Attention-Free (AF) mechanism [Zhai et al., 2021]. SAF's spiking operations has the benefit of low energy consumption. More specifically, the vanilla AF linearly transforms the input X0 ∈ R1×N×D into Q = X0W Q, K = X0W K , and V = X0W V , where W Q, W K and W V ∈ RD×D′ . D′ is the hidden dimension of SAF. Subsequently, it performs the operations through Equations (10) and (11):

where σq and softmax are activation function Sigmoid and Softmax respectively; ⊙ is the element-wise product.

AF uses element-wise operations, which completely gets rid of the need for dot product operations. Its complexity is O(Ld) instead of O(L2d). The dot product operation violates the original design of SNN as it involves complex floatingpoint MAC operations. AF has the potential to replace the existing spiking dot attention. However, the vanilla AF could not be directly applied to SNN, since the fire value of the activation functions (σq and softmax) are floating-point, yet the process involves numerous floating-point exponent multiplication and division operations.

In comparison, our proposed SAF has much fewer floatingpoint operations, as shown in Figure 1(b). Following the same procedure as AF, we first obtain Q, K, and V by linear transformation from an input sequence X0. After that, we convert Q and K into spiking sequences QS and KS by spiking neurons. Finally, we employ the spiking sequences QS, KS along with the floating-point sequence VF to execute the SAF, as in Equations (12) and (13):

where QS,KS ∈ R1×N×D , and LIF is the spiking neuron. Since VF is a floating-point matrix, the outcome of the element-wise product constructed in SAF is also a floatingpoint matrix. Subsequently, we feed this matrix into the Compensation module to obtain the spiking sequences output.

SGFF. It contains three linear transformations and two LIF spiking neurons, as in Equations (14) and (15):

where W G,W P ∈ RD×D′′ , and W O ∈ RD′′ ×D . D′′ is the hidden dimension of SGFF. For example, if the embedding dimension D is 128 and D′′ = 4D, for an input sequence XF, the hidden dimension of GS and PF are expanded fourfold to 512 by linear transformations W G and W P . After computing the element-wise product between GS and PF, the embedding dimension of the sequence is reversed back to 128 by the linear transformation W O , allowing the sequence to pass to the next encoder.

The reason that we propose SGFF instead of using MLP in spiking transformers is to improve the performance. SGFF is a variant of gMLP [Liu et al., 2021], with the additional incorporation of a gating mechanism. The gating mechanism enables the network to selectively learn and control the delivery of information. By controlling the open and closed states of the gate, gMLP can more effectively capture significant features and correlations in the input data, thus enhancing the representational ability of the model. However, the gate state in gMLP relies on an activation function like GELU, which involves a substantial number of floating-point computations. Furthermore, the element-wise product ⊙ in gMLP is also calculated using floating-point multiplication, which is contrary to the intention of SNN. Consequently, we replace the GELU with the spiking neuron LIF to avoid excessive floating-point computations. Additionally, the SGFF mechanism downgrades the element-wise product ⊙ to a binary logical AND (&) operation, which is certainly a more energy-efficient alternative to floating-point multiplication.

3.4 Theoretical Energy Consumption Analysis

This section analyzes the energy consumption of OST from a theoritical perspective. For the vanilla transformer, energy consumption is calculated as the product of the energy per floating-point operation (e.g., EMAC, EM) and the number of operations. In contrast, OST's energy consumption is determined by multiplying the energy per binary spike operation by the timestep, fire rate, and total number of operations. As shown in the first block (top two rows) of Table 2, all convolution (Conv) operations in the vanilla transformer are MAC operations. In comparison, only the first Conv in OST used to generate binary spikes is a MAC operation. The second block of Table 2 compares the self-attention of the vanilla transformer and AFT component of OST. The complexity of f(Q,K, V ) is linear, which is achieved by the element-wise product. Given that the elements in element-wise product are binary spikes (either 0 or 1), the operation can be executed as a mask operation, which requires no energy consumption [Yao et al., 2023a]. Lastly, the third block (bottom three rows) of Table 2 addresses the MLP of the vanilla transformer and the SGFF component of OST. Unlike the vanilla transformer, OST incorporates an additional linear layer (Linear3) to implement a binary gating mechanism. A more detailed analysis can be found in Supplementary S1.

4 Experiments

OST is validated on both static image classification, involving ImageNet, CIFAR10, and CIFAR100, and neuromorphic image classification, using CIFAR10-DVS [Li et al., 2017] and DVS128 Gesture [Amir et al., 2017]. These experiments are detailed in Section 4.1 and Section 4.2 respectively, while the ablation study is presented in Section 4.3.

4.1 Static Image Classification

ImageNet. Following Spikformer, OST training utilizes 224 × 224 images from ImageNet and Adam [Kingma and Ba, 2014]. The learning rate is initially set to 6e-5 and progressively reduced using a cosine decay. The batch size and the epochs are 16 and 310 respectively. Table 1 presents the results, where the methods for comparison are all the SOTA convolutional SNNs (top block), vanilla Transformer (middle row), and spiking Transformers (bottom block, where OST is). OST achieved a SOTA accuracy of 74.97%. In particular, OST narrows the accuracy gap compared to the vanilla Transformer while achieving a 5.5× reduction in energy consumption. OST reduces energy consumption by 46.35% and improves accuracy by 5.91% compared to SEW-ResNet152. Furthermore, OST outperforms Spikformer by 1.56% in terms of accuracy despite having a comparable size. It also requires 4.94G fewer operations and consumes 4.66mJ less power when predicting a single image. Furthermore, we analyzed the average sparsity of different matrices of OST-8-512, as shown in Table 3, the output of SGFF exhibits high sparsity, indicating that the spike gating mechanism effectively selects critical information (Details in Supplementary S2).

CIFAR 10 & CIFAR 100. Also following Spikformer, the input image size here is 32×32 with a batch size of 128. The results are presented in Table 4. OST achieved the best performance on both CIFAR10 and CIFAR100 using one timestep. In particular, OST-4-256 and OST-2-384 achieved accuracy improvements of 1.19% and 0.67%, respectively, compared to Spikformer-4-256 and Spikeformer-2-384. In addition, OST-4-384 achieves the accuracy of 95.64% on CIFAR10, which is 1.1% and 0.35% higher than TET and Spikformer4-384, respectively. On CIFAR100, OST-4-256 and OST-2384 gain 1.09% and 0.52% compared to Spikformer-4-256 and Spikeformer-2-384, respectively, while OST-4-384 outperforms TET and Spikformer-4-384 by 4.31% and 0.95%, respectively. Compared to the vanilla Transformer, OST further narrows the accuracy gap to 1.19% and 2.26% on CIFAR10 and CIFAR100, respectively. Such a gap is supposedly due to the information loss during the binarization of spike communication in OST, while Vanilla Transformer uses floating point communication, and its energy consumption is 4.25mJ, which is more than 9× higher than that of OST. We visualize the output of three SNNs, as in Figure 2. OST can provide more consistent coverage of the target object in comparison with Spikformer and SDT. Our results are statistically better as the P-values are all way below 0.05, reject the null hypothesis. More detailed experiments on different model scales with different parameters are in Supplementary S3.

4.2 Neuromorphic Image Classification

CIFAR10-DVS. This is an event-stream neuromorphic dataset containing 10,000 images, each of which was converted from CIFAR10 using a Dynamic Vision Sensor (DVS). The input image size is 128 × 128, with batch size 16. The learning rate is 1e-3 initially and reduces using a cosine decay. The initial timesteps here is not 4 as in Section 4.1, but 16 due to the increased task difficulty in classifying neuromorphic images. The training epoch is set to 106. The number of transformer encoder blocks N is set to 2, while the embedding dimension D is 256. We also use neuromorphic data augmentation, which is consistent as that in the literature. The results are presented in Table 5, which shows that OST, still with ONLY one timestep, outperforms others, including two SOTA SNNs, SDT (80.0%) and Spikeformer (80.9%).

DVS128 Gesture. This dataset contains 11 hand gesture categories from 29 individuals under 3 illumination conditions. The input size here is 128×128 and batch size 16. The initial learning rate is 1e-3 with a training epoch of 200. To minimize overfitting, we use two transformer encoder blocks with 256 embedding dimensions in the experiments. The results are presented in Table 6, where we achieved an accuracy of 99.0%. Compared to the TA-SNN (60 timesteps), we achieved the same accuracy but again, with ONLY one timestep. Moreover, to achieve a similar performance as OST, SDT and Spikformer need 16 timesteps instead of OST's 1 timestep. The cases of successful classification of three methods are analyzed in Figure 3. The analysis shows that OST places greater emphasis on the target, whereas SDT and Spkiformer exhibit a more dispersed pattern of spiking activations. More detailed experiments and analysis on this part of the study are presented in Supplementary S4.

4.3 Ablation Study

This ablation study of OST uses CIFAR100 to verify the effectiveness of TDCC and SLT. Note that all experiments follow the same setup as the above sections, if not specified. In addition, the impact of timesteps on spiking transformers, including SDT, Spikformer and our OST, are studied as well.

TDCC. To evaluate the impact of Compression and Compensation, we apply them on SDT (with SDT's SDSA and MLP), Spikformer (with Spikformer's SSA and MLP) and OST (with SAF and SGFF). The results are in Table 7, where '�' means the presence of the method, yet '�' is the absence. Looking at the first two rows of each method, Compression clearly reduces accuracies but lowers energy consumption as well. When Compensation is introduced (the third row of each method), accuracy bounces back with a slight increase in energy use, showing TDCC can achieve a good balance between accuracy and energy consumption.

SLT. The effectiveness study of SLT's two parts, SAF and SGFF, is presented in Table 8. Compared with the first row, where both SAF and SGFF are absent, the second row shows SAF improves in both accuracy and energy consumption. The third row shows SGFF brings even more improvement in accuracy but results in much higher energy as well. The results suggest that SAF can reduce power consumption, while SGFF can improve accuracy.

Initial Timesteps. The impact of the initial timesteps T is investigated on various tasks. In general, fewer timesteps would affect other spiking transformers quite negatively, especially on more difficult tasks like neuromorphic image classification. See details in Supplementary S5, S6.

Summary of Experiments. The experiments in the main paper and Supplementary have demonstrated that: (1) OST

performs as well as or better than SOTA methods across static and neuromorphic datasets. (2) OST runs faster, especially noticeable when timesteps T ≥ 4. The speed of OST is 2× more than SOTA method with 8 timesteps. This is particularly important for tasks that may require more timesteps when using other methods. (3) OST performs better, particularly with small timesteps (T < 4). When T = 2, we achieved 5.6% higher accuracy compared to the SOTA methods. OST exhibits a very short inference time yet has high performance. (4) OST is less sensitive to the reduction of timesteps (even with only one timestep). It is crucial for deployment on mainstream neuromorphic hardware in practice.

5 Conclusion

This paper proposes One-step Spiking Transformer (OST), which only needs one timestep in the transformer block, and is of a linear complexity. To achieve this goal, two key components, Time Domain Compression and Compensation (TDCC) and Spiking Linear Tranfomration (SLT) are introduced. The former can compress information from multiple timesteps into one, reducing the spatio-temporal overhead while maintaining performance. The latter further reduces the floating-point MAC operations of the model. Its complexity is as low as linear. The efficacy of OST is verified on both static and neuromorphic image sets. OST can achieve SOTA performance without increasing timesteps even for more difficult tasks. The ablation studies show that TDCC and SLT are indeed effective in reducing energy consumption while improving accuracy. Hence, OST is a strong candidate in SNN.

Acknowledgments

This work was supported by National Natural Science Foundation of China under Grant 62276175.