[SpikeFormer 논문리뷰] SPIKFORMER: WHEN SPIKING NEURAL NETWORK MEETS TRANSFORMER

https://openreview.net/pdf?id=frE4fUwz_h

ICLR 2023

1저자 : Zhaokun Zhou

Peking University, Peng Cheng Laboratory, Sea AI Lab, Shenzhen EEGSmart Technology Co., Ltd

 

 

ABSTRACT

 

- 우리는 생물학적으로 그럴듯한 구조 2개, Spiking Neural Network (SNN)와 self-attention mechanism을 고려한다.

- SNN은 deep learning을 위한 에너지 효율적이고 event-driven paradigm을 제공하고, Self-attention mechanism은 feature dependencies를 포착할 수 있는 능력을 가지고 있기 때문에 Transformer는 좋은 성능을 성취할 수 있게 합니다.

- 이 둘을 결합하는 탐구는 직관적으로 유망하다.

- 본 논문에서, 우리는 self-attention의 능력과 SNN의 생물학적 특성을 활용하는 것을 고려하며, 새로운 Spiking Self Attention (SSA) 뿐만 아니라 강력한 framework인 Spiking Transformer (Spikfromer)를 제안한다.

- Spikformer에서 SSA mechanism은 softmax 없이 spike 형식의 Query, Key 그리고 Value를 사용하여 sparse visual feature를 모델링한다.

- 이것의 computation은 sparse하고 multiplication을 제거하기 때문에, SSA는 효과적이고 낮은 computational energy 소비를 갖는다.

- SSA를 갖는 Spikformer는 neuromorphic과 static datasets에 대한 image 분류에 있어서 기존 sota SNN 형식의 frameworks들 보다 성능이 뛰어날 수 있음을 보인다.

- Spikformer (66.3M parameters)는 SEW-ResNet-152(60.2M, 69.26%)와 유사한 크기이면서, 4번의 steps만에 ImageNet에 대한 top1 정확도 74.81%를 달성하였고, 이는 훈련된 SNNs models들 중에 sota이다.

 

1 INTRODUCTION

 

ANN의 발전과 함께, SNN는 ANN의 발전된 구조들을 차용하여 성능을 향상시킬 수 있다.

- ResNet-like SNNs (Hu et al., 2021)

- Spiking Recurrent Neural Networks (Lotfi Rezaabad & Vishwanath, 2020)

- Spiking Graph Neural Networks (Zhu et al., 2022)

 

원래, 자연어 처리 (NLP)를 위해 고안된 Transformer (Vaswani et al., 2017)는 computer vision의 다양한 task에까지 번성하였다. image classification, object detection, semantic segmentation 그리고 low-level image processing

Transformer의 핵심 부분인 Self-attention은 관심 있는 정보에 선택적으로 집중하는데, 이는 인간 생물학적 시스템의 중요한 특성 중 하나이기도 하다.

직관적으로, 이 두 메커니즘의 생물학적 특성들을 고려할 때, SNNs에서 self-attention을 적용하는 탐구는 흥미로운 일이다.

 

SNN에 self-attention 메커니즘을 적용하는 것에 어려움

vanilla self-attention, VSA 에서 3가지 구성요소 Query, Key, Value가 있다.

Figure 1(a)에서 볼 수 있듯이, VSA의 표준 inference는 먼저 실수 형태의 Query와 Key의 내적(dot-product)를 계산하여 행렬을 얻는다.

이후, 지수 계산과 나눗셈 연산을 포함한 softmax를 사용하여 matrix를 normalize하고 이를 attention map을 생성한다. 그리고 이것은 Value를 가중치로 적용하는데 사용된다.

이러한 과정에서 볼 수 있듯이, VSA의 막대한 computational overhead를 SNN에 직접 적용하는 것은 불가능하다.

따라서, SNN에 Transformer를 개발하기 위해서, multiplications를 피할 수 있는 새롭고 효과적이며 계산 효율적인 self-attention 변형이 필요했다.

 

우리는 그래서, Figure 1(b)에서처럼 Spiking Self Attention (SSA)를 나타낸다. SSA는 self-attention 메커니즘을 SNN에 처음 도입하였으며, 이는 spike sequences를 사용하여 interdependence를 모델링한다.

SSA에서 Query, Key, Value는 spike 형태로 표현되며, 0 또는 1의 값 만을 갖는다

 

SNN에서 self-attention을 적용하는데 주요 어려움은 softmax 때문이다.

1) 그림1에서 보듯이, spike 형태의 Query와 Key로 계산된 attention map은 자연스럽게 non-negative를 가지며, 이는 irrelevant features를 무시해버린다. 따라서 softmax가 우리에게는 필요하지 않다.

2) SSA의 input과 Value는 spike 형태로, 0과 1만 포함하고 있으며, 이는 ANN의 VSA에서의 실수형 input과 Value보다 세밀한 feature를 포함하지 않는다. 따라서, 실수형 Query와 Ket 및 softmax 함수는 이러한 spike sequences를 모델링하는데 있어서 중복된다.

이러한 인사이트를 기반으로, 우리는 SSA에서 attention map을 위한 softmax normalization을 제거한다.

이전에 몇몇 Transformer 관련 논문들도 softmax를 제거하거나 linear function으로 대체하였다. 예를들어, CosFormer (Qin et al., 2022)에서는 softmax를 ReLU와 코사인 함수로 대체한다.

 

