[논문 리뷰] Spikformer V2: Join the High Accuracy Club on ImageNet with an SNN Ticket

Zhaokun Zhou, Kaiwei Che외 6명|arXiv (Cornell University)|2024. 01. 04.

Advanced Memory and Neural Computing인용 수 12

한 줄 요약

Spikformer V2는 Spiking Self-Attention과 Spiking Convolutional Stem을 도입하여 SNN으로 ImageNet 정확도 80%를 초과하는 자기지도 학습(pre-training)을 가능하게 하며, 16 계층과 SSL에서 81.10%를 포함합니다.

ABSTRACT

Spiking Neural Networks (SNNs), known for their biologically plausible architecture, face the challenge of limited performance. The self-attention mechanism, which is the cornerstone of the high-performance Transformer and also a biologically inspired structure, is absent in existing SNNs. To this end, we explore the potential of leveraging both self-attention capability and biological properties of SNNs, and propose a novel Spiking Self-Attention (SSA) and Spiking Transformer (Spikformer). The SSA mechanism eliminates the need for softmax and captures the sparse visual feature employing spike-based Query, Key, and Value. This sparse computation without multiplication makes SSA efficient and energy-saving. Further, we develop a Spiking Convolutional Stem (SCS) with supplementary convolutional layers to enhance the architecture of Spikformer. The Spikformer enhanced with the SCS is referred to as Spikformer V2. To train larger and deeper Spikformer V2, we introduce a pioneering exploration of Self-Supervised Learning (SSL) within the SNN. Specifically, we pre-train Spikformer V2 with masking and reconstruction style inspired by the mainstream self-supervised Transformer, and then finetune the Spikformer V2 on the image classification on ImageNet. Extensive experiments show that Spikformer V2 outperforms other previous surrogate training and ANN2SNN methods. An 8-layer Spikformer V2 achieves an accuracy of 80.38% using 4 time steps, and after SSL, a 172M 16-layer Spikformer V2 reaches an accuracy of 81.10% with just 1 time step. To the best of our knowledge, this is the first time that the SNN achieves 80+% accuracy on ImageNet. The code will be available at Spikformer V2.

연구 동기 및 목표

• SNNs로 고성능 이미지 분류를 가능하게 하고 동기를 부여합니다.
• 곱셈 없이도 에너지 효율적인 Spiking Self-Attention(SSA)를 도입합니다.
• 추가 합성곱 층으로 바이어스되는 구조를 갖춘 Spiking Convolutional Stem(SCS)을 개발합니다.
• Self-Supervised Learning을 활용하여 더 크고 깊은 Spikformer V2를 훈련합니다.
• 순수 SNN에 대해 최첨단 ImageNet 정확도를 입증합니다.

제안 방법

• softmax 기반 주의(attention)를 스파이크 형태의 Q, K, V와 희소하고 곱셈이 없는 계산으로 대체하기 위한 Spiking Self-Attention(SSA)을 제안합니다.
• SPS를 대체하고 추가 합성곱 층을 통해 귀납적 편향을 주입하는 Spiking Convolutional Stem(SCS)을 도입합니다.
• SCS와 더 깊은 아키텍처를 통합하면서 훈련 안정성을 모니터링하며 Spikformer를 Spikformer V2로 확장합니다.
• 비대칭적 SNN-ANN 디코더를 갖는 Mask Autoencoder 형식의 Self-Supervised Learning(SSL)을 적용하여 대형 Spikformer V2 모델을 사전 학습합니다.
• SSL 후에 ImageNet에서 모델을 미세조정하여 한정된 타임 스텝으로 높은 정확도를 달성합니다.

실험 결과

연구 질문

• RQ1 softmax와 곱셈 없이도 SNN에서.self-attention을 효과적으로 구현할 수 있는가?
• RQ2 합성곱 스템이 Spikformer의 스파이크 기반 특징 추출과 ImageNet 성능을 개선하는가?
• RQ3 SSL 사전 학습이 더 크고 깊은 Spikformer 모델이 ImageNet에서 80% 이상의 정확도를 달성하도록 할 수 있는가?
• RQ4 ImageNet에서 Spikformer V2의 타임 스텝, 모델 깊이, 정확도 간의 트레이드오프는 무엇인가?
• RQ5 사전 학습에만 유리하며 SNN 추론에는 불리하지 않은 ANN-Decoder가 있는가?

주요 결과

• 8 계층 및 4 타임 스텝으로 구성된 Spikformer V2는 지도 학습에서 ImageNet에서 80.38% 정확도를 달성합니다.
• 자기지도식 사전 학습은 16-계층 Spikformer V2를 단 1 타임 스텝으로 81.10% 정확도까지 달성하게 해줍니다.
• Spikformer V2는 SNN 중에서 ImageNet에서 처음으로 80% 이상의 정확도를 달성합니다.
• SSA는 소프트맥스를 제거하고 spike-form Q, K, V를 활용하며 곱셈이 없는 연산으로 계산을 줄입니다.
• SCS의 합성곱 블록은 SPS보다 성능이 향상되며 효과적인 SSL 사전 학습을 가능하게 합니다.
• SSL 사전 학습은 더 큰 모델의 학습을 안정화하고 ImageNet에서의 성능 향상을 제공합니다.