Tensor Decomposition Based Attention Module for Spiking Neural Networks

Abstract

[SUMMARY] The attention mechanism has been proven to be an effective way to improve the performance of spiking neural networks (SNNs). However, from the perspective of tensor decomposition to examine the existing attention modules, we find that the rank of the attention maps generated by previous methods is fixed at 1, lacking the flexibility to adjust for specific tasks. To tackle this problem, we propose an attention module, namely Projected-full Attention (PFA), where the rank of the generated attention maps can be determined based on the characteristics of different tasks. Additionally, the parameter count of PFA grows linearly with the data scale. PFA is composed of the linear projection of spike tensor (LPST) module and attention map composing (AMC) module. In LPST, we start by compressing the original spike tensor into three projected tensors with learnable parameters for each dimension. Then, in AMC, we exploit the inverse procedure of the tensor decomposition process to combine the three tensors into the attention map using a so-called connecting factor. To validate the effectiveness of the proposed PFA module, we integrate it into the widely used VGG and ResNet architectures for classification tasks. Our method achieves state-of-the-art performance on both static and dynamic benchmark datasets, surpassing the existing SNN models with Transformer-based and CNN-based backbones. Code for PFA is available at https://github.com/RisingEntropy/PFA .

1 Introduction

Spiking neural networks (SNNs) are currently attracting the interest of academics due to their lower energy consumption and greater bio-interpretability compared to traditional artificial neural networks (ANNs)  [19; 36]. The brain-inspired computation model gives it significant potential in processing temporal data. Though training a deep SNN is still a challenge, recent developments in SNNs introduce backpropagation [3; 55; 65; 50; 53; 41] to relieve this issue to some degree. Moreover, it makes the extending of ANN modules into SNNs possible, such as batch normalization and residual blocks.

As it is now feasible to incorporate ANN modules into SNNs while maintaining the inherent power efficiency, we can leverage these modules to augment the overall performance of SNNs. A good practice for this is the introduction of residual architecture to SNNs. The well-known residual block proposed by He et al. [14] makes it possible to train an ultra-deep ANN network. To introduce the residual block into the field of SNN, many variations [17; 11; 18] are presented from different perspectives, solving problems like gradient vanishing. Moreover, strategies in ANN, such as normalization [67] and architecture, could also effectively favor SNNs’ performance, thereby motivating us to explore more useful techniques in ANNs for better-conducting SNNs.

Figure 1: Accuracy on CIFAR10 (left) and CIFAR100 (right). Compared with other methods, PFA significantly improves network performance.

Apart from improving network architectures, numerous other approaches can improve the performance of a network. As one of the most promising techniques in ANNs, the attention mechanism can mimic the human ability to focus on certain things while ignoring others. Three representative works [62; 70; 4] have recently proven that the attention mechanism can successfully work in SNNs. Especially, Yao et al. [62] switch channel attention to the temporal dimension, revealing that paying attention to the temporal dimension is effective for performance enhancement. Zhu et al. [70] apply two 1D convolutions to a generated 2D tensor to achieve temporal-channel attention. Inspired by experimental observations on predictive attentional remap**, Cai et al. [4] design the SCTFA module to assess the input tensor from temporal, channel, and spatial dimensions, achieving attention among these three dimensions.

Current software frameworks, i.e., PyTorch and Tensorflow, are boosted by GPU using parallel computation. This makes it hard to realize event-driven SNNs as events occur asynchronously. Based on these frameworks, as a trade-off, a popular approach is splitting the input event streams into slices and composing tensors for convenient later computation. This approach is adopted by all attention-mechanism-related works, showing a promising outcome. Since we are essentially dealing with tensor-formatted data flow, it is natural for us to consider involving tensor analysis theories to implement an attention module for SNNs, which is ignored by previous work. By utilizing mathematical tools such as tensor decomposition theory, we can examine attention modules from a more abstract perspective, thereby identifying the shortcomings and limitations of previous methods.

By expressing the composition of attention maps in previous methods through CANDECOMP/PARAFAC (CP) decomposition form, we found that they are essentially ‘rank-1’ methods. When applied to various tasks, they lack a certain degree of flexibility. In this paper, we introduce a module called Projected-Full Attention (PFA), which is capable of generating attention maps with a rank other than 1. Specifically, PFA consists of two components: Linear Projection of Spike Tensor (LPST) and Attention Map Composing (AMC). LPST is responsible for generating $R$ sets of projections from the input tensor using a small number of parameters. AMC is responsible for creating attention maps using these $R$ sets of projections, where $R$ is called The connecting factor representing the rank concept of the CP decomposition and is a hyper-parameter that can be adjusted based on the specific task.

The key contributions of this paper are outlined as follows:

• We propose PFA, a module that can achieve temporal-channel-spatial attention. In contrast to previous approaches, the ranks of the attention maps generated by PFA are not fixed at 1; instead, the rank can be chosen based on the specific task. Furthermore, the parameter count of PFA grows linearly with the data scale, and the computational cost is equivalent to that of a single standard convolution layer.

• We present a comprehensive theoretical analysis, focusing on two critical aspects: the rank of tensors and $R$, the connecting factor. These analyses lead to a tailored selection criterion, providing effective insights into experimental outcomes and offering a valuable guidance for future applications.

