[논문 리뷰] Spike-driven Transformer V2: Meta Spiking Neural Network Architecture Inspiring the Design of Next-generation Neuromorphic Chips

Man Yao, Jiakui Hu외 6명|arXiv (Cornell University)|2024. 02. 15.

Advanced Memory and Neural Computing인용 수 19

한 줄 요약

논문은 spike-driven인 메타 Transformer 기반 SNN인 Meta-SpikeFormer를 제시하며, 55M 파라미터로 ImageNet-1K에서 Top-1 정확도 80.0%를 달성하고 단일 직접 학습 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}.

연구 동기 및 목표

• 전력 효율적인 뉴로모픽 계산을 위해 spike-driven SNN을 활용한 Transformer 기반 비전 모델의 가능성을 탐색한다.
• Conv 기반 SNN 블록과 Transformer 기반 SNN 블록을 spike-driven 제약 하에 결합한 메타 Transformer 아키텍처를 개발한다.
• SNN의 성능과 활용성을 개선하여 표준 비전 벤치마크에서 CNN 기반 SNN을 능가하도록 한다.
• 향후 Transformer 기반 뉴로모픽 칩 설계를 위한 설계 원칙을 제시한다.

제안 방법

• 초기 단계에서 Conv 기반 SNN 블록을 사용하고 후반 단계에서 Pyramid Transformer 기반 SNN 스테이지를 활용하는 메타아키텍처(Meta-SpikeFormer)로 Spike-driven Transformer를 확장한다.
• softmax나 스케일링이 없고 N에서 선형 또는 준-선형 계산 복잡도를 갖는 spike-driven self-attention(SDSA) 연산자를 도입한다.
• 마이크로 레벨 블록을 설계한다: SepConv 및 ChannelConv를 사용하는 Conv 기반 SNN 블록; RepConv 기반 인코딩을 통해 Q_S, K_S, V_S를 생성하고 SDSA를 적용하는 Transformer 기반 SNN 블록.
• 파라미터를 제어하기 위해 피라미드 구조와 더 높은 채널의 스테이지(Stage 4)를 갖는 4단계 아키텍처를 사용한다.
• MS(Membrane Shortcut)를 포함한 Vanilla, SEW, Membrane Shortcut의 세 가지 단축 경로 변형을 탐색하고 Membrane Shortcut이 강력한 정확도를 제공합니다.
• 정적 데이터와 이벤트 기반 데이터에 대해 엔드 투 엔드 학습을 가능하게 하는 직접 학습(대리 기울기) 방식으로 SNN 백본을 학습한다.

실험 결과

연구 질문

• RQ1 메타 Transformer 기반 SNN 아키텍처가 spike-driven 제약 하에서 고성능과 저전력을 함께 달성할 수 있는가?
• RQ2 아키텍처 선택(Conv 블록 vs Transformer 블록, SDSA 변형, 단축 경로)의 차이가 spike-driven SNN의 정확도, 파라미터 수, 전력에 어떤 영향을 미치는가?
• RQ3 직접 학습된 SNN 백본이 분류, 검출, 분할을 동시에 처리할 수 있는가?
• RQ4 Meta-SpikeFormer의 설계 원칙이 미래의 Transformer 기반 뉴로모픽 칩 개발에 어떤 지침을 제공할 수 있는가?

주요 결과

• Meta-SpikeFormer는 55M 파라미터로 ImageNet-1K에서 80.0% top-1 정확도를 달성한다(T=4가 distillation 사전학습 후 80.0%를 달성).
• 현재의 SOTA SNN 기준선을 3.7% 포인트 상회하며 파라미터 수는 17% 더 적은 55M vs 66M로 감소한다.
• 직접 학습된 SNN 백본으로 분류, 검출, 분할을 동시에 처리할 수 있는 최초의 사례이며 테스트 데이터셋에서 SNN의 SOTA 결과를 달성한다.
• ImageNet-1K에서 Meta-SpikeFormer는 다중 작업에서 SNN 도메인에서 최첨단 성능을 보여주며 분석은 Conv 기반 SNN보다 정확도와 다재다능성 측면에서 이점을 시사한다.
• Conv 기반 및 Transformer 기반 SNN 블록, SDSA, Membrane Shortcut를 갖춘 메타아키텍처 설계는 미래의 뉴로모픽 칩 설계에 실용적인 지침을 제공한다.