NeurIPS 2023 | 结合脉冲神经网络和Transformer的纯加法Transformer

2023-10-15 13:10

©作者 | 李国齐课题组

单位 | 中国科学院自动化研究所

论文地址：

https://arxiv.org/abs/2307.01694

代码地址：

https://github.com/BICLab/Spike-Driven-Transformer

受益于基于二进制脉冲信号的事件驱动（Spike-based event-driven，Spike-driven）计算特性，脉冲神经网络（Spiking Neural Network，SNN）提供了一种低能耗的深度学习选项 [1]。本文提出一种 Spike-driven Transformer 模型，首次将 spike-driven 计算范式融入 Transformer。整个网络中只有稀疏加法运算。具体地，所提出的 Spike-driven Transformer 具有四个独特性质：

事件驱动（Event-driven）：网络输入为 0 时，不会触发计算。
二进制脉冲通信（Binary spike communication）：所有与脉冲张量相关的矩阵乘法都可以转化为稀疏加法。
脉冲驱动自注意力（Spike-Driven Self-Attention，SDSA）算子：脉冲形式Q，K，V 矩阵之间运算为掩码（mask）和加法。
线性注意力（Linear attention）：SDSA 算子的计算复杂度与 token 和 channel 都为线性关系。

本文所提出的 SDSA 算子能耗比原始 self-attention 的能耗低 87.2 倍。所提出的Spike-driven Transformer 在 ImageNet-1K 上取得了 77.1% 的 SNN 领域内 SOTA 结果。

背景

当前 SNN 模型的任务性能较低，难以满足实际任务场景中的精度要求。如何结合 Transformer 模型的高性能和 SNN 的低能耗，是目前 SNN 域内的研究热点。

现有的 spiking Transformer 模型可以简单地被认为是异构计算模型，也就是将 SNN 中的脉冲神经元和 Transformer 模型中的一些计算单元（例如：dot-product, softmax, scale）相结合，既有乘加运算（Multiply-and-ACcumulate，MAC），也有加法运算（ACcumulate，AC）。虽然能保持较好的任务精度，但不能完全发挥出 SNN 的低能耗优势。

近期的一项工作，SpikFormer [2]，展示了在 spiking self-attention 中，softmax 操作是可以去掉的。然而，SpikFormer 中保留了 spiking self-attention 中的 scale 操作。原因在于，脉冲形式 Q，K，V 矩阵之间运算会导致输出中会包含一些数值较大的整数，为避免梯度消失，SpikFormer 保留了 scale 操作（乘法）。

另一方面，SpikFormer 采用 Spike-Element-Wise（SEW）[3] 的残差连接，也就是，在不同层的脉冲神经元输出之间建立 shortcut。这导致与权重矩阵进行乘法操作的脉冲张量实际上是多 bit 脉冲（整数）。因此，严格来说，SpikFormer是一种整数驱动 Transformer（Integer-driven Transformer），而不是脉冲驱动 Transformer。

方法

本文提出了 Spike-driven Transformer，如下图所示，以 SpikFormer [2] 中的模型为基础，做出两点关键改进：

提出一种脉冲驱动自注意力（SDSA）算子。目前 SNN 领域中仅有 Spike-driven Conv 和 spike-driven MLP两类算子。本文所提出的 Spike-driven Self-attention 算子，为 SNN 领域提供了一类新算子。
调整 shortcut。将网络中的 SEW 全部调整为 Membrane Shortcut（MS）[4,5]，也就是在不同层的脉冲神经元膜电势之间建立残差连接。

SDSA 算子。ANN 中的原始自注意力（Vanilla Self-Attention，VSA）机制的表达式为：

是 scale 因子。VSA 的计算复杂度为，为 token 个数，为 channel 个数。本文所提出的 SDSA 算子的表达式为：

其中是脉冲形式的矩阵，为哈达玛积。SDSA 算子的计算复杂度为。如下图所示。

总体来说，SDSA 算子有两个特点：

使用哈达玛积替换了之间的矩阵乘法。
逐列求和以及一个额外的脉冲神经元层代替了softmax和scale 操作。

SDSA 算子本身几乎不消耗能量。首先，由于都为二进制脉冲矩阵，它们之间的哈达玛积可以看作是 mask 操作。其次，整个公式（2）中只有带来的稀疏加法，根据统计，后的非零值比例约为 0.02，因此这部分加法的能耗几乎可以忽略不计。