• We conduct thorough experiments on both static and dynamic datasets. The results demonstrate the effectiveness of PFA which achieves state-of-the-art (SOTA) accuracy on both dynamic and static datasets. Ablation studies further validate the structural reasonableness of PFA. Additionally, we visualize the attention map generated by PFA to intuitively illustrate attention distribution across temporal, channel, and spatial dimensions.

2 Related Works

Spiking Neural Network: Spiking Neural Networks (SNNs), the third generation of neural networks [19; 36], offer a closer emulation of the human brain’s efficiency by utilizing discrete spikes for information transmission [64]. This unique method of communication allows for enhanced energy-saving feature compared to traditional Artificial Neural Networks (ANNs), which rely on continuous signals. The temporal representation capabilities of SNNs have been a focal point in research [62; 8; 64], drawing from seminal neurodynamic models such as the Hodgkin-Huxley (H-H) model [15] and further developed in works like Izhikevich [19] and Leaky Integrate and Fire (LIF) [31]. These models underscore SNNs’ potential in capturing time-dependent patterns in data, a feature not inherently present in ANNs. The challenges in training SNNs, due to their non-differentiable spiking behavior, have also been addressed in recent literature, with gradient-based optimization methods [54; 56; 51] and ANN-to-SNN conversion techniques [45; 9; 49; 52] being key developments. In addition to research on training methods, there are many other studies dedicated to improving the performance of SNNs, such as attention mechanisms.

Attention Mechanism in SNNs: Since Google proposed recurrent models of visual attention in 2014 [38], the attention mechanism become a potent tool to increase the performance of a neural network, which gives distinct parts of the input data varied weights. For instance, the squeeze-and-excitation network (SENet) [16] gives different weights to different channels of input data and significantly improves the performance. From this fact, it is of vital significance to introduce the attention mechanism into SNNs. Qiu et al. apply the attention mechanism to the encoding layers of SNNs [42]. To achieve temporal-wise attention, inspired by SENet, TA-SNN [62] is proposed. Besides, TCJA [70] uses two 1D convolutions to pay attention to both channel dimension and temporal dimension to prolong the attention mechanism. Nevertheless, they fail to consider three crucial aspects of SNNs, i.e. temporal, channel, and spatial. By changing SENet and leaving affection in the temporal direction, Cai et al. [4] and Yao et al. [63] successfully apply temporal-channel-spatial attention to SNNs. While the methods of predecessors have been very successful, they do not take into account the fact that the input data is a high-order tensor. Examining data from a tensor perspective can provide a more mathematically abstract viewpoint, revealing additional properties. As described in Section 3.1, the methods employed by previous researchers result in ‘rank-1’ attention maps, lacking the flexibility and specificity for different tasks.

Tensor Decomposition: Tensor decomposition methods [20; 46; 25] have undergone years of development and achieved remarkable success. Many prior studies [28; 61; 39; 26; 2; 35; 48] have embraced these techniques to optimize neural network architectures or uncover versatile modules with multiple functionalities. Kossaifi et al., for instance, conducted a comprehensive review of convolution kernels from the perspective of tensor decomposition and introduced an innovative convolution module that can seamlessly extend to higher dimensions. Lau et al. propose a new decomposition method to enable the direct use of the depth-wise convolutional layer with large kernels in the attention module, without requiring any extra blocks [32]. Apart from using tensor decomposition methods to improve network structures, tensor decomposition theory can also impose certain constraints on the data itself [58; 60; 5]. Xu et al. leverage tensor decomposition to predict real-time traffic flow, demonstrating its effectiveness in handling complex patterns [60]. Chen et al. employ tensor CP decomposition to generate an attention map for segmentation, achieving significant improvements [5]. These methods either use tensor decomposition to reduce the parameter count in processing high-dimensional data or leverage low-rank properties to enhance performance. In this paper, drawing inspiration from these prior endeavors, we leverage tensor decomposition methods to enhance the field of SNNs. We design a lightweight attention module, designed to be both parameter-efficient and computationally efficient.

3 Motivation and Method

3.1 Motivation

Tensor CP decomposition [24], is a mathematical technique employed to decompose a high-order tensor into a sum of rank-one tensors. Mathematically, for a tensor denoted as $\mathcal{X}\in\mathbb{R}^{D_{1},D_{2},\dots,D_{n}}$, the CP decomposition can be expressed as:

$$\mathcal{X}\approx\sum_{r=1}^{R}U_{1,r}^{(D_{1})}\circ U_{2,r}^{(D_{2})}\circ\dots\circ U_{n,r}^{(D_{n})}.$$ (1)

In this equation, $U_{1,r}^{(D_{1})}\circ U_{2,r}^{(D_{2})}\circ\dots\circ U_{n,r}^{(D_{n})}$ represents the $r$-th rank-one tensor, where the symbol $\circ$ signifies the outer product of vectors, and $R$ denotes the rank. Specifically, $U_{n,r}^{(D_{1})}$ represents a vector corresponding to the $n$-th dimension in the $r$-th rank-one tensor, with a dimensionality of $D_{1}$.

Rethinking previous attention modules [62; 70; 4] from the perspective of CP decomposition, we can find that they are all rank-one methods. For the sake of clarity, we denote an all-one vector with a dimensionality of $D$ as $I^{(D)}$. Specifically, for TA-SNN [62], the temporal attention module for SNNs yields the attention map $\mathcal{A}_{\text{TA-SNN}}$ for the input tensor $\mathcal{X}$ as follows:

