[論文レビュー] STNDT: Modeling Neural Population Activity with a Spatiotemporal Transformer

Trung Le, Eli Shlizerman|arXiv (Cornell University)|Jun 9, 2022

Neural dynamics and brain function被引用数 5

ひとこと要約

STNDTは、神経集団内の個々のニューロン間の時間的ダイナミクスと空間的共変動をモデル化することで、Neural Data Transformersを拡張した時空間変換器アーキテクチャを導入する。空間的自己注意と対照的微調整損失を組み合わせることで、4つの多様な神経データセットにおいて未観測の神経活動を予測する上で最先端の性能を達成するとともに、空間的自己注意メカニズムにより解釈可能な影響力を持つニューロンサブセットを特定する。

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は、空間的（ニューロン間）および時間的（時刻ステップ間）依存性を別々に処理する二重自己注意メカニズムを採用する。
• モデルは、$N$ をニューロン数、$T$ を時刻点数とする行列 $X \in \mathbb{N}^{N \times T}$ として神経スパiking活動を処理する。
• 空間的自己注意は、時間的自己注意の特徴行列に適用され、ニューロンが互いに注目するだけでなく、時間軸に跨って注目する能力を有する。
• データオーグメンテーション（ノイズやマスキングなど）に対して一貫した予測を促進するように、モデルの耐性を高めるために対照的学習損失を導入する。
• マスクモデリング目的による事前学習と、対照的学習による微調整を組み合わせることで、汎化性能と予測精度を向上させる。
• 複数のモデルインスタンスの予測を統合するアンサンブルモデルを用いて性能を向上させる。

実験結果

リサーチクエスチョン

• RQ1 時間的および空間的依存性を両方モデル化する変換器ベースのアーキテクチャが、単一試行神経ダイナミクスの予測において、既存手法を上回る性能を示せるか？
• RQ2 空間的自己注意を組み込むことで、集団応答を駆動する機能的に重要なニューロンサブセットの同定が向上するか？
• RQ3 対照的微調整は、ノイズ混在またはオーグメント済みデータ下でのモデルの耐性および予測性能をどの程度向上させるか？
• RQ4 STNDTは、自律的および非自律的ダイナミクスを有する多様な神経データセットに一般化可能か？
• RQ5 STNDTにおける空間的自己注意重みは、神経集団計算の解釈可能な一貫性のあるパターンをどのように明らかにするか？

主な発見

• STNDTは、Neural Latents Benchmarkスイートの4つの神経データセットにおいて、アンサンブルレベルで最先端の性能を達成し、ホールドアウトされたニューロンおよび時刻点の活動予測において、先行手法を上回った。
• MC_Mazeデータセットでは、対照的損失を用いたアンサンブルSTNDTがco-bpsスコア0.3862を達成し、対照的損失なしのモデルに比べ0.0189の向上を示した。
• 対照的損失は、単体およびアンサンブルの両レベルで性能向上をもたらし、特にMC_MazeおよびDMFC_RSGデータセットで顕著な向上が観察された。
• STNDTにおける空間的自己注意は、集団応答を駆動する上で重要な役割を果たすインパルスなニューロンサブセットを一貫して同定し、神経計算の解釈可能性を提供した。
• モデルはレート予測からも強力な復号性能を維持しており、脳-コンピュータインターフェース応用における有用性を示した。
• 小規模または構造のないデータセット（例：Area2_Bump、MC_RTT）では、空間的自己注意のみでは性能向上が得られなかったことから、有効な空間的特徴学習には十分なデータスケールと構造的行動が不可欠であることが示唆された。