STNDT: Modeling Neural Population Activity with a Spatiotemporal Transformer
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.