$$\mathcal{A}_{\text{TA-SNN}}=I^{(HW)}_{S}\circ I^{(C)}\circ\mathcal{F}_{T}^{(T)}(\mathcal{X}).$$ (2)

Here, $\mathcal{F}_{T}^{(T)}(\mathcal{X})$ represents a function that generates the temporal attention vector with a dimensionality of $T$. The same mathematical form can also be used to reconfigure the TCJA [70] and SCTFA [4] method. For TCJA, the attention map $\mathcal{A}_{\text{TCJA}}$ can be written as:

$$\mathcal{A}_{\text{TCJA}}=I^{(HW)}_{S}\circ\mathcal{F}^{(C)}_{C}(\mathcal{X})\circ\mathcal{F}^{(T)}_{T}(\mathcal{X}).$$ (3)

For SCTFA, the attention map $\mathcal{A}_{\text{SCTFA}}$ is:

$$\mathcal{A}_{\text{SCTFA}}=\mathcal{F}^{(HW)}_{S}(\mathcal{X})\circ\mathcal{F}^{(C)}_{C}(\mathcal{X})\circ\mathcal{F}^{(T)}_{T}(\mathcal{X}).$$ (4)

By comparing Equation 1 and Equation 2, 3, 4, we can easily see that the $R$ of Equation 2, 3, 4 is fixed at 1, making them special cases of Equation 1. This motivates us to design an attention module that the rank of the attention map is not fixed at 1, allowing us to control the attention effect by adjusting the rank of the attention map; namely, we wish our proposed module could be written in a form like:

$$\mathcal{A}_{PFA}=\sum_{r=1}^{R}U^{(HW)}_{s,r}\circ U^{(C)}_{c,r}\circ U^{(T)}_{t,r}.$$ (5)

Through this way, we can use $R$ to control certain properties of the attention map, such as low rank and the degree of compression of the original tensor. From the perspective of tensor decomposition, previous methods can be seen as special cases of our method when $R=1$. However, due to the absence of summation terms in previous methods, adjustments tailored to the specificity of the dataset cannot be made, lacking a certain degree of flexibility. The flexibility mentioned is crucial, as demonstrated by the experiments presented in Section 4.4. The results indicate that the optimal choice for $R$ varies across different datasets.

3.2 Projected Full-Attention (PFA)

Figure 2: A simple comparison curve of parameter quantity growth among TCJA [70], TA-SNN [62], and PFA. The parameter scale of PFA increases linearly. In this figure, the channel number is fixed at 128.

Figure 3: The detailed workflow of PFA. The input tensor is first sent to Linear Projection of Spike Tensor (LPST) module to generate three projections and split the three projections into corresponding vectors. In Attention Map Composing (AMC) module, these vectors are composed into the final attention map through the inverse process of CP decomposition. The attention is fused with the input tensor to obtain the refined tensor by Hadamard product.

In this section, we will present an overview of our proposed PFA module and detailed descriptions of two sub-modules of PFA. PFA is composed of two sub-modules, Linear Projection of Spike Tensor (LPST) and Attention Map Composing (AMC). LPST, as its name suggests, produces projections from input tensor $\mathcal{X}$ for AMC in the form of matrices. The AMC module splits the matrices from LPST to vectors to compose the final attention map.

Figure 4: Comparison of accuracy between the vanilla VGG network and VGG network with PFA modules on training and validation sets. While the training set accuracy remains consistently high for both models, the validation set accuracy shows a significant improvement with the addition of PFA modules, suggesting its efficacy in mitigating overfitting.

The core of the attention mechanism is to give essential parts of the input data higher weights while ignoring less-related parts to some degree. This is particularly important for SNNs, especially for data collected by DVS cameras. Data from DVS camera often contains more noise, which greatly affects network performance. Thanks to the feature of CP-decomposition, PFA can effectively suppress noise. Also, as shown in Figure 4, the inexact representation of the input tensor can help reduce overfitting, which is very severe in SNNs. In our proposed PFA, we first design the LPST module to project the higher-order tensor (generally 3D) to three lower-order tensors, i.e., 2D matrices. Using learnable parameters, LPST can selectively reserve essential information while ignoring irrelevant ones. In AMC, we split the matrices into $R$ (called connecting factor here) vectors. Three vectors from three matrices of temporal, channel, and spatial dimensions respectively are used to conduct the outer product to compose a part of the final attention map. The final attention map is the sum of these $R$ parts. We fuse this attention map with the origin input to obtain the refined tensor.

The workflow of PFA is plotted in Figure 3. In what follows, we will give the details of LPST and AMC.

3.2.1 Linear Projection of Spike Tensor (LPST)

Tensor decomposition is a vital approach to represent a tensor and its latent features. Therefore, it is natural for us to utilize tensor-decomposition-related techniques to explore the data characteristics. A typical technique is the CP decomposition. Whatever way we adopt, it is essential to generate a ‘projection’ to extract significant components of the input tensor to yield an attention map later.

In this part, we need to construct three linear projections of the input data stream along spatial, channel, and temporal dimensions. In LPST, we apply two fully connected (FC) layers and a convolution layer to project the input tensor for simplicity. Note that here we simply use single-layer FCs and a convolution layer. What we do here is just a simple projection with learnable parameters rather than complex function fitting. In this way, we retain the required information and reduce the number of parameters.

For the temporal dimension, we first squeeze the input tensor ($H\times W\times C\times T$) in the following way:

