# STNDT: Modeling Neural Population Activity with a Spatiotemporal Transformer

arXiv: https://arxiv.org/abs/2206.04727 · HTML: https://mlsanity.com/paper/2206.04727/

## Key facts

- Cited by 23 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 67 references that are in the map.
- First posted on arXiv 9 Jun 2022.

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization](https://mlsanity.com/paper/1610.02391/) (2016)
- [Axiomatic Attribution for Deep Networks](https://mlsanity.com/paper/1703.01365/) (2017)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [Revealing the Dark Secrets of BERT](https://mlsanity.com/paper/1908.08593/) (2019)
- [Spatio-Temporal Graph Transformer Networks for Pedestrian Trajectory Prediction](https://mlsanity.com/paper/2005.08514/) (2020)
- [Visualizing and Measuring the Geometry of BERT](https://mlsanity.com/paper/1906.02715/) (2019)
- [A Spatio-temporal Transformer for 3D Human Motion Prediction](https://mlsanity.com/paper/2004.08692/) (2020)
- [Open Sesame: Getting Inside BERT's Linguistic Knowledge](https://mlsanity.com/paper/1906.01698/) (2019)
- [Spatial-Temporal Transformer for Dynamic Scene Graph Generation](https://mlsanity.com/paper/2107.12309/) (2021)
- [Neural Latents Benchmark '21: Evaluating latent variable models of neural population activity](https://mlsanity.com/paper/2109.04463/) (2021)
- [Representation learning for neural population activity with Neural Data Transformers](https://mlsanity.com/paper/2108.01210/) (2021)
- [Linear dynamical neural population models through nonlinear embeddings](https://mlsanity.com/paper/1605.08454/) (2016)
- [What does Attention in Neural Machine Translation Pay Attention to?](https://mlsanity.com/paper/1710.03348/) (2017)
- [Drop, Swap, and Generate: A Self-Supervised Approach for Generating Neural Activity](https://mlsanity.com/paper/2111.02338/) (2021)
- [NAST: Non-Autoregressive Spatial-Temporal Transformer for Time Series Forecasting](https://mlsanity.com/paper/2102.05624/) (2021)
- [Deep inference of latent dynamics with spatio-temporal super-resolution using selective backpropagation through time](https://mlsanity.com/paper/2111.00070/) (2021)

## Built on by (most cited first, 23 of 23)

- [A Unified, Scalable Framework for Neural Population Decoding](https://mlsanity.com/paper/2310.16046/) (2023)
- [Population Transformer: Learning Population-level Representations of Neural Activity](https://mlsanity.com/paper/2406.03044/) (2024)
- [The Brain's Bitter Lesson: Scaling Speech Decoding With Self-Supervised Learning](https://mlsanity.com/paper/2406.04328/) (2024)
- [Neural decoding from stereotactic EEG: accounting for electrode variability across subjects](https://mlsanity.com/paper/2411.10458/) (2024)
- [Generalizable, real-time neural decoding with hybrid state-space models](https://mlsanity.com/paper/2506.05320/) (2025)
- [Learning Time-Invariant Representations for Individual Neurons from Population Dynamics](https://mlsanity.com/paper/2311.02258/) (2023)
- [SPINT: Spatial Permutation-Invariant Neural Transformer for Consistent Intracortical Motor Decoding](https://mlsanity.com/paper/2507.08402/) (2025)
- [Extraction and Recovery of Spatio-Temporal Structure in Latent Dynamics Alignment with Diffusion Models](https://mlsanity.com/paper/2306.06138/) (2023)
- [BaRISTA: Brain Scale Informed Spatiotemporal Representation of Human Intracranial Neural Activity](https://mlsanity.com/paper/2512.12135/) (2025)
- [V1T: large-scale mouse V1 response prediction using a Vision Transformer](https://mlsanity.com/paper/2302.03023/) (2023)
- [BACE: Behavior-Adaptive Connectivity Estimation for Interpretable Graphs of Neural Dynamics](https://mlsanity.com/paper/2510.20831/) (2025)
- [BLEND: Behavior-guided Neural Population Dynamics Modeling via Privileged Knowledge Distillation](https://mlsanity.com/paper/2410.13872/) (2024)
- [Attention for Causal Relationship Discovery from Biological Neural Dynamics](https://mlsanity.com/paper/2311.06928/) (2023)
- [Neuroprobe: Evaluating Intracranial Brain Responses to Naturalistic Stimuli](https://mlsanity.com/paper/2509.21671/) (2025)
- [Cross-Modal Representational Knowledge Distillation for Enhanced Spike-Informed LFP Modeling](https://mlsanity.com/paper/2512.12461/) (2025)
- [Know Thyself by Knowing Others: Learning Neuron Identity from Population Context](https://mlsanity.com/paper/2512.01199/) (2025)
- [Dynamical modeling of nonlinear latent factors in multiscale neural activity with real-time inference](https://mlsanity.com/paper/2512.12462/) (2025)
- [JGAT: a joint spatio-temporal graph attention model for brain decoding](https://mlsanity.com/paper/2306.05286/) (2023)
- [Functional embeddings enable Aggregation of multi-area SEEG recordings over subjects and sessions](https://mlsanity.com/paper/2510.27090/) (2025)
- [When predict can also explain: few-shot prediction to select better neural latents](https://mlsanity.com/paper/2405.14425/) (2024)
- [Exploring Biological Neuronal Correlations with Quantum Generative Models](https://mlsanity.com/paper/2409.09125/) (2024)
- [DCHO: A Decomposition-Composition Framework for Predicting Higher-Order Brain Connectivity to Enhance Diverse Downstream Applications](https://mlsanity.com/paper/2509.09696/) (2025)
- [RPNT: Robust Pre-trained Neural Transformer -- A Pathway for Generalized Motor Decoding](https://mlsanity.com/paper/2601.17641/) (2026)

## Abstract (arXiv)

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.