Neural Dynamics Self-Attention for Spiking Transformers

1. 연구 배경 및 문제 정의 (Problem)

스파이크 신경망 (SNN) 과 Transformer 의 결합은 에지 비전 애플리케이션에서 에너지 효율성과 성능을 균형 있게 맞추기 위한 유망한 접근법입니다. 그러나 기존 스파이크 Transformer는 두 가지 치명적인 한계에 직면해 있습니다.

1. 성능 격차 (Performance Gap):
   • 기존 Transformer 의 핵심인 Softmax 연산이 SNN 의 **스파이크 셀프 어텐션 (SSA)**에서 제거되면서, 어텐션 분포가 급격히 변했습니다.
   • VSA (Vanilla Self-Attention): 인접한 토큰 간 강한 상관관계를 반영하여 **국소적 (Local)**이고 **희소 (Sparse)**한 어텐션 분포를 보입니다.
   • SSA: Softmax 가 없어 어텐션 점수가 거의 균일하게 분포하며, **전역적 (Global)**인 패턴을 보입니다. 이로 인해 SNN 이 중요한 지역적 정보에 집중하는 능력이 떨어지고, ANN 기반 Transformer 대비 성능이 저하됩니다.
2. 높은 메모리 오버헤드 (High Memory Overhead):
   • SSA 는 $N \times N$ 크기의 어텐션 행렬 (또는 $QK$행렬) 을 명시적으로 저장해야 하므로, 추론 시 메모리 사용량이$O(N^2)$또는$O(d^2)$ 수준으로 급증합니다. 이는 리소스가 제한된 에지 장치나 뉴로모픽 칩에서의 배포를 어렵게 만듭니다.

2. 제안 방법론 (Methodology)

저자들은 생물학적 시각 뉴런의 **국소 수용野 (Local Receptive Field, LRF)**와 막 전위 (Membrane Potential) 동역학에서 영감을 받아 LRF-Dyn을 제안했습니다. 이 방법은 두 단계로 구성됩니다.

2.1 국소 수용野를 활용한 SSA (LRF-SSA)

• 개념: SSA 메커니즘에 국소 수용野 (LRF) 모듈을 통합하여 인접한 영역에 더 높은 가중치를 부여합니다.
• 구현:
  • 전역 어텐션 (Global Receptive Fields) 에 더해, 다중 스케일 확장 컨볼루션 (Multi-scale dilated convolutions, $3\times3$ 커널, dilation $d=3, 5$) 을 사용하여 국소적 상호작용을 모델링합니다.
  • 수식적으로 어텐션 가중치를 $Q \times K^T$ (전역) 와 국소 컨볼루션 항의 합으로 정의합니다.
• 효과:
  • 이론적 분석 (Theorem 1 & 2): LRF-SSA 는 VSA 와 유사한 낮은 엔트로피 (Low-entropy) 분포를 가지며, 인접 토큰에 대한 민감도를 높여 국소 모델링 능력을 강화함을 증명했습니다.
  • Softmax 를 사용하지 않으면서도 VSA 의 국소성 특징을 모방하여 성능을 개선합니다.

2.2 뉴런 동역학을 통한 메모리 제거 (LRF-Dyn)

• 개념: LRF-SSA 에서 발생하는 어텐션 행렬 저장 오버헤드를 제거하기 위해, 전하 - 발화 - 리셋 (Charge-Fire-Reset) 동역학을 기반으로 어텐션 계산을 재구성합니다.
• 구현:
  • 메모리 잠재력 (Memory Potential): 이전 토큰들의 $K^T V$ 곱셈 결과를 누적하여 '막 전위'로 저장합니다.
  • 시냅스 입력 (Presynaptic Input): 현재 토큰의 입력을 추가합니다.
  • 동역학 방정식: $X_n[t] = A \odot X_{n-1}[t] + \Gamma \cdot Token_n[t]$ 형태의 재귀적 구조를 사용하여, 명시적인 $N \times N$ 어텐션 행렬 저장 없이 어텐션 계산을 수행합니다.
  • 다중 가지 (Dendritic) 구조: 광수용체 뉴런의 다중 시간 척도 행동을 모방하여, 다양한 가지 (dendrites) 를 통해 어텐션 점수를 할당하고 soma 에서 통합합니다.