$$Y^{T}_{(c,t)}=\frac{1}{HW}\sum\limits_{i=1}^{H}\sum\limits_{j=1}^{W}\mathcal{X}_{(c,t,i,j)},$$ (6)

where $Y^{T}_{(c,t)}$ denotes the output matrix, $H$ and $W$ denote the spatial sizes, $\mathcal{X}_{(i,j,c,t)}$ denotes the input tensor. After the squeeze operation, the tensor is converted to a matrix in the shape of $C\times T$. Then, we apply an FC layer to reserve $R$ necessary information in each column. A sigmoid activation function is applied after the FC layer. Finally, we obtain a projection matrix in the shape of $r\times T$.

For the channel dimension, the operation is similar to that of the temporal dimension. We squeeze the input tensor by:

$$Y^{C}_{(c,t)}=\frac{1}{HW}\sum\limits_{i=1}^{H}\sum\limits_{j=1}^{W}\mathcal{X}_{(c,t,i,j)},$$ (7)

where $Y^{C}_{(c,t)}$ denotes the output matrix, $C$ denotes the number of channel. Then, we pass the $T\times C$ matrix through an FC layer to acquire the projected matrix in the shape of $r\times T$ with the sigmoid activation function applied.

For spatial dimension, the channel is squeezed by:

$$Y^{S}_{(t,x,y)}=\frac{1}{C}\sum\limits_{k=1}^{C}\mathcal{X}_{(k,t,x,y)}.$$ (8)

Afterward, we employ a convolution to the squeezed tensor, treating $T$ as the input channel and $R$ as the output channel. Then, we reshape the output tensor from the convolution into a matrix in the shape of $HW\times r$. The sigmoid activation function is applied.

3.2.2 Attention Map Composing (AMC)

Once we obtain the three projections corresponding to three dimensions, a necessary step is to compose an attention map using the three projections. Here, we adopt the reverse process of the tensor CP decomposition. In the CP decomposition, a tensor $\mathcal{X}\in\mathcal{R}^{HW\times C\times T}$ can be represented as:

$$\mathcal{A}_{s,c,t}=\sum_{r=1}^{R}U^{(HW)}_{s,r}\circ U^{(C)}_{c,r}\circ U^{(T)}_{t,r},$$ (9)

where $U_{si}$, $U_{ci}$ and $U_{ti}$ denote the $i$-th vectors in the spatial, channel, and temporal projection tensors respectively, $\mathcal{A}$ is the attention map, and $\circ$ is the outer product of vectors. In particular, $R$, the connecting factor, can be viewed as the rank concept in the tensor CP decomposition (this process is plotted in Figure 3). It is worth noting that the tensor CP decomposition is just an approximate representation of tensors, which perfectly meets our demands. $R$ can control the accuracy of our representation of the input tensor. That is, the larger $R$ is, the more accurate the representation is, and vice versa, which perfectly realizes the attention mechanism we need. In the experiment part, we discuss the choice of connecting factor $R$ and its effect in detail.

3.3 Parameter and Computational Cost Analysis

This section presents an analysis of the parameter and computational cost of PFA. PFA has an advantage in terms of parameter quantity. For an input tensor with the shape of $H\times W\times C\times T$, the overall parameter amount is:

$$\underbrace{C\times R}_{\text{MLP of }Y^{T}}+\underbrace{T\times R}_{\text{MLP of }Y^{C}}+\underbrace{k^{2}\times T\times R}_{\text{Conv of }Y^{S}},$$ (10)

where $k$ is the convolution kernel size. Note that $R$ is limited to a small fixed value according to our theoretical analysis in Section 3.4, and the parameter size grows linearly with $T$ and $C$. Figure 2 shows the parameter growth curve of PFA and compares several other techniques.

PFA is also lightweight in computational burden. The time complexity of PFA can be calculated in the following way:

$$\underbrace{3HWTC}_{\text{Obtain }Y^{T},Y^{C},Y^{S}}+\underbrace{2TCR}_{\text{MLP of }Y^{T},Y^{C}}+\underbrace{HWk^{2}TR}_{\text{Conv of }Y^{S}}+\underbrace{RHWTC}_{\text{AMC}}.$$ (11)

It is worth noting that a standard convolution operation typically entails a computational burden of $HWk^{2}TC_{in}C_{out}$. In this context, given that $R$ is not large, a PFA module incurs a computational load that is smaller than that of a convolution operation since $C$ is usually in the range of several hundred.

3.4 Theoretical Analysis on $R$

Figure 5: A toy example of the effect of different ranks on the approximation outcome. error is measured under $\ell_{2}$ norm. The norm starts to cover when the rank exceeds 30.

In light of the fact that CP decomposition provides only an approximate representation of a tensor, and the quality of this approximation relies on the chosen factor, denoted as $R$, it is essential to establish a clear guiding principle for determining the optimal value of $R$. This section delves into our approach for gaining insights into the characteristics of PFA and streamlining its future applications.

Our exploration begins by adopting the methodology put forth by Acar et al. [1] to investigate the rank of the input data. However, due to the recognized NP-hard nature of computing the CP decomposition of a tensor [13; 24], a direct calculation from the tensor becomes impractical. To surmount this challenge, we employ a gradient descent technique to approximate the original tensor by systematically experimenting with different rank values, ultimately selecting the one that yields the best approximation for the tensor. To provide further clarity, consider an element of the tensor represented as:

