ICLR 2024 - Spike-driven Transformer V2 - Meta Spiking Neural Network Architecture Inspiring the Des

Spike-driven Transformer V2: Meta Spiking Neural Network Architecture Inspiring the Design of Next-generation Neuromorphic Chips

V1 版本：Spike-driven transformer，NeurIPS 2023
论文：https://arxiv.org/abs/2404.03663
代码：https://github.com/BICLab/Spike-Driven-Transformer-V2

神经形态计算（neuromorphic computing），这利用神经形态芯片上脉冲神经网络（Spiking Neural Network，SNN），是传统人工智能的一种有前景的节能替代方案。基于卷积神经网络（CNN）的 SNNs 是当前神经形态计算的主流。相比之下，目前还没有专门为基于 Transformer 的 SNNs 设计的神经形态芯片，这些 SNNs 刚刚出现，它们的表现仅与基于 CNN 的 SNNs 相当，没有明显优势。

这项工作提出了一种通用的基于 Transformer 的 SNN 架构，称为“Meta-SpikeFormer”，其目标是：

低功耗，支持脉冲驱动范式，网络中只有稀疏的加法运算；
多功能性，处理各种视觉任务；
高性能，显示出对基于 CNN 的 SNNs 的压倒性性能优势；
元架构，为未来下一代基于 Transformer 的神经形态芯片设计提供灵感。

具体来说，本文将 NeurIPS 2023 发表的第一版工作中的脉冲驱动 Transformer 扩展为元架构，并探索了结构、脉冲驱动自注意力和跳跃连接对其性能的影响。在 ImageNet-1K 上，Meta-SpikeFormer 实现了 80.0% 的 top-1 准确率（5500 万参数），超过了当前最先进的（SOTA）SNN 基线（6600 万参数）3.7%。这是第一个可以直接训练的 SNN 骨干网络，能够同时支持分类、检测和分割，获得 SNNs 中的 SOTA 结果。

引言

SNNs 是为了成为传统机器智能的低功耗替代品。独特的脉冲驱动是实现这一宏伟概念的关键，即在神经形态芯片上运行 SNNs 时，只有部分脉冲神经元被激活，以执行稀疏的突触累积（synaptic accumulate，AC）操作。神经形态计算本质上是一种“算法—硬件共同设计”的范式（Bottom-up and top-down approaches for the design of neuromorphic processing systems: Tradeoffs and synergies between natural and artificial intelligence）。生物神经元被建模为脉冲神经元，并在算法层面上以某种方式形成 SNNs。然后，神经形态芯片在硬件层面上配备了脉冲驱动的 SNNs。基于 CNN 的 SNNs 目前是常见的脉冲驱动设计。因此，典型的神经形态芯片，如 TrueNorth、Loihi、Tianjic 等，都支持脉冲驱动的卷积和多层感知器（MLP）操作。几乎所有 CNN 时代的架构，例如 VGG、ResNet 等，都可以开发成相应的 SNN 版本。由于 ANN 中的 Transformer 在各种任务中显示出巨大潜力，在过去两年中，一些基于 Transformer 的设计在 SNNs 中出现。大多数基于 Transformer 的 SNNs 未能利用 SNNs 的低功耗优势，因为它们不是脉冲驱动的。通常，它们保留了由普通 Transformer 主导的能耗高的乘 - 累积（MAC）操作，如缩放点积（Complex dynamic neurons improved spiking transformer network for efficient automatic speech recognition）、softmax（Online transformers with spiking neurons for fast prosthetic hand control）、scale（Spikformer: When spiking neural network meets transformer）等。

本作的第一代中开发了一种脉冲驱动的自注意力操作器，首次将脉冲驱动集成到 Transformer 中。然而，尽管这种脉冲驱动 Transformer 仅通过稀疏 AC 在 ImageNet-1K 上实现了 SOTA 结果，但它尚未显示出对基于卷积的 SNNs 的明显优势。这一版本中，提出的 Meta-SpikeFormer 则在性能和多功能性方面推进了 SNN 领域。由于 Vision Transformer（ViT）显示出出 Transformer 在视觉方面可以表现出色，已经产生了许多研究。最近 MetaFormer 中总结了各种 ViT 变体，并认为通过不指定 token 混合器（自注意力）从 ViTs 中抽象出通用架构。受这项工作的启发，本文研究了 Transformer-based SNNs 的元架构设计，涉及三个方面：网络结构、跳跃连接（skip connection，shortcut）、带有完全 AC 操作的脉冲驱动自注意力（SDSA）。

所提算法

Spiking Neuron Layer

脉冲神经元层将时空信息包含到了膜电位中（membrane potential），然后将其转换为二值脉冲，以用于下一层的脉冲驱动的计算。本文使用了标准的 Leaky Integrate-and-Fire（LIF）脉冲神经层。

LIF 脉冲神经元的动力学状态

其神经动力学方程如下（t 是脉冲神经元内部的时间步），整体可以理解为神经元在不断地充电，直到达到阈值后放电，放电后重置。之后再不断地“充电 - 放电 - 重置”。

充电： $U [t] = H [t - 1] + X [t]$ ，这里的 X[t] 可以理解为脉冲神经元的输入，来自于例如卷积、MLP 和自注意力等脉冲驱动算子。而 $H [t - 1]$ 是来自于上一个时刻的时序状态。
放电： $\text{Hea}(U[t] + u_{th}])$ ，Hea 是 Heaviside 阶跃函数，只有在输入≥0 的时候输出为 1，否则为 0。所以这里的脉冲 tensor $S [t]$ 是二值，即 0 或 1。 $u_{th}$ 是放电的阈值。
重置： $V_{reset} S[t] + (\beta U[t])(1 - S[t])$ 这里通过 $S [t]$ 实现了一个选择，要么是重置电压 $V_{reset}$ ，要么是当前的经过因子 $\beta$ 衰减的膜电位 $\beta U[t]$ 。

Meta-SpikeFormer

输入，对于 SNN，由于其具有时序处理的建模需求，所以对于图像任务，输入也需要先转换为时序形式，本文使用了最简单的复制方式进行处理，即直接额外构建一个时序维度复制总的时间步，即 T 份。这种设定相当于是将时序信息的编码能力交给了模型的浅层。在 SNN 领域中，也有独立对图像进行时序编码。

基于脉冲的静态图像编码，但本文用的是直接复制T份后送入模型学习隐式编码。

首先，本文在宏观层面上将第一代脉冲驱动 Transformer 的结构与 MetaFormer 中的 CAFormer 对齐。原始的四个脉冲编码层扩展为四个基于卷积的 SNN 块。通过实验验证了早期卷积块对 SNNs 的性能和多功能性很重要。最终模型中，前两个 stage 基于 CNN-based SNN 块，后两个基于 Transformer-based SNN 块。

然后，在微观层面上设计了基于卷积和基于 Transformer 的 SNN 块。

Conv-Based SNN Block

使用了 MobileNet 中的 SepConv 作为 TokenMixer，这与 MetaFormer 中的设计一致，但是 ChannelMLP 替换为了具有 3x3 卷积的 ChannelConv 以增强归纳偏置。公式中的 SN 表示脉冲神经元层。

Transformer-based SNN Block

这里使用重参数化多分支卷积结构生成脉冲形式的查询 Q、键 K、值 V。并构建了 ChannelMLP。

Spike-Driven Self-Attention (SDSA)：与标准自注意力的差异主要在三个方面：Q、K、V 都是脉冲 tensor；Q、K、V 之间的操作不包含 Softmax、Scale；SDSA 整体的计算复杂度与序列程度呈线性关系。

MobileVit-v2		EdgeNeXt

这里的星号是元素乘法。		这里的点号是矩阵乘法。

这些都是通道注意力。在第一版工作中提出的 SDSA-1 的基础上，本文重新设计了三种不同的组件：

SDSA-2：与 MobileVit-v2 中的注意力、SENet 中的 Block 设计非常相似。
SDSA-3 $SN_s((Q_sK_S^\top)V_s)$ ：与 EdgeNeXt 中的设计非常相似。论文中提到借鉴自 Spikformer: When spiking neural network meets transformer 中的形式 $SN(Q_S K_S^\top V_S \star s)$ 。这里使用一个 $s$ 作为放缩因子归一化数值尺度，以避免梯度问题，而在 SDSA-3 中，直接将其合并到了脉冲神经元的激发阈值中，从而避免了这个乘法操作。
SDSA-4：与 SDSA-3 的差异主要在于其中的脉冲神经元形式，这里的 $s$ 设置为可学习参数。

值得注意的是，由于这里涉及到二值脉冲信号，所以不同类型操作本身对于二值脉冲信号的效率就是需要关注的了，考虑到脉冲驱动的矩阵乘法可以通过寻址算法被转化为加法运算（Bottom-up and top-down approaches for the design of neuromorphic processing systems: Tradeoffs and synergies between natural and artificial intelligence），以及**脉冲驱动的哈达玛乘积可以看做是一种掩码操作，即逻辑与运算。这成本很低。**所以四种 SDSA 陡只有稀疏加法运算。

Shortcut 的设计

SNN 中的残差学习主要包含两个方面：

可以实现恒等映射（identity mapping），从而缓解退化问题。
可以保证脉冲驱动的计算（spike-driven computing），这是 SNN 低功耗的基础。

SNN 中现有三种 shortcut 设计：

Vanilla Shortcut（VS）：来自 Going deeper with directly-trained larger spiking neural network。形式类似于 ResNet。这可以实现脉冲驱动，但是无法实现恒等映射。
Spike-Element-Wise（SEW）：来自 Deep residual learning in spiking neural network，这可以实现恒等映射，但是脉冲的加法会生成整数。因此 SEW-SNN 是一个“整数驱动”而非脉冲驱动的 SNN。
Membrane Shortcut（MS）：来自 Advancing spiking neural networks toward deep residual learning，直接连接两个神经元的膜电位，并且可以实现脉冲驱动的恒等的映射。

实验细节

图像分类和目标检测
- We set the timestep T=1 for 200 epoch training to reduce training cost, then finetune it to T=4 with an additional 20 epochs. Moreover, we use the model trained on ImageNet classification to finetune the detection or segmentation heads. This is the first time that the SNN domain has been able to process dense prediction tasks in a unified way.
- No backbone with direct training in SNNs can handle classification, detection, and segmentation tasks concurrently. Only recently did the SNN domain have the first directly trained model to process detection.
- We evaluate Meta-SpikeFormer on the COCO benchmark. We first transform mmdetection codebase into a spiking version and use it to execute our model. We exploit Meta-SpikeFormer with Mask R-CNN. ImageNet pre-trained weights are utilized to initialize the backbones, and Xavier to initialize the added layers.
- No SNN has yet reported processing results on ADE20K. In this work, Meta-SpikeFormers are evaluated as backbones equipped with Semantic FPN. ImageNet trained checkpoints are used to initialize the backbones while Xavier is utilized to initialize other newly added layers. We transform mmsegmentation codebase into a spiking version and use it to execute our model.
事件驱动的行为识别：基于事件的视觉（也称为“神经形态视觉”）是神经形态计算最典型的应用场景之一。著名的神经形态相机，动态视觉传感器（DVS），仅在亮度变化时将视觉信息编码成稀疏的事件（带有地址的脉冲）流。由于脉冲驱动的特性，SNNs 在处理事件时具有固有的低功耗和低延迟优势。
- We use HAR-DVS to evaluate our method. HAR-DVS (Wang et al., 2022) is the largest event-based Human Activity Recognition (HAR) dataset currently, containing 300 classes and 107,646 samples, acquired by a DAVIS346 camera with a spatial resolution of 346 × 260. The raw HAR-DVS is more than 4TB, and the authors convert each sample into frames to form a new HAR-DVS.
- We handle HAR-DVS in a direct training manner with T = 8. Meta-SpikeFormer achieves comparable accuracy to ANNs and is better than the Conv-based SNN baseline.
Energy Cost: The energy cost of ANNs is FLOPs times EMAC . The energy cost of SNNs is FLOPs times EAC times spiking firing rate. EMAC = 4.6pJ and EAC = 0.9pJ are the energy of a MAC and an AC, respectively.

对未来神经形态芯片（neuromorphic chip design）的启发

Meta-SpikeFormer 对芯片设计的技术创新启示主要体现在三个方面：

Conv+ViT 设计。这种混合的渐进式局部 - 全局建模可以利用 CNN 和 Transformer 的优势，前者建模特征，后者捕捉长距离依赖。通过实验验证了这种设计对 SNNs 的性能和多功能性是有益的。
脉冲驱动自注意力（SDSA）算子是 Transformer 基础 SNN 块中长距离依赖建模的核心设计，但这是当前神经形态芯片所缺乏的设计。
元架构。鉴于元 Conv-based 和 Transformer-based 块，研究人员可以根据在准确性、参数和功耗方面的要求，对元 SNN 块内的设计细节进行针对性优化。

Meta-SpikeFormer 对芯片设计的意义也体现在三个方面：

Algorithm-hardware co-design：大多数神经形态芯片设计开始于计算堆栈的底部，即材料和设备。所提算法展现的卓越特性可能会吸引并启发算法驱动的芯片设计。
对大规模神经形态计算的信心。小规模神经形态计算已经显示出显著的功耗和性能优势。本文展示了更大规模 SNNs 在性能和多功能性方面的潜力。
降低芯片设计成本。元设计便于后续研究人员进行优化，有助于 SNN 领域快速缩小与 ANNs 的差距，并降低了算法驱动硬件设计前所需的算法探索成本。