이러한 SSA에 대한 설계를 통해, spike 형식의 Query, Key 그리고 Value의 계산은 multiplications를 피하고 AND와 덧셈 연산으로 수행될 수 있다. 또한, 이는 계산적으로도 매우 효과적이다.

Sparse spike-form Query, Key, Value 그리고 간단한 계산 덕분에, SSA에서 연산의 수는 작고 이는 SSA의 에너지 소비를 매우 적게 만들어준다. 더욱이, softmax가 제거되어 우리의 SSA는 분해 가능하며, 시퀀스 길이가 하나의 head의 feature dimension보다 더 클 때, 계산 복잡성이 더욱 줄어든다.

요약하자면, 우리는 다음과 같이 3개의 컨트리뷰션으로 요약할 수 있다.

 

- SNN의 특성을 갖춘 Spiking Self Attention (SSA)의 새로운 spike-form self-attention을 설계하였다. Softmax 없이 sparse spike-form Query, Key, and Value를 사용하여, SSA의 계산은 곱셈을 피하고 효율성을 높일 수 있었다.

 

- 우리는 우리가 제안한 SSA를 기반으로 Spiking Transformer (Spikformer)를 개발하였다. 우리가 아는한, SNN에서 self-attention과 Transformer를 구현한 최초의 사례이다.

 

- Extensive experiments에서 우리가 제안한 구조가 static과 neuromorphic datasets 둘 다에 대해 SNN 분야 sort를 달성함을 보인다. 특히, 직접 훈련된 SNN 모델을 사용하여 4번의 steps로 ImageNet에서 74% 이상의 정확도를 달성한 것은 이번이 처음이다.  

 

나의 요약:

기존 Transformer는 Query랑 Key를 내적해서 행렬계산을 하고 만든 Attention Map을 또 Softmax하고 또 Value랑 행렬 연산을 해야해서 계산량이 많음. 이런 과정은 SNN의 특성에 맞지 않음

그래서 Spike 형태의 0과 1만 갖는 Query, Key, Value를 포함하는 Spike Self-Attention (SSA)을 만듦으로써 self-attention 메커니즘을 SNN에 처음 도입함.

근데, Spike 형태의 0과 1만 갖는 Query, Key, Value를 갖게 하니깐 attention map을 위한 softmax normalization 과정이 필요없게 됨. 그래서 soft 연산을 제거함. 이는 기존 Transformer 논문에서도 softmax 대신 ReLU와 Cosin 함수로 대처하는 등 사례가 있음. 그리고 이렇게 softmax를 제거함으로써 SSA는 분해가 가능함

 

-> 의문: 뭔가 당연한 소리를 길게 늘여놓은거 같다. 과정과 결과의 순서가 바뀐 느낌…? 이렇게 간단하면 왜 지금까지 힘들게 어려운 연산을 했겠는가…. 0과 1로만 하면 Gradient Vanishing 안 일어나나?

 

 

3 METHOD

3.1 OVERALL ARCHITECTURE

 

2D image sequence를 입력 받으면, Spiking Patch Splitting (SPS) 모듈이 이것을 D 차원 spike-form feature vector로 선형 변환하고 이를 N개의 flattened spike-form patches x의 시퀀스로 분할한다.  

실수형 임베딩은 SNN에 사용될 수 없다.

우리는 a conditional position embedding generator (Chu et al., 2021)를 사용하여, spike-form relative position embedding (RPE)를 생성하고 RPE를 patches sequence x에 더하여 X0을 얻는다.

conditional position embedding generator은 3x3크기의 kernel을 갖는 2D convolution layer (Conv2d), batch normalization(BN) 그리고 spike neuron layer (SN)을 포함한다.

 

그런 다음, 우리는 X0을 L-block Spikformer encoder에 통과시킨다. standard ViT encoder block과 유사하게, Spikformer encoder block은 Spiking Self Attention (SSA)와 MLP block으로 구성되어 있다.

Residual connections가 SSA와 MLP block 둘 다에 적용된다.

Spikformer encoder block의 주요 구성 요소로써, SSA는 softmax 없이 spike 형태의 Query (Q), Key (K), Value (V)를 사용하여 이미지의 지역-전역 정보(local-global information)를 효율적으로 모델링하는 방법을 제공한다.

global average-pooling (GAP)은 Spikformer encoder에서 처리된 feature에 적용되고 D 차원 feature를 output한다, 이 D 차원 feature는 prediction Y를 output하기 위해 fully-connected-layer classification head (CH)로 보내진다.

 

(4): Spiking Patch Splitting (SPS) 모듈이 이미지 시퀀스 I를 입력받아 D 차원 spike-form feature vector로 선형 변환하고 이를 N개의 flattened spike-form patches x의 시퀀스로 분할한다

(5): 이 x를 conditional position embedding generator가 받아, spike-form relative position embedding (RPE)를 생성

(6): RPE를 patches sequence x에 더하여 X0을 얻는다.

(7, 8)그런 다음, 우리는 X0을 L-block Spikformer encoder에 통과시킨다. Spikformer encoder block은 Spiking Self Attention (SSA)와 MLP block으로 구성되어 있다.

(9): global average-pooling (GAP)은 Spikformer encoder에서 처리된 feature에 적용되고 D 차원 feature를 output한다, 이 D 차원 feature는 prediction Y를 output하기 위해 fully-connected-layer classification head (CH)로 보내진다.