$$x_{ijk}\approx\sum\limits_{r=1}^{R}a_{i,r}b_{j,r}c_{k,r}.$$ (12)

We treat this as an optimization problem, with $a_{i,r}$, $b_{j,r}$, and $c_{k,r}$ considered as parameters. We then optimize the loss function defined as:

$$\mathcal{L}=\frac{1}{2}\sum_{(i,j,k)\in\Omega}(x_{i,j,k}-\sum\limits_{r=1}^{R}a_{i,r}b_{j,r}c_{k,r})^{2}.$$ (13)

To perform optimization, we compute the first-order partial derivative of the loss with respect to a parameter (e.g., $a_{i,r}$), given by:

$$\frac{\partial\mathcal{L}}{\partial a_{i,r}}=-\frac{1}{2}\cdot 2\cdot(x_{i,j,k}-\sum\limits_{r=1}^{R}a_{i,r}b_{j,r}c_{k,r})\cdot b_{j,r}c_{k,r}.$$ (14)

Finally, we apply the gradient descent method for optimization:

$$a_{i,r}^{\prime}\leftarrow a_{i,r}-\mu\cdot\frac{\partial\mathcal{L}}{\partial a_{i,r}},$$ (15)

Here, $\mu$ represents the step size used in the gradient descent process. This approach allows us to effectively estimate the rank of the tensor while circumventing the computational complexities associated with CP decomposition.

We take a piece of data from the dataset CIFAR10DVS as a toy example and discover its rank. In this case, we adopt $\mu=0.0001$ and iterate the process 1000 times to get the final optimized parameters. We utilize $\ell_{2}$ norm as the metric to judge the approximation effect. As shown in Figure 5, the error reduces as the rank increases, and starts to converge when the rank exceeds 30. This toy example shows that the rank of the input tensor is around 30, indicating that we need a rank exceeding 30 to represent it using CP decomposition precisely. It should be noted that accurately representing the original tensor for the attention mechanism is not conducive to improving network performance. The core idea of the attention mechanism is to focus on the important parts while ignoring the secondaries. Representing raw tensors too precisely is counterproductive. On the one hand, overly precise characterization of input tensors is detrimental to neglecting unimportant parts. On the other hand, it may lead to severe overfitting. As shown in Figure 4, the vanilla VGG network exhibits severe overfitting, achieving nearly 100% accuracy on the training set but less than 80% on the validation set. However, upon incorporating the PFA modules, there is not a substantial improvement in training set accuracy. Nevertheless, there is a significant increase in validation set accuracy. This observation underscores the benefit of utilizing less precise tensor representations with PFA to mitigate overfitting. Therefore, a good choice for $R$ should not be too large. This is a very useful conclusion. On the one hand, it tells us the range of finding $R$. On the other hand, it also ensures that PFA can obtain good results without excessive calculation.

As shown in our experiment (Figure 8), though the rank of the attention map differs on dynamic datasets like CIFAR10DVS and static datasets like CIFAR100, the best option for $R$ is not very big. Further observing the experimental results, we find that when the value of $R$ exceeds a certain threshold (approximately $\frac{T}{2}$, $T$ is the time step), the larger $R$ is, the worse the effect will be. The above analysis of the experimental results proves the correctness of our conclusion. Considering the above analysis, we propose a principle for selecting $R$: for dynamic datasets, the search for the optimal value of $R$ is performed within the range of no more than $T$. Specifically, we focus on identifying the best $R$ value around $\frac{T}{2}$, considering the varying content across different time steps. In the case of static datasets, the search for the optimal $R$ value begins from 1 and extends until the best value is determined. Since static datasets involve duplicated frames, this search aims to find the most suitable $R$ value for the low-rank condition.

4 Experiment

We evaluate the classification performance of our PFA on two static datasets: CIFAR10 [30], CIFAR100 [30] and two dynamic datasets: CIFAR10DVS [33], NCaltech-101 [40]. Details of the datasets, network architecture, data augmentation, loss function, and pre-process procedure is introduced in this section. Extensive experiment results are also presented.

Dataset	Learning Rate	Epoch	Batch Size
CIFAR10	0.1	200	128
CIFAR100	0.1	200	64
CIFAR10DVS	0.0001	200	32
NCaltech-101	0.0001	200	32

Table 1: Hyper-parameter settings of PFA.

4.1 Datasets and Training Details

4.1.1 Datasets

We have conducted experiments on both static and dynamic datasets for object classification. We train and validate PFA on a workstation equipped with one RTX 3090. For different datasets, the hyper-parameters are listed in Table 1. The summaries of datasets and augmentation involved in the experiment are listed below.