MS 残差连接。目前 SNN 领域中一共有三种残差连接。一种是直接参考 ResNet 的 Vanilla Shortcut [6]，在不同层的膜电势和脉冲之间建立捷径；一种是 SEW [3]，在不同层的脉冲之间建立捷径；一种是 MS [4]，在不同层的膜电势之间建立捷径。MS 连接之后会跟随一个脉冲神经元，这可以将膜电势之和转化为 0/1，从而保证网络中所有脉冲张量与权重矩阵之间的乘法可以被转换为加法。因此，本文使用 MS 残差来保证 spike-driven。

SNN中的算子及其能耗评估

Spike-driven 的核心是，与脉冲矩阵相关的乘法运算都可以被转换为稀疏加法。当 SNN 运行在神经形态芯片上时，spike-driven 计算范式能够发挥出低能耗优势。

Spike-driven Conv 和 Spike-driven MLP。脉冲驱动计算有两层含义：事件驱动和二进制脉冲通信。前者保证了输入为 0 时，不会触发计算；后者保证了有脉冲输入时，触发的计算为加法。当前 SNN 领域中，两类典型的算子是 spike-driven Conv 和 spike-driven MLP。在进行矩阵乘法时，如果其中一个矩阵是脉冲形式，那么矩阵乘法可以通过寻址算法被转换为加法。

Spike-driven Self-Attention。SDSA 中之间的运算包含了三个脉冲矩阵之间的运算。那么就有两种可行的方式。第一，与 spike-driven Conv 和 spike-driven MLP 中一致，每两个脉冲矩阵之间的运算为乘法，然后被转换为加法，这也是 [2] 中采用的方法。第二，本文中所采用的方式，两个脉冲之间的运算为哈达玛积，因为脉冲矩阵中只有 0 或者 1，所以这种操作相当于 mask 操作，可以被当前的神经形态芯片 [7] 所支持。

在理论评估 SNN 能耗时，可以简单的认为，SNN 的能耗为：同等架构 ANN 的FLOPs × 脉冲发放率 × 单个加法操作的能耗 × 时间步。下图给出了 ANN 中的自注意力和 SNN 中的脉冲驱动自注意力部分之间的能耗对比。

结果

Spike-driven Transformer 在 ImageNet 上的结果如下所示。本文取得了 SNN 域的 SOTA 结果。

不同规模模型下，本文的 SDSA 与 ANN 中的 VSA 之间的能耗对比如下图所示（包含生成所需的能耗）。可以看出，由于 SDSA 是线性注意力，网络规模越大，SDSA 的能效优势越突出。

本文中关于 SDSA 和 MS 的消融实验。MS 会带来性能提升，SDSA 则会导致性能损失。总体来说，性能增益大于损失。

注意力图可视化如下。

全文到此结束，更多细节建议查看原文。本文所有代码和模型均已开源，欢迎关注我们的工作。

参考文献

[1] Roy, Kaushik, Akhilesh Jaiswal, and Priyadarshini Panda. "Towards spike-based machine intelligence with neuromorphic computing."Nature (2019).

[2] Zhou, Zhaokun, Yuesheng Zhu, Chao He, Yaowei Wang, Shuicheng Yan, Yonghong Tian, and Li Yuan. "Spikformer: When spiking neural network meets transformer."ICLR (2023).

[3] Fang, Wei, Zhaofei Yu, Yanqi Chen, Tiejun Huang, Timothée Masquelier, and Yonghong Tian. "Deep residual learning in spiking neural networks."NeurIPS (2021).

[4] Hu, Yifan, Lei Deng, Yujie Wu, Man Yao, and Guoqi Li. "Advancing Spiking Neural Networks towards Deep Residual Learning."arXiv preprint arXiv:2112.08954 (2021).

[5] Yao, Man, Guangshe Zhao, Hengyu Zhang, Yifan Hu, Lei Deng, Yonghong Tian, Bo Xu, and Guoqi Li. "Attention spiking neural networks."IEEE T-PAMI (2023).

[6] Zheng, Hanle, Yujie Wu, Lei Deng, Yifan Hu, and Guoqi Li. "Going deeper with directly-trained larger spiking neural networks." AAAI (2021).

[7] Pei, Jing, Lei Deng, Sen Song, Mingguo Zhao, Youhui Zhang, Shuang Wu, Guanrui Wang, Guoqi Li et al. "Towards artificial general intelligence with hybrid Tianjic chip architecture."Nature (2019).

更多阅读