[논문 리뷰] STNDT: Modeling Neural Population Activity with a Spatiotemporal Transformer

Trung Le, Eli Shlizerman|arXiv (Cornell University)|2022. 06. 09.

Neural dynamics and brain function인용 수 5

한 줄 요약

STNDT는 신경 집단 내 개별 뉴런 간의 시간적 동적 변화와 공간적 공변동을 모델링하는 공간시적 트랜스포머 아키텍처를 소개한다. 공간 자기주의와 대비 학습 손실을 통합함으로써, STNDT는 네 가지 다양한 신경 데이터셋에서 볼 수 없는 신경 활동을 예측하는 데 최신 기술 수준의 성능을 달성하였으며, 공간 주의 메커니즘을 통해 해석 가능한 영향력 있는 뉴런 집합을 드러낸다.

ABSTRACT

Modeling neural population dynamics underlying noisy single-trial spiking activities is essential for relating neural observation and behavior. A recent non-recurrent method - Neural Data Transformers (NDT) - has shown great success in capturing neural dynamics with low inference latency without an explicit dynamical model. However, NDT focuses on modeling the temporal evolution of the population activity while neglecting the rich covariation between individual neurons. In this paper we introduce SpatioTemporal Neural Data Transformer (STNDT), an NDT-based architecture that explicitly models responses of individual neurons in the population across time and space to uncover their underlying firing rates. In addition, we propose a contrastive learning loss that works in accordance with mask modeling objective to further improve the predictive performance. We show that our model achieves state-of-the-art performance on ensemble level in estimating neural activities across four neural datasets, demonstrating its capability to capture autonomous and non-autonomous dynamics spanning different cortical regions while being completely agnostic to the specific behaviors at hand. Furthermore, STNDT spatial attention mechanism reveals consistently important subsets of neurons that play a vital role in driving the response of the entire population, providing interpretability and key insights into how the population of neurons performs computation.

연구 동기 및 목표

• 단일 시험 동안의 스파iking 활동에서 뉴런 간 시간적 변화와 공간적 공변동을 모두 포괄하는 신경 집단 역학을 모델링하기 위해.
• 시간적 동적 변화와 함께 상호 뉴런 상관관계를 명시적으로 모델링하여, 예측 성능을 향상시키고 볼 수 없는 뉴런과 시간대에 대한 예측 성능을 향상시키기 위해.
• 마스크 모델링 목표와 보완되는 새로운 대비 학습 손실을 통해 노이즈에 대한 모델의 강건성을 향상시키기 위해.
• 공간 주의 메커니즘을 통해 인과적인 집단 반응을 이끄는 핵심 뉴런 집합을 식별함으로써 모델의 해석 가능성을 향상시키기 위해.
• 특정 행동 과제의 구조에 의존하지 않고, 자율적 및 비자율적 동적 특성을 모두 포함하는 다양한 신경 데이터셋에서 성능을 평가하기 위해.

제안 방법

• STNDT는 공간(뉴어론 간)과 시간(시간대 간)의 의존성에 대해 별도의 멀티헤드 자기주의 모듈을 사용하는 이중 주의 메커니즘을 적용한다.
• 모델은 $X \in \mathbb{N}^{N \times T}$ 형태의 행렬로 신경 스파이크 활동를 처리하며, 여기서 $N$은 뉴런의 수이고 $T$는 시간점의 수이다.
• 공간 주의는 시간 주의 특징 행렬에 적용되어, 뉴런들이 서로를 주의를 기울이는 동시에 시간에 걸쳐 주의를 기울일 수 있도록 한다.
• 데이터 증강(노이즈 및 마스킹 포함)에 대해 일관된 예측을 유도함으로써 강건성을 향상시키기 위해 대비 학습 손실이 도입된다.
• 모델은 마스크 모델링 목표로 사전 학습하고, 대비 학습을 통해 미세조정하여 일반화 및 예측 정확도를 향상시킨다.
• 성능 향상을 위해 앙상블 모델을 사용하며, 여러 모델 인스턴스의 예측을 융합한다.

실험 결과

연구 질문

• RQ1 공간적 및 시간적 의존성을 모두 모델링하는 트랜스포머 기반 아키텍처가 기존 방법보다 단일 시험 신경 역학 예측에서 뛰어난 성능을 내는가?
• RQ2 공간 주의를 통합함으로써 집단 반응을 이끄는 기능적으로 중요한 뉴런 집합을 더 잘 식별할 수 있는가?
• RQ3 대비 미세조정이 노이즈가 있거나 증강된 데이터에서 모델의 강건성과 예측 성능을 어느 정도 향상시키는가?
• RQ4 STNDT는 자율적 및 비자율적 동적 특성을 가진 다양한 신경 데이터셋에 대해 일반화 가능한가?
• RQ5 STNDT의 공간 주의 가중치는 신경 집단 계산의 해석 가능한 일관된 패턴을 어떻게 드러내는가?

주요 결과

• STNDT는 신경 라티언트 벤치마크 세트의 네 가지 신경 데이터셋에서 앙상블 수준에서 최신 기술 수준의 성능를 기록하였으며, 보류된 뉴런과 시간대의 활동을 예측하는 데 이전 방법들을 능가하였다.
• MC_Maze 데이터셋에서 대비 손실을 사용한 앙상블 STNDT는 co-bps 점수 0.3862를 기록하여 대비 손실 없이 학습한 모델보다 0.0189 향상되었다.
• 대비 손실은 단일 및 앙상블 수준에서 성능 향상을 이끌었으며, 특히 MC_Maze 및 DMFC_RSG 데이터셋에서 가장 두드러진 성과를 보였다.
• STNDT의 공간 주의는 집단 반응을 이끄는 데 핵심적인 역할을 하는 영향력 있는 뉴런 집합을 일관되게 식별하여 신경 계산의 해석 가능성을 제공하였다.
• 모델은 속도 예측에서 강력한 복원 성능를 유지하여 뇌-컴퓨터 인터페이스 응용 분야에서의 유용성을 입증하였다.
• 작은 규모이거나 비정형 데이터셋(예: Area2_Bump, MC_RTT)에서는 공간 주의만으로 성능 향상가 발생하지 않음을 고려할 때, 효과적인 공간 특징 학습을 위해서는 충분한 데이터 규모와 정형화된 행동이 필요하다는 점이 확인되었다.