• 효과:
  • 추론 시 어텐션 행렬을 저장할 필요가 없어 메모리 복잡도가 $O(Nd^2)$에서 $O(kd)$(여기서$k$는 가지 수) 로 감소합니다.
  • 이벤트 기반 (Event-driven) 특성을 유지하며 에지 장치 배포에 최적화됩니다.

3. 주요 기여 (Key Contributions)

1. 문제 식별: 기존 Spiking Transformer 의 성능 저하 원인이 Softmax 제거로 인한 국소 모델링 능력 부재와 높은 메모리 오버헤드임을 이론 및 실험적으로 규명했습니다.
2. LRF-Dyn 제안:
   • LRF-SSA: 국소 수용野 모듈을 도입하여 SSA 의 국소 모델링 능력을 강화하고 성능을 ANN 수준에 근접시킵니다.
   • LRF-Dyn: 뉴런의 동역학을 활용하여 어텐션 행렬 저장 없이 계산을 수행함으로써 추론 메모리를 획기적으로 줄입니다.
3. 실증적 검증: 다양한 비전 작업 (이미지 분류, 시맨틱 세그멘테이션) 과 아키텍처 (Spikformer, QKFormer, Spike-Driven-v3) 에서 제안 방법이 성능 향상과 메모리 감소 동시 달성을 입증했습니다.

4. 실험 결과 (Results)

이미지 분류 (ImageNet-1k):

• Spikformer-8-512: LRF-SSA 적용 시 정확도 74.62% (기존 대비 +1.24% 향상), LRF-Dyn 적용 시 74.51% (+1.13% 향상).
• QKFormer: LRF-SSA 적용 시 정확도 82.52% (+0.48% 향상).
• SDT-V3: LRF-SSA 적용 시 76.22% 정확도 달성 (매우 적은 파라미터로 경쟁 모델 대비 우수).
• 메모리 효율: LRF-Dyn 은 $O(kd)$의 저장 복잡도를 가지며, Spikformer-8-512 기준 메모리 사용량을 49.4% 감소시키면서도 정확도를 유지/향상시켰습니다.

시맨틱 세그멘테이션 (ADE20K):

• SDT-V3 + LRF-SSA: mIoU 43.5% (기존 대비 +2.2% 향상).
• SDT-V3 + LRF-Dyn: mIoU 43.1% (기존 대비 +1.8% 향상).
• Qualitative: 제안된 방법은 기존 SSA 보다 더 세밀한 (fine-grained) 분할 결과를 보여주며, 중요한 영역에 더 집중하는 어텐션 패턴을 보입니다.

Ablation Study:

• LRF 모듈 (커널 수 증가) 이 성능 향상에 결정적 역할을 함을 확인했습니다.
• LRF-Dyn 은 인과적 SSA 모델보다 동일한 조건에서 일관되게 더 나은 성능을 보였습니다.

5. 의의 및 결론 (Significance)

이 논문은 에지 환경에서 고품질 SNN 모델 배포를 위한 핵심 솔루션을 제시합니다.

• 이론적 통찰: Softmax 없는 어텐션이 왜 국소적 모델링에 실패하는지, 그리고 생물학적 뉴런의 동역학을 어떻게 이를 해결하는 데 활용할 수 있는지에 대한 명확한 이론적 근거를 제공했습니다.
• 실용적 가치:
  • 에너지 효율성: SNN 의 이벤트 기반 특성을 유지하면서 Transformer 의 성능을 확보합니다.
  • 메모리 효율성: 대규모 어텐션 행렬 저장 없이 추론이 가능하여, 제한된 메모리를 가진 뉴로모픽 칩이나 모바일 장치에서의 적용 가능성을 크게 높였습니다.
• 미래 전망: LRF-Dyn 은 차세대 저전력 비전 Transformer 의 기본 구성 요소 (Unit) 로서, 에지 AI 의 성능과 효율성 한계를 확장하는 중요한 이정표가 될 것으로 기대됩니다.