CIFAR 10/100 consist of 50k training images and 10k testing images with the size of $32\times 32$ [30]. We use ResNet-19 for both CIFAR10 and CIFAR100. Random horizontal flips and crops are applied to the training images for augmentation. Moreover, cutout is also used for augmentation which is the same as [8]
CIFAR10-DVS converts 10,000 frame-based images of 1010 classes into event streams with the dynamic vision sensor. Since the CIFAR10DVS dataset [33] does not divide training and testing sets, we split the dataset into 9k training images and 1k test images and reduced the spatial resolution from $128\times 128$ to $48\times 48$ [8]. We use VGGSNN for neuromorphic datasets CIFAR10DVS. In each frame, horizontal flip** and mixup are adopted, where the probability of Flip** is set to 0.5. Then, we randomly select one augmentation among rolling, rotation, cutout, and shear, where the random rolling range is 5 pixels, and the degree of Rotation is sampled from the uniform distribution, which is the same as the Ref [8].
NCALTECH-101 is also converted from the original version of Caltech-101 [40] with a slight change in object classes to avoid confusion. The NCaltech-101 consists of 100 object classes plus one background class. We apply the 9: 1 train-validation split as CIFAR10DVS. The VGGSNN is applied to neuromorphic datasets NCALTECH-101. Moreover, the augmentation is identical to CIFAR10DVS. Fashion-MNIST is a more demanding successor of the famous MNIST dataset. Fashion-MNIST contains 70,000 static grayscale images of 10 different categories of clothing, each image being of size 28x28 pixels. We utilize the simple example network from the Spikingjelly framework, which comprises two convolution layers and a fully-connected layer, to evaluate PFA on this dataset.

4.1.2 Loss function

We used the cross-entropy loss for the neuromorphic datasets CIFAR10DVS and NCALTECH-101. In order to get a better performance in static datasets (CIFAR 10 / 100), we used a Temporal Efficient Training [8] loss function. It can be concluded as follows:

$$\mathcal{L}_{\mathrm{TET}}=\frac{1-\lambda}{T}\cdot\sum_{t=1}^{T}\mathcal{L}_{\mathrm{CE}}[\boldsymbol{O}(t),\boldsymbol{y}]+\frac{\lambda}{T}\cdot\sum_{t=1}^{T}\mathcal{L}_{\mathrm{MSE}}(\mathbf{O}(t),\phi),$$ (16)

where $T$ is the total time steps, $\mathcal{L}_{\mathrm{CE}}$ and $\mathcal{L}_{\mathrm{MSE}}$ denote the cross-entropy loss and mean square error. And $\boldsymbol{y}$ represents the target label. Moreover, $\mathbf{O}(t)$ and $\phi$ are the output and a constant used to regularize the membrane potential distribution. Additionally, we set $\phi=V_{th}$ in our experiments. In practice, we use a hyper-parameter $\lambda$ to adjust the proportion of the regular term.

Figure 6: Schematic diagram of the network structure. On the left is the convolution layer of the VGG network, and on the right is the convolution layer of Resnet.

4.1.3 Network Architectures

To evaluate PFA, we integrate it into VGG-type [47] and ResNet [14] architectures, which are previously utilized in [12; 70; 8; 63]. Figure 6 presents the basic blocks involved in VGG and ResNet architectures. Five PFA modules are inserted into a VGG-type network after each average pooling operation. In the ResNet architecture, PFA is inserted before the residual connection in the last two layers. When splitting the event flow data into frames, we simply integrate the input in the time domain. Particularly, our work is based on the software framework SpikingJelly [10]. Taking notice that frames in CIFAR10-DVS and NCaltech-101 are entirely occupied by their content, we utilized a pooling layer as the first layer of the network to reduce memory consumption.

4.2 Comparison with Existing SOTA Works

Dataset Method Type Architecture Time steps Accuracy CIFAR10/CIFAR100 SIDP  [44]^ICLR-2019 Hybrid training ResNet-20/VGGSNN 250/125 95.4/76.4 Diet [43]^TNNLS-2021 Spike-based BP ResNet-20 10/5 92.5/64.0 tdBN [68]^AAAI-2021 Spike-based BP ResNet-19 6 93.1/71.1 Dspike [34]^NeurIPS-2021 Spike-based BP ResNet-18 6 94.2/74.2 NAS-SNN [23]^ECCV-2022 Spike-based BP NAS 5 92.7/73.0 TET [8]^ICLR-2022 Spike-based BP ResNet-19 6 94.5/74.7 DSR [37]^CVPR-2022 Spike-based BP ResNet-19 20 95.4/78.5 Spikformer [69]^ICLR-2023 Spike-based BP Spiking-ViT 4 95.5/78.2 PFA Spike-based BP ResNet-19 6 95.7/79.1 4 95.7/78.1 2 95.6/76.7 CIFAR10DVS LIAF-Net  [57]^TNNLS-2021 Spike-based BP LIAF-Net 10 70.4 tdBN [68]^AAAI-2021 Spike-based BP ResNet-19 10 67.8 TA-SNN  [62]^ICCV-2021 Spike-based BP LIAF-Net 10 72.0 Dspike [34]^NeurIPS-2021 Spike-based BP VGGSNN 10 75.4 TET [8]^ICLR-2022 Spike-based BP VGGSNN 10 83.2 DSR  [37]^CVPR-2022 Spike-based BP VGGSNN 10 77.3 TCJA  [70]^TNNLS-2024 Spike-based BP VGGSNN 10 80.7 Spikformer [69]^ICLR-2023 Spike-based BP Spiking-ViT 16 80.9 PFA Spike-based BP VGGSNN 14 84.0 NCaltech-101 SALT  [22]^{Neural Netw-2021} Spike-based BP VGGSNN 20 55.0 TCJA  [70]^TNNLS-2024 Spike-based BP VGGSNN 14 78.5 PFA Spike-based BP VGGSNN 14 80.5

Table 2: Compare with existing works. Our method improves network performance across all tasks.

To validate the effectiveness of our PFA, we conducted a comparative analysis with several state-of-the-art SNNs, including those employing hybrid training and spike-based backpropagation. The results for both static image data and dynamic data classification tasks are summarized in Table 2. Specifically, the results for the static image dataset are presented at time steps 2, 4, and 6.

For CIFAR10 and CIFAR100 datasets, our PFA outperforms prior works in terms of accuracy. Notably, PFA achieves superior performance on CIFAR10 and CIFAR100, with respective improvements of 1.2% and 4.4% over the TET method [8] which utilizes the same network architecture and time steps. Furthermore, PFA surpasses the DSR method [37] while requiring only 10$\times$ fewer time steps. Additionally, our attention-based SNN demonstrates better performance than the Spiking-ViT backbone [69].

Method	Time Steps	Accuracy
ST-RSBP[66]${}^{\text{NIPS-2019}}$	400	90.1
LISNN[6]${}^{\text{IJCAI-2020}}$	20	92.1
PLIF[12]${}^{\text{ICCV-2021}}$	8	94.4
TCJA[70]${}^{\text{TNNLS-2024}}$	8	94.8
PFA (Ours)	8	94.5

Table 3: Accuracy on Fashion-MNIST[59] dataset.

For the Fashion-MNIST dataset, as depicted in Table 3, while the accuracy of our proposed PFA falls slightly short of TCJA [70], it still surpasses previous methodologies utilizing longer time steps [66; 6] or incorporating learnable spike thresholds [12]. We hypothesize that the relatively diminutive image size of the Fashion-MNIST dataset curtails the efficacy of spatial attention. This conjecture aligns with the findings from our ablation experiments, indicating that spatial attention plays a pivotal role in determining PFA performance.

For the dynamic dataset CIFAR10DVS, our method outperforms previous approaches using binary spikes by 3.1%, even with fewer time steps. For the NCaltech101 dataset, we achieved an impressive 80.5% top-1 accuracy, far surpassing the performance of prior work [70].

To further illustrate the advantages of our PFA, we compare it with previous works at various time steps. As depicted in Figure 1, we compare PFA with Spike-based BP SNNs. Notably, our PFA exhibits a significant advantage at shorter simulation time steps, owing to its superior representation capability.

We observe from the perspective of simulated time steps that PFA can achieve satisfactory performance even with a smaller number of simulated time steps. For the CIFAR-10 dataset, simply utilizing 2 simulated time steps with PFA can surpass the performance of previous works, including transformer-like structures such as Spikeformer. On the CIFAR-100 dataset, compared to the top-performing Spikeformer, PFA lags by only 0.1% accuracy with 4 time steps. On DVS datasets like CIFAR-10DVS and NCaltech, employing 14 simulated time steps, PFA significantly outperforms previous methods. On CIFAR10-DVS, PFA surpasses Spikeformer by 3.1%, and compared to the TET [8] method with a similar structure, it achieves a higher accuracy by 0.8%. On the NCaltech-101 dataset, PFA outperforms the previous best TCJA method by 2%.

In summary, PFA offers substantial advantages over previous methods. When applied to attention modules like TCJA [70], PFA significantly enhances resistance to overfitting, leading to improved accuracy on the validation set. Although Spikeformer, which is based on the Spike ViT architecture, performs better than many CNN-based techniques, an especially interesting breakthrough is the incorporation of PFA into traditional CNN architectures like VGG, which can even outperform Spikeformer’s [69] performance.

4.3 Experiments on Image Generation Tasks

Figure 7: Generated images of ANN VAE, FSVAE[21], and our PFA on CIFAR10 [29] dataset.

Dataset	Method	IS$\uparrow$	FID$\downarrow$	FAD$\downarrow$
MNIST[7]	ANN[21]AAAI-2022	5.95	112.5	17.09
	FSVAE[21]AAAI-2022	6.21	97.06	35.54
	TCJA[70]TNNLS-2024	6.45	100.8	19.39
	PFA (Ours)	6.97	96.82	15.12
Fashion- MNIST[59]	ANN[21]AAAI-2022	4.58	123.7	18.08
	FSVAE[21]AAAI-2022	4.55	90.12	15.75
	TCJA[70]TNNLS-2024	5.61	93.41	12.46
	PFA (Ours)	5.35	97.3	11.97
CIFAR10[29]	ANN[21]AAAI-2022	2.59	229.6	196.9
	FSVAE[21]AAAI-2022	2.94	175.5	133.9
	TCJA[70]TNNLS-2024	3.73	170.1	100.4
	PFA (Ours)	3.84	166.4	92.83

Table 4: Comparison with existing methods on image generation tasks.

To further validate the efficacy of PFA, we conduct image generation tasks on MNIST [7], Fashion-MNIST [59], and CIFAR10 [29] using fully spike variation autoencoders [21] (FSVAE) . We integrated our proposed PFA as the first layer of FSVAE while maintaining other settings identical to the original network. The experimental results are outlined in Table 4. Notably, PFA exhibits superior performance metrics on both MNIST [7] and CIFAR10 [29] datasets. Additionally, we visually depict the generated images in Figure 7, further corroborating the enhancements facilitated by PFA.

4.4 Discussion on $R$

CIFAR10

CIFAR100

CIFAR10DVS

NCaltech101

Figure 8: Accuracy on four tested datasets with different $R$. PFA achieves best performance when $r=3,1,8,8$ correspondingly.

Figure 8 demonstrates the performance with different $R$. The outcome reveals that for dynamic datasets, the best rank option falls around $\frac{T}{2}$ while for static datasets, the best value of $R$ (also known as connecting factor) is very small.

