[论文解读] Tensor Decomposition Based Attention Module for Spiking Neural Networks

Haoyu Deng, Ruijie Zhu+4 more|arXiv (Cornell University)|Oct 23, 2023

Advanced Memory and Neural Computing被引用 4

一句话总结

该论文提出了一种基于张量分解的投影全注意力（Projected Full Attention, PFA）模块，用于脉冲神经网络（SNNs），实现高效、低秩、三维（时间-通道-空间）注意力，且参数呈线性增长。通过利用CP分解及其通过线性投影和注意力图组合的逆过程，PFA在静态和动态SNN基准测试中实现了最先进性能，且计算开销极低。

ABSTRACT

The attention mechanism has been proven to be an effective way to improve spiking neural network (SNN). However, based on the fact that the current SNN input data flow is split into tensors to process on GPUs, none of the previous works consider the properties of tensors to implement an attention module. This inspires us to rethink current SNN from the perspective of tensor-relevant theories. Using tensor decomposition, we design the extit{projected full attention} (PFA) module, which demonstrates excellent results with linearly growing parameters. Specifically, PFA is composed by the extit{linear projection of spike tensor} (LPST) module and extit{attention map composing} (AMC) module. In LPST, we start by compressing the original spike tensor into three projected tensors using a single property-preserving strategy 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.

研究动机与目标

• 为解决现有SNN中缺乏张量感知注意力机制的问题，这些SNN通常将事件数据视为静态张量，而未利用张量结构。
• 设计一种轻量级注意力模块，在保持低参数增长的同时，实现在时间、通道和空间维度上的全注意力。
• 利用张量分解理论——特别是CP分解及其逆过程——实现在SNN中高效生成注意力图。
• 提供一种理论基础坚实、可解释性强且可扩展的注意力机制，兼容脉冲计算与GPU加速框架。

提出的方法

• 提出脉冲张量的线性投影（LPST）模块，通过可学习的、维度特定的线性投影，将输入脉冲张量压缩为三个低维投影张量。
• 引入注意力图组合（AMC）模块，通过连接因子反转CP分解过程，重建完整的注意力图。
• 基于张量秩分析和数据集特定属性，采用一种秩选择准则来确定最优分解秩 $ r $，确保性能与效率的平衡。
• 使用Hadamard积将优化后的注意力图与原始输入张量融合，实现动态特征重加权。
• 在VGG和ResNet架构中应用PFA模块，对静态（CIFAR10/100）和动态（CIFAR10DVS）数据集进行端到端训练。
• 采用可微分框架，兼容反向传播，实现SNN的端到端注意力训练。

Figure 3: The detailed workflow of PFA. The input tensor is first sent to Linear Projection of Spike Tensor (LPST) unit to generate three projections and split the three projections into corresponding vectors. In Attention Map Composing (AMC) unit, these vectors are composed into the final attention

实验结果

研究问题

• RQ1 张量分解能否被有效用于设计SNN中轻量化、参数高效的注意力机制？
• RQ2 输入张量的秩以及连接因子 $ r $ 的选择如何影响SNN中注意力图的质量与模型性能？
• RQ3 时间、通道和空间注意力组件在SNN分类准确率中的贡献程度如何？
• RQ4 基于张量的注意力机制是否能在静态与动态基准数据集上超越现有的基于Transformer和CNN的SNN模型？

主要发现

• PFA在CIFAR10（94.85%准确率）和CIFAR100（78.65%准确率）上达到最先进性能，优于采用Transformer和CNN骨干网络的现有SNN模型。
• 消融实验表明，若从任一维度（时间、通道、空间）移除注意力机制，性能均会下降，其中空间注意力影响最大（消融后准确率降至79.60%）。
• 即使仅插入一个PFA层，网络准确率即得到提升，且随着PFA层数增加，性能持续单调上升，表明其具备良好的可扩展性与一致收益。
• 即使在稀疏SNN中，注意力图的秩仍保持较低水平，揭示了稀疏表示中尚未开发的优化潜力。
• 可视化结果表明，PFA在三个维度上均学习到有意义的注意力分布，其中空间注意力最为显著，其次是时间与通道维度。
• 基于张量秩与连接因子分析提出的秩选择准则，能有效预测最优 $ r $，并与实际性能表现高度一致。

Figure 8: Attention visualization on CIFAR10DVS dataset. Temporal and channel parts are taken from the deep layer of VGGSNN. 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 ma