STNDT: Modeling Neural Population Activity with a Spatiotemporal Transformer

TL;DR

STNDT is a novel spatiotemporal transformer model that improves neural population activity prediction by explicitly modeling neuron interactions across space and time, outperforming previous methods and offering interpretability.

Contribution

The paper introduces STNDT, a transformer-based architecture that models neuron responses across space and time, incorporating contrastive learning for enhanced prediction and interpretability.

Findings

Achieves state-of-the-art neural activity estimation across four datasets.

Effectively captures autonomous and non-autonomous neural dynamics.

Identifies key neurons influencing population responses.

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.