[논문 리뷰] Tensor Decomposition Based Attention Module for Spiking Neural Networks

Haoyu Deng, Ruijie Zhu외 4명|arXiv (Cornell University)|2023. 10. 23.

Advanced Memory and Neural Computing인용 수 4

한 줄 요약

이 논문은 선형적으로 증가하는 파rameter를 가진 효율적이고 저랭크이며 3차원(시간-채널-공간) 주의를 가능하게 하기 위해 텐서 분해를 사용하는 스파iking 신경망(SNNs)을 위한 투영된 전체 주의(PFA) 모듈을 제안한다. CP 분해와 선형 투영 및 주의 맵 조합을 통한 역과정을 활용함으로써 PFA는 계산 오버헤드를 최소화하면서 정적 및 동적 SNN 벤치마크에서 최신 기술 수준의 성능을 달성한다.

ABSTRACT

The attention mechanism has been proven to be an effective way to improve spiking neural network (SNN). However, based on the fact that the current SNN input data flow is split into tensors to process on GPUs, none of the previous works consider the properties of tensors to implement an attention module. This inspires us to rethink current SNN from the perspective of tensor-relevant theories. Using tensor decomposition, we design the extit{projected full attention} (PFA) module, which demonstrates excellent results with linearly growing parameters. Specifically, PFA is composed by the extit{linear projection of spike tensor} (LPST) module and extit{attention map composing} (AMC) module. In LPST, we start by compressing the original spike tensor into three projected tensors using a single property-preserving strategy with learnable parameters for each dimension. Then, in AMC, we exploit the inverse procedure of the tensor decomposition process to combine the three tensors into the attention map using a so-called connecting factor. To validate the effectiveness of the proposed PFA module, we integrate it into the widely used VGG and ResNet architectures for classification tasks. Our method achieves state-of-the-art performance on both static and dynamic benchmark datasets, surpassing the existing SNN models with Transformer-based and CNN-based backbones.

연구 동기 및 목표

• 기존 SNNs에서 텐서 인식 주의 메커니즘이 부족한 문제를 해결하기 위해, 이는 일반적으로 이벤트 데이터를 텐서로 간주하지만 텐서의 구조를 활용하지 않는 방식이다.
• 시간, 채널, 공간 차원 전역에서 전체 주의를 가능하게 하면서도 저비용 파rameter 증가를 유지하는 경량 주의 모듈을 설계하기 위해.
• 특히 CP 분해와 그 역과정인 텐서 분해 이론을 활용하여 SNNs에서 효율적인 주의 맵 생성을 위해.
• 뉴모포닉 컴퓨팅 및 GPU 가속 프레임워크와 호환되는 이론적으로 탄탄하고 해석 가능하며 확장 가능한 주의 메커니즘을 제공하기 위해.

제안 방법

• 입력 스파이크 텐서를 학습 가능한 차원별 별도 선형 투영을 통해 세 개의 낮은 차원의 투영 텐서로 압축하는 데 사용되는 스파이크 텐서의 선형 투영(LPST) 모듈을 제안한다.
• CP 분해 과정을 역행하여 연결 요소를 사용해 전체 주의 맵을 재구성하는 주의 맵 조합(AMC) 모듈을 도입한다.
• 최적의 분해 랭크 $ r $ 를 결정하기 위해 텐서 랭크 분석과 데이터셋 특성 기반의 랭크 선택 기준을 활용한다. 이는 성능과 효율성을 보장한다.
• 정교화된 주의 맵을 원본 입력 텐서와 히터드르 곱을 통해 융합함으로써 동적 특징 재가중을 가능하게 한다.
• 정적(CIFAR10/100) 및 동적(CIFAR10DVS) 데이터셋에서 VGG 및 ResNet 아키텍처 내부에 PFA 모듈을 적용하여 엔드 투 엔드 학습을 수행한다.
• 백프로파게이션과 호환되는 미분 가능 프레임워크를 사용하여 SNNs의 엔드 투 엔드 학습을 가능하게 한다.

Figure 3: The detailed workflow of PFA. The input tensor is first sent to Linear Projection of Spike Tensor (LPST) unit to generate three projections and split the three projections into corresponding vectors. In Attention Map Composing (AMC) unit, these vectors are composed into the final attention

실험 결과

연구 질문

• RQ1 스파이킹 신경망에서 텐서 분해를 효과적으로 활용하여 저비용이고 파rameter 효율적인 주의 메커니즘을 설계할 수 있는가?
• RQ2 입력 텐서의 랭크와 연결 요소 $ r $ 의 선택이 SNNs에서 주의 맵 품질과 모델 성능에 미치는 영향은 어떠한가?
• RQ3 시간, 채널, 공간 주의 구성 요소가 SNNs의 분류 정확도에 기여하는 정도는 어느 정도인가?
• RQ4 텐서 기반 주의 메커니즘이 정적 및 동적 벤치마크 데이터셋에서 기존의 트랜스포머 및 CNN 기반 SNN 모델을 초월할 수 있는가?

주요 결과

• PFA는 CIFAR10(94.85% 정확도)과 CIFAR100(78.65% 정확도)에서 최신 기술 수준의 성능을 달성하며, 트랜스포머 및 CNN 기반 백본을 사용하는 기존 SNN 모델을 모두 능가한다.
• 추론 분석 결과, 세 차원 중 어느 하나의 주의를 제거할 경우 성능이 저하되며, 특히 공간 주의가 가장 큰 영향을 미친다(제거 시 79.60%로 감소).
• 단 한 개의 PFA 레이어를 삽입하는 것만으로도 네트워크 정확도가 향상되고, PFA 레이어 수가 증가할수록 성능이 단조롭게 향상됨을 확인하여 확장성과 일관된 이점이 있음을 입증한다.
• 희소 SNN에서도 주의 맵의 랭크가 낮게 유지됨을 확인하여 희소 표현에서 여전히 최적화 잠재력이 남아 있음을 시사한다.
• 시각화 결과 PFA가 세 차원 전역에서 의미 있는 주의 분포를 학습하는 것으로 확인되었으며, 공간 주의가 가장 두드러지게 나타나고, 이어서 시간 및 채널 차원이 이어진다.
• 텐서 랭크와 연결 요소 분석 기반의 제안된 랭크 선택 기준이 최적의 $ r $ 를 효과적으로 예측하며, 실험적 성능와 잘 일치함을 확인하였다.

Figure 8: Attention visualization on CIFAR10DVS dataset. Temporal and channel parts are taken from the deep layer of VGGSNN. The spatial layer visualizations are extracted from the shallow layer of the network due to the inherent complexity of the information contained in the deeper layers, which ma