[论文解读] 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 提出了一种时空 Transformer 架构，通过在神经元群体中同时建模时间动态和空间协变性，扩展了神经数据 Transformer。通过引入空间自注意力机制和对比学习微调损失，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.

研究动机与目标

• 建模单次试验中神经元的时间演化与神经元间空间协变性的神经群体动力学。
• 通过显式建模神经元间相关性与时间动态，提升对未见神经元和时间步的预测性能。
• 通过一种新型对比学习微调损失，增强模型对噪声的鲁棒性，该损失与掩码建模目标互补。
• 通过空间注意力机制识别驱动群体响应的关键神经元子集，提升模型可解释性。
• 在涵盖自主与非自主动力学、且不依赖特定行为任务结构的多样化神经数据集上评估模型性能。

提出的方法

• STNDT 采用双注意力机制：分别为空间（神经元间）和时间（时间步间）依赖关系设计独立的多头自注意力模块。
• 模型将神经元放电活动处理为矩阵 $X \in \mathbb{N}^{N \times T}$，其中 $N$ 为神经元数量，$T$ 为时间点数。
• 空间注意力应用于时间注意力特征矩阵，使神经元既能相互关注，也能跨时间关注。
• 引入对比学习损失，通过鼓励在噪声和掩码等数据增强下保持一致的预测结果，提升模型鲁棒性。
• 模型首先使用掩码建模目标进行预训练，再通过对比学习进行微调，以增强泛化能力与预测准确率。
• 采用集成模型提升性能，利用多个模型实例的预测结果。

实验结果

研究问题

• RQ1 基于 Transformer 的架构若同时建模神经群体活动中的空间与时间依赖关系，是否能在预测单次试验神经动力学方面超越现有方法？
• RQ2 引入空间注意力是否能提升对驱动群体响应的功能性重要神经元子集的识别能力？
• RQ3 对比学习微调在噪声或增强数据下在多大程度上提升了模型的鲁棒性与预测性能？
• RQ4 STNDT 是否能在具有不同动力学特性（自主 vs. 非自主）和行为任务结构的多样化神经数据集上实现泛化？
• RQ5 STNDT 中的空间注意力权重如何揭示可解释且一致的神经群体计算模式？

主要发现

• 在神经潜在变量基准套件的四个神经数据集中，STNDT 在集成模型层面实现了最先进性能，优于先前方法在预测保留神经元和时间步活动方面的表现。
• 在 MC_Maze 数据集上，采用对比损失的集成 STNDT 模型取得了 0.3862 的 co-bps 得分，较无对比损失的模型提升了 0.0189。
• 对比损失在单模型与集成模型层面均提升了性能，其中在 MC_Maze 和 DMFC_RSG 数据集上提升最为显著。
• STNDT 中的空间注意力机制持续识别出在驱动群体响应中起关键作用的有影响力神经元子集，为神经计算提供了可解释性。
• 模型从发放率预测中保持了强大的解码性能，证明其在脑机接口应用中的实用性。
• 在小规模或非结构化数据集（如 Area2_Bump、MC_RTT）中，仅使用空间注意力未能带来性能提升，表明有效空间特征学习需要足够的数据规模与结构化行为。