Dynamic datasets exhibit significant variations in content across different time steps. Conversely, static datasets involve the duplication of frames, where each frame is replicated $T$ times as input to the network. Consequently, this duplication results in the attention map being subjected to a low-rank condition. Considering a comprehensive assessment of both experimental outcomes and theoretical insights, our conclusion asserts that the precision of the attention map in faithfully mirroring the input tensor is unnecessary. Such fidelity does not inherently facilitate the effective concentration on pivotal components. Furthermore, our findings unveil an intriguing phenomenon: even in scenarios where SNNs exhibit considerable sparsity, the rank of the attention map can remain markedly low. This revelation implies the presence of untapped optimization potential within the realm of SNN sparsity.

It is worth noting that on the CIFAR10-DVS dataset, when $R=1$, the performance of PFA is comparable to TCJA (see Table 2 and does not show significant advantages. This phenomenon strongly confirms our previous statement that TCJA is a "rank-1" attention, a specal case of PFA, highlighting the necessity and advantages of non-"rank-1" attention modules.

4.5 Ablation Study

Figure 9: Ablation study on PFA layers. The x-axis represents the amount of PFA inserted. The accuracy improves with the insertion of PFA layer.

Figure 10: Ablation study on attention dimensions. T/C/S denotes temporal/channel/spatial correspondingly. SC denotes spatial-channel joint attention and so on. Full means paying attention to three dimensions at the same time, i.e. a standard PFA module.

Figure 11: Attention visualization on CIFAR10DVS dataset. Temporal and channel parts are taken from the deep layer of the VGG network. The spatial layer visualizations are extracted from the shallow layer of the network due to the inherent complexity of the information contained in the deeper layers, which may impede human comprehension and hinder the interpretability of our analysis.

In this section, we perform ablation experiments from both global and local perspectives to validate the efficacy of PFA on CIFAR10DVS [29] Initially, we vary the number of added PFA layers in the network to assess the impact on classification accuracy. Subsequently, we set the three projections obtained within the PFA to constant matrices, with each element assigned a value of 1, aiming to eliminate attention to specific dimensions. This helps explore the influence of different attention dimensions on classification results. Finally, for a more comprehensive insight into the influence of PFA on the training process, we conduct ablation experiments specifically involving the insertion of PFA.

We incorporate varying numbers of PFAs into the VGG network and present the results in Figure 9. In this experiment, the value of $R$ is held constant at 8. It is evident that the network’s accuracy exhibits a noticeable improvement with more PFA layers. Remarkably, even the addition of a single layer leads to enhanced accuracy in the network.

In order to explore how attention in temporal-channel-spatial dimensions affects classification accuracy, we conduct an ablation by fixing the corresponding projection matrix obtained by LPST to an all-one tensor. As shown in Figure 10, whichever dimension is ablated will adversely affect the results. Ablating temporal and ablating channels have similar effects. This means that PFA treats temporal and channel dimensions equally. However, when we ablate the spatial dimension, the accuracy of the network drops significantly (79.60%) but still outperforms the baseline (77.90%). Whichever dimension we ablate, the performance drops. When we restrict the focus dimensions to only one, namely, spatial, temporal, or channel, we observe another decrease in the network’s performance compared to when attention is distributed across two dimensions. The retention of spatial attention significantly impacts accuracy, followed by temporal attention. This observation highlights that PFA effectively directs attention to distinct regions across multiple dimensions, aligning with the intended purpose of attention modules. Specifically, PFA exhibits attention distribution across all three dimensions, albeit with varying weights. Notably, the spatial dimension emerges as the most significant, followed by the time and channel dimensions, in terms of attention importance.

4.6 Attention Visualization

In order to visually represent the attention distribution of the PFA module, we have generated an attention map in PFA and depicted it in Figure 11. The attention maps corresponding to the temporal and channel dimensions in the figure are obtained from the deep layer of the network. This choice was made because deep layers tend to provide a more accurate reflection of the attention’s role. Conversely, the spatial attention map is obtained from the network’s input layer, where the shallow graphic features are more distinct and conspicuous, facilitating analysis and observation. Figure 11 also illustrates the distinctive attention distributions in the temporal, channel, and spatial dimensions achieved by PFA. This finding aligns with the conclusions drawn from our ablation study, further validating our earlier assertions.

5 Conclusions

In this paper, we introduce a novel attention module named PFA, which combines learning and the generation of a temporal-channel-spatial attention map through CP decomposition. Specifically, we propose the LPST to learn a projection of the input tensor that can extract essential information for AMC to compose the attention map. In the AMC, we utilize the reverse process of CP decomposition to generate the attention map. After analyzing the rank of the raw input tensor and the best option for connecting factor $R$, we propose a principle to search for the best $R$. Remarkably, our findings suggest that even a low rank can lead to outstanding performance. On classification tasks, regardless of the form of the dataset, dynamic or static, PFA achieves the SOTA performance. When we extend our evaluation of PFA to image generation tasks, it still shows its capacity to enhance performance on this task as well. Furthermore, additional ablation studies bolster the rationale behind PFA’s efficacy, while visualizing the attention distribution of PFA offers intuitive insights into its substantial impact.

In the future, we plan to extend the applicability of the PFA model to various demanding intelligent applications and deploy the model on a neuromorphic chip.