STNDT: Modeling Neural Population Activity with Spatiotemporal Transformers
Authors: Trung Le
Eli Shlizerman
Venue: NeurIPS 2022 Accept (2022)
Keywords: neuroscience
systems neuroscience
computational neuroscience
neural population dynamics
brain-computer interfaces
neuroprosthetics
electrophysiology
neural coding
transformers

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: Modeling Neural Population Activity with Spatiotemporal Transformers
Trung Le University of Washington Seattle, WA tle45@uw.edd edu
Eli Shlizerman University of Washington Seattle, WA shlizee@uv.en edu
Abstract
Modeling neural population dynamics underlying noisy single-trial spiking activi- ties is essential for relating neural observation and behavior. A recent non-recurren 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 activ- ity while neglecting the rich covariation between individual neurons. In this paper introduce Spatio Temporl Neural Data Transformer (STNDT)), architecture that explicitly models responses of individual neurons in the population across time and space uncover their underlying firing rates. In addition, we propose a contrastive learning loss that works in accordance with mask modeling objective further improve the predictive performance. We show that our model achieves state-of-the-ae arr performance ensemble level in estimating neural activi- ties across four neural datasets, demonstrating its capability to capture autonomous and non-autonn aa--attoooouus dynamics spanning different cortical regions while being completely agnostic to the specific behaviors hand. Furthermore, STNDT spatial attention mechanism reveals consistently important subsets of neurons that play a vital role driving the response of the entire population, providing interpretability and key insights into how the population of neurons performs computation
NDT--based
1 Introduction
One of the most prominent questions in systems neuroscience is how neurons perform computations that give rise to behaviors. Recent evidence suggests that computation in the brain could governed at the population level 21 Population of neurons are proposed to obey internal dynamical rule that drives their activities over time
Inferring these dynamics on a single trial basis is crucial for understanding the relationship between neural population responses and behavior, subsequently enabling the development of robust decoding schemes with wide applicability in comptrr interfaces (BCT) (5 However, modeling population dynamics single trials is challenging due to the stochasticity of individual neurons making their spiking activity vary from trial to trial even when they subject to identical stimuli recorded under repeated behavior conditions.
direct approach to reduce the trial-to-trial variability neural responses could be to average responses over repeated trials of the same behavior [ 9. to convolve the neural response with Gaussian kernel [10], or general, to define a variety of neural activity measures [T. However, more success found in approaches that explicitly model neural responses dynamical system, including methods treating the population dynamics as being linear
12 [3], switched linear [14),
non-linear [15] [6], or reduced projected nonlinear models
Recent approaches leveraging
Code is available at http:://gthbb.com//sliiee////
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

recurrent neural networks (RNN) have shown promising progress in modeling distinct components of a dynamical system- neural latent states, initial conditions and external inputs on moment-
These sequential methods rely on continuous processing of neural inputs successive timesteps, causing latency that hampers applicability in real-time decoding of neural signals. Consequently to RNN-based approaches, Neural Data Transformer (NDT) [16] was proposed as a non-recurrent approach to improve inference speed by leveraging the transformers architecture which learns and predicts momentary inputs in parallel [19]. While successful, NDT has only focused on modeling the relationship neural population activity between timesteps while ignoring the rich covariation among individual neurons. Neurons in population have been shown to have heterogeneous tuning profiles where each neuron has a different level of preference to a particular muscle movement direction [20 21]. Neuron pairs also exhibit certain degree of correlation in terms of trial-to-trial variability (noise correlation) that affects the ability to decode the behaviors they represent 2 22]. These spatial correlations characterize the amount of information that can be encoded in the neural population [22]. necessitating the need to model the neural population activity
to-moment basis [5
17 18
across both time and space dimensions.
In this work, we propose to incorporate the information distributed along the spatial dimension to improve the learning of neural population dynamics, and introduce Spatio Neural Data Transformer, an architecture based on Neural Data Transformer which explicitly learns both the spatial covariation between individual neurons and the temporal progression of the entire neural population. We summarize our main contributions as follows:
We introduce STNDT which allows the transformer to learn both the spatial coordination between neurons and the temporal progression of the population activity by letting neurons attend each other while also attending over temporal instances.
We propose a contrastive training scheme, complementary to the mask modeling objective, to ensure the robustness of model prediction against induced noise augmentations.
We validate our model's performance on four neural datasets in the publicly available Neural Latents Benchmark suite [23] and show that ensemble variants of our model outperforms other sfate-of-the- the- methods, demonstrating its capability to model autonomous and non-autonomous neural dynamics in various brain regions while being agnostic to external behavior task structures.
We show that the spatial attention, a feature unique to STNDT, identifies consistently important subsets of neurons that play essential role in driving the response of the entire population. This exclusive attribute of STNDT provides interpretability and key insights into how the neural population distributes the computation workload among the neurons.
2 Related Work
Modeling spatial covariation in neural population: Neurons act an orchestrated system which collectively encodes behaviors in a distributed and redundant manner. Many previous works have studied and incorporated neural variability actoss neurons to closely match firing statistics observed in multi-channel neural recordings 30 (25) simulated population responses within a Dichotomized
Gaussian framework and solved for signal and noise correlations numerically. [26] 27 developed
Generative Adversarial Networks that were able to capture pairwise correlations among the neurons and generate realistic firing patterns. [28 modeled the population responses being generated from a latent variable with learnable covariance matrix reflecting covariability among the neurons.
While these methods resemble our work in the overarching motivation of capturing interactions among ncurons, they rely on the knowledge of the respective stimulicconditions that the trials belong to when modeling the interaction. On the other hand, STNDT is trained in unsupervised manner and learns the rich covariation among neurons encompassing all recorded behaviors without access to any external observation apart from the population spiking activity. In addition, while the goal of aforementioned methods is to generate realistic firing activities associated with induced stimuli, oftentimes with some assumptions regarding their statistics (e.g. noise correlation is shared across time bins and trials), STNDT aims to uncover the denoised firing patterns behind the noisy single-trial spiking activity and does not depend on any prior assumptions regarding their firing statistics.
Transformers for modeling spatiotemporal data: Transformers were initially developed to model the relationship between words in sentence, which can be thought of temporal progression of

---

Temporal Attention
Q
Temporal Features 2,
KT
(Tx
Spatiotemporal Features Zsr
p
(TxN)
Spatial Attention
Q5
X
K5
T)
(N N)
Figure I: Spatiotemporal Neural Data Transformer (STNDT) architecture. Separate multihead self-attention modules are trained to learn spatial covariation and temporal progression of neural activities. Temporal attention feature matrix is treated as the matrix V upon which spatial attention is multiplied to give the final spatiotemporal features. Colors represent entities over which self attention is performed. The complete STNDT consists of multiple layers of such spatiotemporal attention modules.
sequence of tokens. Recent works have leveraged the self-attention mechanism in transformers to model spatiotemporal data types where there exist an additional interacting dimensions possessing distinct dynamics, such as trajectories traffic agents [31+33], dynamic scene graph video [34], or 3D human motion [35]. However, in these works the spatial interaction at each timestep and the temporal dynamics for each entity are captured independently, treating the other dimension as the batch dimension at each attention block. In contrast, STNDT interleaves spatial and temporal attention in unified framework, using spatial attention to re- re-weight temporal features and enabling direct study of each individual neuron's 's role in driving the population dynamics.
Interpretability self-attention mechanism: Several approaches have been proposed to probe the inner workings of black-box deep learning models 38) Unlike our work, these approaches attempted to attribute importance of visual inputs to the model prediction did not take into account interaction between input features. For attention- based models, the weights of attention matrix have been used as a tool to provide certain level of interpretability (39]42). The interpretability is built upon the fact that attention weights signify how much influence other inputs have on a particular input in deciding its final outcome in a self-supervision manner. This influence might align with some human interpretable meaning, such as linguistic patterns (43]. In our work, further leverage attention weights to gain insights into the interaction of neurons from multi-channel neural recordings.
supervised setting and
3 Methods
Single-trial spiking activity of a neural population can be represented as spatiotemporal matrix X E NTNN, where each column X, E NT is the time series of one neuron, T the number of time bins for each trial, and N is the number of neurons in the population. Each element Xmn in the matrix is the number of action potentials (spikes) that neuron n fires within the time bin t. Spike counts
Problem formulation:
assumed to be samples of an inhomogeneous Poisson process P(A(t, n)) where X(t, n) is the underlying true firing rate of neuron n at time t. The matrix Y E RTTN containing X(t, n) fully represents the dynamics of the neural population and explains the observable spiking data of the respective trial. We propose to learn the mapping o(;W)) X + Y by the Spatiotemporal Transformer with the set of weights W.
Spatiotemporal Neural Data Transformer: the core of the transformer architecture is the multihead attention mechanism, where feature vectors learn to calibrate the influence of other feature
3

---

vectors in their transformation. Spike trains are embedded into feature matrices X with added sinusoidal positional encoding to preserve order information as initially proposed in [9]. We em- played separate embeddings to encode positions in each temporal and spatial dimension individually, resulting in two distinct feature embeddings Xy- = Emb(X) + Pr and Xs = Emb(X") ))+ Ps.
A set of three matrices WH, embedding X7 ( 22, . queries QT XTWA. keys KT W and values Vr XY upon which latent variable Z is computed as:
T W WY E RNXN are learned to transform T N-dimensional
@rKi (VN
Zr = (Qrr VT) F (softmax
Vr
VN
The outer product of QT KT represents the attention each pays to all other zy and determines how much influence their values vy have on its latent output 2,. Fís the sequence of concatenating multiple heads and feeding through feedforward network with ReLU activation [19]. heads for all reported models.
We used
Implementations of transformers in popular applications such as in natural language processing literature consider each feature vector 2; as an N-dimeesional token sequence, equivalent to word in sentence. Elements in the N v-ddimnnioona vector therefore serve as convenient numerical representation and do not have inherent relationships among them. The attention mechanism thus only models the relationship between tokens in sequence. In our application, each feature vector i;; a collection of firing activities of N physical neurons among which there exists an interrelation neuronal population acts coordinated structure with complex interdependencies rather than standalone individuals. We therefore propose to model both the temporal relationship the evolution of neural activities and the spatial relationship- covariability of neurons by learning two separate multihead attention blocks (Figure D. The temporal latent state Zy is computed with temporal attention block in Equation I In parallel, spatial attention block operates the spatial embedding Xs and learns an attention weights matrix signifying the relationship between neurons:
QsKs S
As softmax
VT
where Qs = Xx and Ks = X
This As matrix is then multiplied with the transpose of temporal latent state ZT to incorporate the influence of spatial attention on the final spatiotemporal latent state ZST:
ZSt = =F(s
(3)
For stable training, as in 19 we used layer normalization before Ny, Xs. As and feedforward layers. Residual connections are also employed around temporal attention, feedforward layers and A.
Mask modeling and contrastive losses: Similar [66. train the spatiotemporal transformer in an unsupervised way with BERT's mask modeling objective 44. During training, a random subset of spike bins along both spatial and temporal axes of input X are masked (zero-ed out or altered) and the transformer is asked to reconstruct the log firing rate at the masked bins such that the Poisson negative log likelihood is minimized:
T Lmaek ΣΣ
expl (zg) -iay
where Eig and žyy are the log output firing rate and input spike of neuron i at timestep j if location ij masked.
Neural dynamies are shown be embedded in a diw-dimensional space, i.e. model prediction should be fairly consistent when smaller subset of neurons are used compared to when the entire population taken into account. Furthermore, stereotyped behaviors often found neuroscience experiments, trials with the same condition should yield similar output firing rate profiles. Therefore, to enhance robustness of model prediction to neural firing variability we further constrain model firing rate outputs by contrastive loss, such that different augmentations of the same trial input remain closer

---

A
0.395
neuron 137
neuron 140
neuron 144
c0-Ops
50
200ms
200ms
200ms
ensemble size
B
observed
STNDT
200ms
200ms
200ms
Figure 2: A: co-bps metrics improves when multiple models are ensembled together. B: STNDT facilitates accurate inference of behavior from spiking data. Decoded hand trajectories from trials (dashed line) closely match the ground truth trajectories (solid line). C: STNDT uncovers the stereotyped feature of neural activity structured behaviors. Firing rate prediction and PSTHs of three example neurons are shown. Trials belonging to the same condition are plotted with the same color (4 trials per condition shown). All results are shown for MC_Maze dataset.
to each other and stay distant other trial inputs. We adopt the NT-XEnt contrastive loss introduced in 45
Coontrastive
Σ log
(5)
Joml 1xx:exp(sim(z,, 2x))
where sim(u, v) u (TTp([uu]]ll)) the cosine similarity between two predictions u and v on two different augmentations of input and is the temperature parameter.
Transformations such as dropping different views of neural data [46). In our work, we define the augmentation transformation as random dropout and alteration of spike counts at random elements in the original input matrix X, similar to how masking is done, i.e. zero out or change spike counts to random integers at random neurons and timesteps. See Appendix for details on probabilities used to create these augmentations.
neurons and jittering samples in time have been used create
We follow 47 to use Bayesian optimization for hyperparameters tuning. We observe that the primary metrics co-smoothing bits/spike (co-bps) are not well correlated with the mask loss (see Figure in the Appendix, while co-bps, vel R², psth R² and fi-bps are more pairwise correlated. Therefore, we run Bayesian optimization to optimize co-bps for M models then select the best N models as ranked by validation co-bps, and ensemble them by taking the mean of the predicted rates of these N models.
Bayesian hyperparameter tuning:
Experiments and results
Datasets and evaluation metries: We evaluate model performance four neural datasets in the publicly available Neural Latents Benchmark (23]: MC_Maze, MC_RTT, @eaa__umpp, and DMFC RSG. The 4 datasets cover autonomous and non-autonomous neural population dynamics recorded on rhesus macaques in a variety of behavioral tasks (delayed reaching, self-paced reaching, reaching with perturbation, time interval reproduction) spanning multiple brain regions (primary motor cortex, dorsal premotor cortex, somatosensory cortex, dorso-medial frontal cortex). The diverse scenarios and systems offer comprehensive evaluation of a latent variable model and serve as standardized benchmark for comparison between different modeling approaches. We use different metrics to measure performance
our model depending on the particular behavior task of each dataset, following the standard evaluation pipeline in [23 We evaluate and report our model performance on the hidden test split held by NL.B to have a fair comparison with other otate-of-the-art (SOTA) methods. See 23 for further details of evaluation strategy and how the metrics are calculated.
5

---

Table 1: Performance of STNDT compared to SOTA methods
Maze and MC_RTT
RTT datasets
MC Maze
MC RTT
Methods
co-bpst
vel R²t
psth R²T
fp-bpst
co-bpst
vel R²1
fp-bpst
GPFA
0.1872
0.6399
0.5150
0.1548
0.5339
Smoothing
0.2109
0.6238
0.1853
0.1468
0.4142
SL.DS
0.2249
0.7947
0.5330
1.1579
0.1649
0.5206
0.0620
MINT
0.3304
0.9121
0.7496
0.2076
0.1676
0.5953
0.1012
AutoLFADS
0.3364
0.9097
0,6360
0.23499
0.1868
0.6167
0.1213
iLQR-VAL VAE
0.3559
0.8840
0.6062
0,1480
AESMTEI (single)
0.3599
0.9105
0.6641
0.2470
0.1927
0.6627
0.1229
AESMTE3 (ensemble)
0.3676
0.9114
0.6683
0,2589
0.2053
0.6334
0.1344
STNDT single (ours)
0.3691
0.8985
0.6567
0.2505
0.1938
0.6143
0.0988
STNDT ensemble (ours)
0.3862
0.9095
0.6693
0.2686
0.2095
0.6270
0.1244
Table 2: Performance of STNDT as compared to SOTA methods on Area2 Bump and DMFC RSG
datasets
Area2 Bump
DMFC_RSS RSG
Methods CO- bps
vel R²t
psth R²t
fp- bpst
CO- bps
tp-corr]
psth R²t
fp-bpst
GPFA
0.1680
0.5975
0.5289
0.1176
-0.3763
0,2142
Smoothing
0.1544
0.5736
0.2084
0.1202
-0.5139
0,2993
SLDS
0.1960
0.7385
0.5740
0.0242
0.1243
-0.5412
0.3372
-0.0418
MINT
0.2735
0.8877
0.9135
0.1483
0.1821
-0.6929
0. 0.7013
0.1650
AutoLFADS
0.2569
0.8492
0.6318
0.1505
0.1829
-0.8248
0.6359
0.1844
iLQR-VAE
AESMTEI (single)
0.2801
0.8675
0.6367
0.1523
0.1733
0.6189
0.5267
0.1511
AESMTE3 (ensemble)
0.2860
0.8999
0.7109
0.1603
0.1886
0.7601
0.6064
0.1828
STNDT single (ours)
0.2818
0.8766
0.6454
0.1357
0.1859
-0.2055
0.6051
0.1601
STNDT ensemble (ours)
0.2898
0.8913
0.7368
0.1476
0.1940
-0.8857
0.6452
0.1910
Co-smoothing (co-bps): the primary metric, measuring the ability of the model to predict activity of held-out neurons it has not seen during training. Co-bps is tied the goodness of mask loss evaluated for held-out neurons.
Behavior decoding (vel R or tp-corr): measures how useful the model firing rates prediction can be used to decode behavior (the velocity of primate's hand the cases of MC_Maze and Areas_ Bump datasets, the correlation between neural speed and time between Set cue and Go response DMFC_RSC dataset).
indicates how well predicted firing rates match the peri-stimuls time histogram in repeated, stereotyped task structures.
Match to peri-stimulus time histogram (psth R²)::
Forward prediction (fp-bps): measures model's 's ability to predict unseen future activity of the neural population. It is computed in the similar manner as co-bps but on the held-out time points of all neurons.
We compare STNDT against the following baselines, all of which have been evaluated using the same held-out test split.
Baselines:
A simple method where Gaussian kernel is convolved with held-in spikes to produce smoothed held-in firing rates. Then Poisson Generalized Linear Model (Poisson GLM) is fitted from the held-in smoothed rates held-out rates.
Smoothing 23
extracts population latent states as a smooth and low dimensional evolution by combining smoothing and dimension reduction in common probabilistic framework.
GPFA [0]:
models neural dynamics as a switching linear dynamical system, which breaks down nonlinear data into sequences of simpler dynamical modes.
SLDS 14:

---

models population activity as non- non-linear dynamical system with bi- bi-directional recurrent neural networks at the core and scalable framework of hyperparameter tuning.
AutoLFADS [77]
MINT 48 an interpretable decode algorithm that exploits the sparsity and stereotypy of neural activity to interpolate neural states using a library of canonical neural trajectories.
iLQQR-VA 49 improves upon LFADS with iterative linear quadratic regulator algorithm, an optimization-baser recognition model to replace RNN the inference network.
NDT [6] leverages transformer architecture with some adaption neural data model temporal progression neural activity across time. AESMTEI is the best single model and AESMTE3 is the best emsemble of multiple models found result of Bayesian hyperparameter tuning 47
4.1 Spatiotemporal transformer achieves state-of-the-art performance in modeling autonomous dynamics
We first tested STNDT on recordings of dorsal premotor (PMd) and motor cortex (M1) of monkey performing a delayed reaching task (MC_Mar Maze dataset) to evaluate the ability of STNDT to uncover single-trial population dynamics in highly structured behavior. The dataset has been studied extensively in previous work [15 It consists of 2869 trials of monkey performing a center-out reaching task in a maze with obstructing barriers, composing 108 different conditions for straight and curved reaching trajectories. The monkey is trained to hold the cursor at the center while the target is presented and only move the cursor to reach the target after a 'Go' cue. The neural dynamics during the preparation and execution periods is well modeled as an autonomous dynamical system [1].
We observed that by explicitly modeling spatial interaction, STNDT outperformed other state-of-the. art methods and improved NDT's ability to model autonomous single-tring trial dynamics as measured by the negative log likelihood of unobserved neural activity. The single STNDT model improved both Poisson log likelihood of heldout neurons (co-bps) and heldout timesteps (fp-bps). The performance is further increased by aggregating multiple STNDT models as shown in Table I and Figure A
Since MC_Maze features repeated trials, the prediction of any latent variable models should uncover stereotypical patterns of neuronal responses for trials belonging to the same condition. Therefore, we computed PSTH which is the average of neural population response across trials of the same condition, and measure R matching of model prediction to this PSTH. We observed that with the help of spatial modeling and contrastive loss, STNDT boosts NDT ability to recover this stereotyped firing pattern] [I] We show in Figure 2 several responses of example neurons. STNDT firing rates prediction of trials under the same condition exhibit a consistent, stable PSTH as desired. These predicted rates also decode behaviors accurately when mapped to hand velocity via linear regression model (Table [ Figure 2)..
4.2 Spatiotemporal transformer improves inference of non-autonomous neural dynamics underlying naturalistic behaviors
There is much interest in systems neuroscience to study neural dynamics in unconstrained, naturalistic behaviors it is crucial for developing ubiquitous BCI decoders. We evaluated STNDT's applicability to this setting via recordings in primary motor cortex during self- self-paced reaching task IMC_RTV dataset) 23 50 Unlike MC [C__ dataset, the monkey in this task continuously acquires targets which appear randomly in an 8x8 grid without preparatory periods, resulted in a wide variety of hand trajectories and trial lengths. We observe that STNDT achieves SOTA performance on the primary metric co- and performs on par with NDT on remaining metrics, while maintaining a more robust performance against random initializations of model weights (Table and Appendix).
4.3 Spatiotemporal transformer better captures input-driven dynamics underlying sensory processes
We next tested STNDT in a setting where unexpected input perturbations affect the neural dynamics in somatosensory cortex to probe whether STNDT can leverage spatial interaction to improve modeling of non-autonomous dynamics in this brain region. Area2_ reaa_Bump dataset consists of recordings from the Area 2, which was shown in previous works to be driven by mechanical perturbation to the arm and contains information about whole-arm kinematics 23 [5). The task comprises of active and passive trials with center hold period at the start. During active trials, the monkey performs classic
7

---

center-out reaching task. In passive trials, force is applied on the monkey's ' hand in random direction via a manipulandum. after which the monkey has to return the center target and proceed with the task as in active trials. Despite the relatively small scale of the dataset, STNDT brings about further improvements to NDT performance in terms of co-bps and psth- R2, on both single and ensemble levels.
4.4 Spatiotemporal transformer enhances prediction of neural population activity during cognitive task
Dorsomedial frontal cortex (DMFC) is believed to serve as an intermediate layer between low-level sensory and motor areas, and possess distinct confluence of internal dynamics and inputs 52 53 We are therefore interested to see if characterizing spatial relationship alongside temporal relationship and incorporating contrastive loss could help STNDT better model the dynamics in this brain region. We tested STNDT on the DMFC IMFC_RSG dataset [23 53 consisting of recordings from a rhesus macaque performing a time-interval reproduction task. The monkey is presented two 'Ready" and 'Set' stimuli separated by a specific time interval l, while fixating eye and hold the joystick at the center position. It then has to execute 'Go" response by either an eye saccade or joystick movement such that the time interval fo between its reponse and the 'Set' cue is sufficiently close to t,. STNDT successfully captures the dynamics in this cognitive task, outperforming NDT by large margin across co-bps, psth- R²2 and fp-bps on both single and ensemble level (Table 2).
4.5 Spatial attention mechanism identifies important subsets of neurons driving the population dynamics
In Figure 3 we visualize spatial attention weights obtained from STNDT on the MC_ dataset in the first and last attention layers. Attention map for remaining datasets are provided in Appendix. Interestingly, spatial attention shows that in early layers, only a small subsets of neurons in the population are consistently attended to by all neurons. The spatial attention tends to disperse as the model goes to deeper layers. Strikingly, the subset of geavily-attender neurons stays relatively identical across different trials, hinting that these neurons might play a crucial role in driving the population response to the behavior task. We further tested this hypothesis by incrementally dropping the neurons heavily attended to (i.e. zeroing out their spiking activity input to the model) in a descending order of their attention weights identified in the first layer. We observed that dropping these important neurons identified by STNDT caused a significant decline in the model performance (Figure 4. The performance decline was sigmificantly more than the case where the same number of random neurons are dropped. To rule out the possible case that dropping neurons only has adverse effect on the spatial attention module but that effect propagates to the subsequent modules and indirectly impacts the performance of the overall STNDT pipeline, we repeated the experiment on the vanilla NDT model which, unlike STNDT, lacks a spatial attention structure. Interestingly. we observed the same performance deterioration when we dropped the spiking activity of STNDT- identified important neurons and asked a pretrained vanilla NDT to make inference on the resulting inputs. This finding suggests that the impact of the important neurons that only STNDT can identify might potentially generalize to other latent variable models that without input from these neurons, some latent variable models might not function optimally. We provide additional results from similar analyses on GPFA and Smoothing models in the Appendix.
We further examine whether important neurons were selected by the spatial attention mechanism based on some criteria more sophisticated than simple firing statistics, as more active neurons tend to have higher signal-to-noise ratio and might encode more useful information with regard to behaviors. We find that the important neurons are not the ones with the highest spike counts or the least variability in spiking activity. In fact, attention weights of a neuron do not correlate or only correlate weakly to its firing activity statistics, as we show in 3 the Pearson's correlation of a neuron's 's attention weight with the mean and variance of its spiking activity. All correlation values have p-value < le-4. These results indicate that STNDT's spatial attention has picked up on meaningful population features that are more significant than firing statistics of the neurons.
8

---

trial
trial
trial
trial
320-22
layer
160
100
160
100 320 380 280
53 1103200 140:16
100 1201404400000 38000000000
20
100 140 140
20-
-44
80-
100
100
333
100
352
Figure 3: Visualization of STNDT's spatial attention weights in the first and last layers of four example trials. Attention weights in layer reveal consistent subset of neurons that are heavily attended by all neurons the population. The attention becomes more dispersed in deeper layers. Results are shown for 182 neurons in MC_Maze dataset.
0.4
0.8
co-ops
0.0
vel R2
-0.2
Important, STND
0.6
Important,
randore, NND NDT random,
rendorn, SIND
0.5
importntt,
07%
10%
15%
20%
25%
30%
5%
10%
15%
20%
25%
30%
% neurons dropped
% neurons dropped
0.2
-2
PSTH R2
bps
0.0
STND
randon,
STNDT
-10-
randor,
important, NDT
-0.2
important, ND
056
Sis
200 25% 76
10te
15%
S/
10%
15%
20%
25%
neurons dropped
96 neurons dropped
Figure 4: Spatial attention module, unique to STNDT, identifies important neurons that are the main driving force of population response to behavioral task. Performance of STNDT as measured by four evaluation metrics are plotted neurons are incrementally dropped from input neural population. Performance significantly deteriorates when important neurons identified by STNDT are dropped, while only decreases slightly when random neurons dropped. The effect of important neurons indentified by STNDT generalizes to vanilla NDT, which lacks a spatial attention structure. Shaded region represents 2 standard error of the mean. Results are shown for Maze dataset.
Table 3: Pearson's 's correlation between spatial attention weight of a neuron versus mean and variance of its spiking activity.
MC Maze
MC RIT
Area2 Bump
DMFC RSG
pspike mean, attn weight) Alpike var, attn weight)
0.0164
0.2217 0.2189
0.0327 0.0353
0.0852 0.0937
0.0124

---

Table 4: Ablation Study:
Performance of STNDT MC_ Maze and MC_ RTT datasets with and without contrastive loss (CL) on single and ensemble levels.
Maze
MC_RTT
Methods
co-bpst
vel R²t
psth R²1
fp-bpst
co-bpst
vel R²*
fp-bpst
AESMTEI (single)
0.3599
0.9105
0.6641
0.2470
0.1927
0.6627
0.1229
AESMTE3 (ensemble)
0.3676
0.9114
0.6683
0,2589
0.2053
0.6334
0.1344
STNDT single w/o CT.
0.3668
0.8979
0.6549
0.2471
0.1865
0.5988
0.0964
STNDT single
0.3691
0.8985
0.6567
0,2505
0.1938
0,6143
0.0988
STNDT ensemble w/o CL
0.3843
0.9090
0.6686
0.2675
0.2065
0.6352
0,12
STNDT ensemble w/ CL
0.3862
0.9095
0.6693
0.2686
0.2095
0.6270
0.1244
Table 5:
Ablation Study:
Performance of STNDT on Area2 Bump and DMFC RSG datasets with
and without contrastive loss (CL) on single and ensemble levels.
Area2 Bump
DMFC RSG
Methods
CO-
vel R
psth 1t"t
fp- bpst
CD- bpst
Ip- com]
psth Rt*
fp-
bps*
bpst
AESMTEI (single)
0.2801
0.8675
0.6367
0.1523
0.1733
-0.6189
0.5267
0.1511
AESMTE3 (ensemble)
0.2860
0.8999
0.7109
0.1603
0.1886
-0.7601
0.6064
0.1828
STNDT single w/o CL
0.2765
0.8773
0.7169
0.1498
0.1824
0.5059
0.6134
0.1473
STNDT single CL
0.2818
0.8766
0.6455
0.1357
0.1859
0..200
0.6051
0.1601
STNDT ensemble w/o CL
0,2904
0.8937
0.7303
0.1491
0.1931
-0.5186
0. 0.6429
0. 18188
STNDT ensemble w/ CL
0.2898
0.8913
0.7368
0.1476
0.1940
-0.4857
0.6452
0.1910
4.6 Ablation Study: Contrastive loss encourages consistency of model prediction and improves performance
We conduct an ablation study to assess the effectiveness of contrastive loss on the overall performance of STNDT. Tables 4 and repport how the model scores on different metrics across all four datasets on the single and ensemble levels. In general, we observe that having contrastive loss further improves the performance of STNDT on predicting neural activity of heldout neurons (co-bps) and heldout timesteps (fp-bps). The contribution of contrastive loss is most eminent on Ma_e dataset.
5 Discussion
In this paper we presented STNDT, novel architecture based upon NDT [16] that explicitly learns the covariation among individual neurons in the population alongside the momentary evolution of the population spiking activity in order to infer the underlying firing rates behind highly variable single-trial spike trains. By incorporating self-attention along both spatial and temporal dimensions as well as a contrastive loss, STNDT enhances NDT's ability to model dynamics spanning a variety of tasks and brain regions, most notably by the accurate prediction of activity for unseen neurons (co-bps). Although STNDT does not consistently outperform NDT on other secondary metrics, we show in the Appendix that STNDT is more robust random initializations and performs better than NDT on average across random seeds. Moreover, the improvement STNDT contributes on co-bps is the direct reflection of the spatial attention's 's success. Since the spatial attention module aims to learn the relationship between all (observed and unobserved) neurons at training time, it will leverage this information to infer activities of unobserved neurons based those of observed neurons at testing time, which is exactly what co-bps measures. Finally, the novel spatial attention mechanism unique STNDT brings about valuable interpretability as discovers influential subsets of neurons whose activities contain salient information about the response of the entire neural population without which some latent variable models might not function optimally.
Acknowledgment: This work was supported in part by National Science Foundation grant OAC-2117997 and Washington Research Fund to ES, Authors also acknowledge the partial support by the Departments of Electrical Computer Engineering (TL and ES). Applied Mathematics (ES), the Center of Computational Neuroscience (ES), and the eScience Center (ES) at the University of Washington.
10

---

References
[1] Rafael Yuste. From the neuron doctrine neural networks. Nature reviews neuroscience, 16(8):487-497. 2015.
[2] Shreya Saxena and John P Cunningham. Towards the neural population doctrine. neurohiology, 55: 103-111, 2019.
Current opinion in
[3] Krishna V Shenoy, Mancesh Sahani, and Mark M Churchland. Cortical control of arm movements: dynamical systems perspective. Anmual review of neuroscience, 36:337-359, 2013.
[4] Krishna V Shenoy and Jonathan C Kan. Measurement, manipulation and modeling of brain-wide neural population dynamics. Nature Communications, 12(1):1-5, 2021.
[5] Francis R Willett, Donald T Avansino, Leigh R Hochberg. Jaimie M Henderson, and Krishna V Shenoy. High-performance brain-to-text communication via handwriting. Nature, 593(7858):299-24, 2021.
[6] Jennifer L Collinger, Brian Wodlinger, John E Downey, Wei Wang, Elizabeth C Tyler- Kabara, Douglas J Weber, Angus JC McMorland, Meel Velliste, Michael L Boninger, and Andrew B Schwartz. High- performance neuroprosthetic control by an individual with tetraplegia. The Lancet, 38 19966)):55577664 2013.
[7] Beata Jarosiewicz, Anish Sarma, Daniel Bacher. Nicolas Y Masse, John Simeral, Brittany Sorice, Erin M Oakley, Christine Blabe, Chethan Pandarinath, Vikash Gilja, et al. Virtual typing by people with tetraplegia using self-calibrating intracortical brain-computer interface. Science translational medicine. 3 37-3133ral111111 319 2015.
[8] Rafael Levi, Pablo Varona, Yuri Arshavsky, Mikhail I Rabinovich, and Allen I Selverston. The role of sensory network dynamics in generating motor program. Journal of Neuroscience, 25(42):9807 9815, 2005.
[9] Miguel AL Nicolelis, Luiz A Baccala, Rick CS Lin, and John K Chapin. Sensorimotor encoding by syn- chronous neural ensemble activity multiple levels of the somatosensory system. Science, Science, 15555 1358, 1995.
[10] Byron M Yu, John P Cunningham, Gopal Santhanam, Stephen Ryu, Krishna V Shenoy, and Mancesh Sahani.
factor analysis for low-dimensional single-trial analysis of neural population activity. Advances in neural information processing systems, 21, 2008.
[11] Shlizerman, Konrad Schroder, and Nathan Kutz. Neural activity measures and their dynamics. SIAM Journal on Applied Mathematics, 72(4): (2(4):126-1125) 12910119112220011 2012.
[12] Jomathan Kao, Paul Nuyujukian. Stephen Ryu, Mark Churchland, John P Cunningham, and Krishna Shenoy. Single-trias trial dynamics of motor cortex and their applications to brain- main-mechine interfaces. Nature communications, 6(1):11 12, 2015.
[13] Yuanjun Gao, Evan W Archer. Liam Paninski, and John P Cunningham, Linear dynamical neural population models through nonlincar embeddings. Advances in neural information processing systems, 29, 2016.
[14] Scott Linderman, Matthew Johnson, Andrew Miller, Ryan Adams, David Blei, and Liam Paninski. Bayesian learning and inference in recurrent switching linear dynamical systems. In Artificial Intelligence and Statistics, pages 914-922, PMLR, 2017.
[15] Chethan Pandarinath, Daniel J O' O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew T Kaufman, Stephen I Ryu, Leigh R Hochberg, et al. Inferring sinale-trial neural population dynamics using sequential auto-encoders. Nature methods, 15(10)800-8815 2018.
[16] Joel Ye and Chethan Pandarinath. Representation learning for neural population activity with neural data transformers. arXiv preprint 1.Xi::2008.0120000 2021.
[17] Mohammad Reza Keshtkaran, Andrew R Sedler, Raeed H Chowdhury, Raghav Tandon, Diya Basrai, Sarah I. Nguyen, Hansem Sohn, Mchrdad Jazayeri, Lee E. Miller, and Chethan Pandarinath. A large-scale neural network training framework for generalized estimation of single-trial population dynamics. bioRxiv. 2021.
[18] Feng Zhu, Andrew Sedler, Harrison Grier, Nauman Ahad, Mark Davenport, Matthew Kaufmann Andrea Giovannucei, and Chethan Pandarinath. Deep inference of latent dynamics with spatio- temiooa super- resolution using selective backpropagation through time. Advances in Neural Information Processing Systems, 34, 2021.
11

---

[19] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
[20] Margaret Yvnne Mahan and Apostolos P Georgopoulos. Motor directional tuning across brain areas: directional resonance and the role of inhibition for directional accuracy. Frontiers in neural circuits, 7:92, 2013.
[21] Adam Kohn, Ruben Coen-Cagli, Ingmar Kanitscheider, and Alexandre Pouget. Correlations and neuronal population information, Annual review of neuroscience, 39:237-256, 2016.
[22] Bruno B Averbeck, Peter E Latham, and Alexandre Pouget. Neural conrelations, population coding and computation. Nature reviews neuroscience, 7(5):358-366, 2006.
[23] Felix Pei. Joel Ye, David Zoltowski, Anqi Wu, Raeed H. Chowdhury. Hansem Sohn, Joseph E O'Doherty, Krishna Shenoy, Matthew Kaufiman, Mark Churchland, et al. Neural latents benchmark'2 21: Evaluating latent variable models of neural population activity. arXiv preprint arXiv:2109 2021.
[24] Elad Schneidman, Michael Berry, Ronen Segev, and William Bialek. Weak pairwise correlations imply strongly correlated network states in a neural population. Nature, (10007087):1007-1100 2006.
[25] Dmitry Lyamzin, Jakob Macke, and Nicholas A Lesica. Modeling population spike trains with specified time-varying spike rates, trial-to-trial variability, and pairwise signal and noise correlations. Frontiers in computational neuroxcience, 4: 4:144, 2010.
[26] Manuel Molano-Mazon Arno Onken, Eugenio Piasini, and Stefano Panzeri. Synthesizing realistic neural population activity patterns using generative adversarial networks. arXiv preprint arXiv: 003388 2018.
[27] Poornima Ramesh, Mohamad Atayi, and Jakob H Macke. Adversarial training of neural encoding models on population spike trains. 2019.
[28] Jakob H Macke, Lars Buesing, John Cunningham, Byron M Yu, Krishna V Shenoy, and Mancesh Sahani. Empirical models of spiking in neural populations. Advances in neural information processing systems, 24, 2011.
[29] Stephen Keeley, Mikio Aoi, Yiyi Yu, Spencer Smith, and Jonathan W Pillow. Identifying signal and noise structure in neural population activity with gaussian process factor models. Advances in Neural Information Procensing Systems, 33:13795-13805, 2020.
[30] Mohammad Bashiri, Edgar Walker, Konstantin- Manstantin-Klemmns Lurz, Akshay Jagadish, Taliah Muhammad, Zhiwei Ding. Zhuokun Ding, Andreas Tolias, and Fabian Sinz. A flow-based latent state generative model of neural population responses to natural images. Advances in Neural Information Processing Systems, 34:15801-15815 2021.
[31] Weihuang Chen, Fangfang Wang, and Hongbin Sun. S2met: Spatio-temport transformer networks for trajectory prediction autonomous driving. Asian Conference on Machine Learning, pages PMIRR, 2021.
[32] Cunjun Yu, Xiao Ma, Jiawei Ren, Haiyu Zhao, and Shuai Yi. Spatio-temporal graph transformer networks pedestrian trajectory prediction. In European Conference on Computer Vision, pages 507-523. Springer. 2020.
[33] Kai Chen, Guang Chen, Dan Xu, Lijun Zhang, Yuyao Huang, and Alois Knoll. Nast: non- autorrrrrreeiev spatial-temporal transformer for time series forecasting, arXiv preprint vrXiv:2102.05022 2021.
[34] Yuren Cong, Wentong Liao, Hanno Ackermann, Bodo Rosenhahn, and Michael Ying Yang. Spatial- temporal transformer for dynamic scene graph generation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 16372-16382, 2021.
[35] Emre Aksan, Manuel Kaufmann, Peng Cao, and Otmar Hilliges. A spatio-temporal transformer for 3d human motion prediction. In 2021 International Conference on 3D Vision (3DV), pages 565-574. IEBE, 2021.
[36] Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based localization. In Proceedings of the IEEE international conference on computer vision, pages 618-626, 2017.
[37] Mukund Sundararajan, Ankur Taly, and Qiqi Yan. Axiomatic attribution for deep networks. In International conference machine learning, pages 3319-3322 PMLR, 2017.
12

---

[38] Niru Maheswaranathan, Lane T McIntosh, David B Kastner, Josh Melander, Luke Brezovec, Aran Nayebi, Julia Wang, Surya Ganguli, Stephen A Baccus, et al. Deep learning models reveal internal structure and diverse computations in the retina under natural scenes. biorxiv. URL: http://www. biorxiv. wgg/oonenneeatyy200/////////4440 hup://dx, doi. orgglomteet/erryyy001/0////////0 full. pdf. 2018.
arXiv: http:://ww. biorxiv.
[39] Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D) Manning. What does bert look at? an analysis of bert's 's attention. arXiv preprint arXiv: 04341, 2019.
[40] Olga Kovaleva, Alexey Romanov, Anna Rogers, and Anna Rumshisky. Revealing the dark secrets of bert. arXiv preprint arXiv:1908.088: 085933, 2019.
[41] Yongjie Lin, Yi Chem Tan, and Robert Frank. Open sesame: getting inside bert's linguistic knowledge. arXiv preprint arXiv:1906.0160 2019.
[42] Hamidreza Ghader and Christof Monz. What does attention in neural machine translation pay attention to? arXiv preprint 17X1v:17/0.03348, 2017.
[43] Emily Reif, Ann Yuan, Martin Wattenberg, Fernanda B Viegas, Andy Coenen, Adam Pearce, and Been Kim. Visualizing and measuring the geometry of bert. Advances Neural Information Processing Systems. 32, 2019.
[44] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirec- tional transformers for language understanding. arXiv preprint arXiv: 1810.04805, 2018.
[45] Ting Chen, Simon Komblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607. PMLR, 2020.
[46] Ran Liu, Mehdi Azabou, Max Dabagia, Chi-Heng Lin, Mohammad Cheshlaghi Azar, Keith Hengen, Michal Valko, and Eva Dyer. Drop, swap, and generate: A self-supervised approach for generating neural activity. Advances in Neural Information Processing Systems, 34:10587-10599, 2021.
[47] Darin Sleiter, Joshua Schoenfield, and Mike Vaiana. ac-nIb-2021.
httpa://githun
Satpyeattrrrraaa/a-nn------- 2021. 2021.
[48] Sean Perkins. Mint: Mesh of idealized neural trajectories. httpss//giituub..cm/uuuuuaaaaaat nlb_ wwr.sshop/blo/mminn.n MMINT. pdf 2022.
[49] Marine Schimel, Ta-Chu Kao, Kristopher Jensen, and Guillaume Hennequin. ilqr-vae: control-based learning of input-driven dynamics with applications to neural data. bioRxiv, 2021.
[50] Joseph G Makin, Joseph E O' Doherty, Mariana MB Cardoso, and Philip N Sabes. Superior arm-movement decoding from cortex with a new, unsupervised-lcaming algorithm. Journal of neural engineering, 15())020000, 2018
[51] Raeed H Chowdhury, Joshua Glaser, and Lee E Miller, Area 2 of primary somatosensory cortex encodes kinematics of the whole arm. Elife, 2020.
[52] Mattia Rigotti, Omri Barak, Mclissa R Warden, Xiao-Jing Wang, Nathaniel D Daw, Earl K Miller, and Stefano Fusi. The importance mixed selectivity complex cognitive tasks. Nature, 497(7451):585-355 590, 2013.
[53] Hansem Sohn, Devika Narain, Nicolas Meirhaeghe, and Mehrdad Jazayeri. Bayesian computation through cortical latent dynamics. Neuron, 103(5):934 (145)):34-947,0000 2019.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contribu- tions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] Please see Section S
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
13

---

(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? Yes]] Please see Appendix.
(c) Did you report enror hars (e.g.. with respect to the random seed after running experiments multiple times)? [Yes] Please see Appendix.
(d) Did you include the total amount of compute and the type of resources used (e.g.. type of GPUs, internal cluster, cloud provider)? [Yes] Please Appendix.
4. If you are using existing assets (c.g., code, data, models) or curating/relcasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A] The data and code provided the Neural
Latent Benchmark are unlicensed
(c) Did you include any assets either in the supplemental material URL? [Yes] The code included as URL.
(d) Did you discuss whether and how consent was obtained from people whose data you're us- ing/curating? [Yes]
(e) Did you discuss whether the data you are using/curating contains personally identifiable informa- tion or offensive content? [N/A] Personally identifiable information and offensive content docs not apply to the neurophysiologica data from rhesus macaques.
If you used crowdsoureing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "mp6qIcpo9q", "reviewer": "Reviewer_uE5F", "summary": "The paper introduces a spatio-temporal neural transformers: a transformer-based generative model of neural population activity that captures correlations between neurons, in addition to stimulus-driven and temporal correlations in the activity.\n", "questions": "1) It would be good to have a discussion of methods that have been developed to capture noise correlations in neural population acitivity. It would also be good to have a comparison of STNDT to at least one of the methods as well: e.g. Lyamzin et. al. (2010).\n\n2) It would be good to have a more thorough (or at least more measured) discussion of the interpretability of attention weights. Statements as in line 222: \"The interpretability...final outcome.\";  in line 241 \"This finding suggests...not function optimally\" and in line 262 \"Finally, the novel...not function optimally\" may not be warranted when similar analyses on other network architectures exist, with arguably similar results. \nIt would also be good to have some further analysis on the subset of \"important\" neurons identified by the attention weights: for example, are these neurons with the highest firing rates, or the least variability in firing rates?\n\n3)  It would be good to have information about training costs, computational and sample efficiency of the menthod\n\n4) Performance of the different methods are compared based on several metrics: it would be good to have some intuition on why these particular metrics were chosen, and also how these are computed.\n\nMinor comments:\n1) Figures 1-4 have grey lines around the border that don't nescessarily overlap with the division of subpanels\n2) The colourbars in Figure 4 have no label, and the colours are hard to distinguish against the dark background in the plots", "limitations": "Overall, although the paper presents an interesting use case of transformers, the value they add over existing methods is not clear. In particular, it is not clear whether the transformers truly are beneficial without a comparison to other methods explicitly set up to capture noise correlations in neural population data. Furthermore, this paper argues for the benefits of STNDT based on the interpretability of the attention weights: however the analysis does not convincingly show whether (a) the attention weights truly pick up on data features more significant than say, the firing rates of the neurons (b) whether they are any more interpretable than weights obtained from linear regression or attribution methods. Finally, without information about the computational costs of this method, it is hard to judge its value, even if it outperforms other methods on the Neural Latents Benchmark tasks.\n\nTaken together, there needs to be more analysis to show whether STNDT can provide more scientific insight and add more value to generative modeling of neural population data, over existing methods.\n", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall the paper is well-written, the methods and results are presented clearly. In particular, this seems to be an interesting use-case for transformers, and the ablation study with the contrastive loss is also interesting.\n\nHowever, there are some major concerns regarding the paper:\n1) There is a substantial body of work concerning capturing correlations between neurons (many of which also concurrently capture stimulus-driven /  temporal variability in the population activity -- it would be good to have a discussion of these papers, and also, comparisons of STNDT against some of them. A non-exhaustive list includes:\n- Schneidman et al (2006): https://europepmc.org/backend/ptpmcrender.fcgi?accid=PMC1785327&blobtype=pdf\n- Macke et al. (2011): http://www.gatsby.ucl.ac.uk/~maneesh/papers/macke-etal-2011-nips-preprint.pdf\n- Lyamzin et al. (2010): http://journal.frontiersin.org/article/10.3389/fncom.2010.00144/abstract\n- Molano-Mazon et al. (2018): http://arxiv.org/abs/1803.00338\n- Ramesh et al (2019): https://openreview.net/forum?id=S1xxRoLKLH\n- Bashiri et al. (2021): https://openreview.net/pdf?id=1yeYYtLqq7K\n\n2) While STNDT certainly achieves better performance compared to other methods, it is not clear what scientific insight can be gained from using a transformer. While section 3.3 and 3.4 describe how attention weights can be used to find sub-networks of \"important\" neurons and model consistency, it is not clear whether these weights are any more interpretable than say, regression weights, or those computed on CNNs with attribution methods. Indeed, attribution methods such as GradCAM (https://arxiv.org/abs/1610.02391) and axiomatic attribution (https://arxiv.org/abs/1703.01365) have been used in conjunction with other deep neural network architectures trained on neural population activity in precisely the same manner described here, with arguably similar success (e.g. Maheswaranathan et al. 2018: https://www.biorxiv.org/content/early/2018/06/08/340943).\n\n3) It would also be nice to have an estimate of the training costs, computational and sample efficiency of STNDT in comparison with other methods -- otherwise it is hard to truly judge the benefits of using this method, over others for datasets beyond those described in the paper.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "f4hvhwHxOs", "reviewer": "Reviewer_457j", "summary": "This paper presents a transformer architecture for modelling neural spike trains. The model considers both the spatial (dependency between the activity of different neurons) and temporal aspects of neural activities. The previous model based on transformers only considers the temporal aspect. The method is evaluated on four benchmark datasets and compared with several baselines, indicating the model's effectiveness.", "questions": "Among the four assessed metrics, the model is the best one, mainly from the aspect of co-bps; is this a consequence of hyper-parameters being tuned based on co-bps? Are the hyper-parameters of the baseline models also tuned based on co-bps?\n", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nThe proposed architecture is intuitive and interesting. It also addressed an important shortcoming with the previous transformer model.\n\nThe focus on interpretability is interesting and can potentially lead to important discoveries in neuroscience about the role of each brain region.\n\nThe paper is well presented and easy to follow.\n\nWeaknesses:\nSome of the architectural details are missing from the paper. For example, how many heads were used? What was the exact architecture for each layer? How is the test performance calculated? Is it based on masked inputs on test data?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "KV8BoBtOSWw", "reviewer": "Reviewer_gfnD", "summary": "This paper addresses the task of modelling noisy single-trial spiking activities and extends the Neural Data Transformer (NDT) model by including an attention mechanism between neurons that treats each neuron as a token. Additionally, the authors add the SimCLR contrastive learning objective to the model trying to improve generalization. The paper investigates the new model’s performance on spiking neuron activity of four publicly available motor control datasets.", "questions": " 1. Tables 1 + 2: I do not see an overall improvement by your additional attention model over AESMTE1/3. The metrics are above the baseline as frequently as they are below i, suggesting that we're looking at random fluctuations. Am I missing something?\n\n 1. Tables 3 + 4: Again, the ablation study seems to suggest to me that the contrastive objective does not actually improve the model. Am I missing something?\n\n 1. Does the model use positional encodings? It seems like that would be a useful thing to encode properties of the tokens such as neuron identity. Otherwise the tokens are permutation invariant, which means the attention mechanism has to infer the token identity from the current spatial (or temporal) pattern of spikes, which probably limits the model quite substantially (and unnecessarily). Please explain whether you used positional encodings and, if so, how or, if not, why.\n\n 1. L 102ff.: Please describe in more detail how exactly you perform the masking. Do you drop individual entries X_ij or entire rows and/or columns?\n\n 1. L 119: Please expand on how exactly you alter the spike counts for the contrastive loss and provide the dropout probability.\n\n 1. Fig 4: It is not clear to me what the attention maps show or what insights we can gain from them. I suspect the \"important\" neurons in layer 1 are just the ones with the most spikes, which form good basis functions for filling in since their response vectors are (a) dense and (b) tend to have higher signal-to-noise ration due to the (very approximate) Poisson statistics of spike trains.\n\n\n\n\n\n\n## Details\n\n - Fig 3A: Why does performance drop if you increase the ensemble size beyond approx. 8 models? Please also state how many individual models you use to form one ensemble.\n\n - Paragraph at l. 123: I think it is reasonable to focus on co-bps for the Bayesian hyperparameter search. At the same time I do think that the mask loss – which is probably the Poisson log-likelihood? – correlates with the 4 evaluation metrics you compare it to in Fig 2A. If they did not correlate at all, it would be unclear how training the model with the log-likelihood objective could improve the model’s performance on an uncorrelated evaluation metric.\n\n - You describe two out of the four datasets you used in Sec 3.1 and 3.2 in detail. I think your paper would improve if you would additionally describe the two other datasets you used. You might be able to cut down on the level of detail you provide in 3.1 and 3.2 in case the pages limit would be a problem.\n\n - Fig 1: The dimensions of X and X^T are stated wrong. According to line 69 X is N x T, and the spike rasters in Fig. 1 suggest the same. Accordingly, the size of the temporal attention matrix in the figure should be 5 x 5 and the spatial one 3 x 3. For X and X^T the columns should be colored, not the rows.\n\n - The model is well described, although you could make explicit in the Methods section and maybe Fig. 1 that the model is made of 4 layers of the transformer block depicted in Fig. 1\n\n - L 107: you could make clear that \\tilde{z} is the “output *log* firing rate”\n\n - L 131: Introduce later used abbreviation NLB\n\n - L 125 and l 150: you could indicate here that these two metrics correspond to the log-likelihood you use for model training to improve clarity\n\n - Fig 3C: This Fig might improve if you plot the ground truth and the predicted traces into one plot for better comparability. Please add a description what the different colors mean (different examples?).\n\n - L 226: instead of “across 4 attention layers” is not quite correct, as you only depict layers 1 and 4. As a matter of fact, it would be interesting to see the same plots for layer 2 and 3.\n\n - L 258: I do not understand the meaning of the part of the sentence starting “as well as the discovery…”. \n\n - The paragraph starting at l. 264 is hard to comprehend as the sentences are a bit long and complicated to read. This paragraphs readability could improve a lot by rewording / restructuring it.\n\n", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths \n + Research question is important and the presented modeling idea makes sense\n + The paper is well written and easy to understand\n + The motivation in the Introduction section as well as the comparison to related work is decent\n + Compares the model to a variety of baselines + metrics\n\n## Weaknesses\n - New attention mechanism does not lead to measurable improvement\n - Contrastive objective does not lead to measurable improvement\n - Unclear how attention maps provide \"interpretability\"\n - Some of the methods not clearly described (augementation, masking, positional encodings?)\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept