[论文解读] Spike-driven Transformer V2: Meta Spiking Neural Network Architecture Inspiring the Design of Next-generation Neuromorphic Chips

Man Yao, Jiakui Hu+6 more|arXiv (Cornell University)|Feb 15, 2024

Advanced Memory and Neural Computing被引用 19

一句话总结

该论文提出 Meta-SpikeFormer，这是一个基于元 Transformer 的脉冲驱动型 SNN，脉冲驱动，在 ImageNet-1K 上以 55M 参数达到 80.0% 的 top-1 准确率，并在单一直接训练的 SNN 主干上实现分类、检测和分割。此外还讨论了对神经形态芯片设计的影响。

ABSTRACT

Neuromorphic computing, which exploits Spiking Neural Networks (SNNs) on neuromorphic chips, is a promising energy-efficient alternative to traditional AI. CNN-based SNNs are the current mainstream of neuromorphic computing. By contrast, no neuromorphic chips are designed especially for Transformer-based SNNs, which have just emerged, and their performance is only on par with CNN-based SNNs, offering no distinct advantage. In this work, we propose a general Transformer-based SNN architecture, termed as ``Meta-SpikeFormer", whose goals are: 1) Lower-power, supports the spike-driven paradigm that there is only sparse addition in the network; 2) Versatility, handles various vision tasks; 3) High-performance, shows overwhelming performance advantages over CNN-based SNNs; 4) Meta-architecture, provides inspiration for future next-generation Transformer-based neuromorphic chip designs. Specifically, we extend the Spike-driven Transformer in \citet{yao2023spike} into a meta architecture, and explore the impact of structure, spike-driven self-attention, and skip connection on its performance. On ImageNet-1K, Meta-SpikeFormer achieves 80.0\% top-1 accuracy (55M), surpassing the current state-of-the-art (SOTA) SNN baselines (66M) by 3.7\%. This is the first direct training SNN backbone that can simultaneously supports classification, detection, and segmentation, obtaining SOTA results in SNNs. Finally, we discuss the inspiration of the meta SNN architecture for neuromorphic chip design. Source code and models are available at \url{https://github.com/BICLab/Spike-Driven-Transformer-V2}.

研究动机与目标

• 通过利用脉冲驱动的 SNN 实现能耗更低的 Transformer-基础视觉模型。
• 开发一个元 Transformer 架构，将 Conv 基 SNN 块与 Transformer 基 SNN 块在脉冲驱动约束下结合。
• 提高 SNN 的性能和通用性，超越标准视觉基准上的 CNN 基 SNN。
• 提供架构指引，激发未来基于 Transformer 的神经形态芯片设计。

提出的方法

• 将 Spike-driven Transformer 扩展为元架构（Meta-SpikeFormer），在早期阶段使用 Conv 基 SNN 块，在后期阶段使用 Pyramid Transformer 基 SNN 阶段。
• 引入 spike-driven 自注意力（SDSA）算子，无 softmax 或缩放，在 N 维度上具有线性或接近线性计算复杂度。
• 设计微观级别块：使用 SepConv 和 ChannelConv 的 Conv 基 SNN 块；使用 RepConv 编码并应用 SDSA 的 Transformer 基 SNN 块生成 Q_S、K_S、V_S。
• 采用四级架构，采用金字塔结构并在较高通道的 Stage 4 控制参数量。
• 探索三种捷径变体（Vanilla, SEW, Membrane Shortcut），其中 MS（Membrane Shortcut）提供强大准确性。
• 直接训练（代理梯度）以实现 SNN 主干在静态和事件数据上的端到端学习。

实验结果

研究问题

• RQ1 在脉冲驱动约束下，元 Transformer 基 SNN 架构是否能够实现高性能与低功耗的统一？
• RQ2 架构选择（Conv 与 Transformer 块、SDSA 变体与捷径）如何影响脉冲驱动 SNN 的准确性、参数量和功耗？
• RQ3 直接训练的 SNN 主干能否同时处理分类、检测和分割任务？
• RQ4 来自 Meta-SpikeFormer 的设计原则能够为未来基于 Transformer 的神经形态芯片开发提供哪些指导？

主要发现

• Meta-SpikeFormer 在 ImageNet-1K 上以 55M 参数实现 80.0% 的 top-1 准确率（T=4 在蒸馏预训练后达到 80.0%）。
• 它比当前 SOTA SNN 基线高 3.7 个百分点，且参数更少 17%（55M 对 66M）。
• 它是第一款能够直接训练并同时处理分类、检测和分割的 SNN 主干，在测试数据集上实现了 SNN 的 SOTA 结果。
• 在 ImageNet-1K 上，Meta-SpikeFormer 展示了在多个任务中的 SNN 领域的最先进性能，分析显示在准确性和通用性方面优于基于 Conv 的 SNN。
• 一种包含 Conv 基和 Transformer 基 SNN 块、SDSA 和 Membrane Shortcut 的元架构设计为未来神经形态芯片设计提供了实际指南。