Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: 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.

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.

However, there are some major concerns regarding the paper:
1) 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:
- Schneidman et al (2006): https://europepmc.org/backend/ptpmcrender.fcgi?accid=PMC1785327&blobtype=pdf
- Macke et al. (2011): http://www.gatsby.ucl.ac.uk/~maneesh/papers/macke-etal-2011-nips-preprint.pdf
- Lyamzin et al. (2010): http://journal.frontiersin.org/article/10.3389/fncom.2010.00144/abstract
- Molano-Mazon et al. (2018): http://arxiv.org/abs/1803.00338
- Ramesh et al (2019): https://openreview.net/forum?id=S1xxRoLKLH
- Bashiri et al. (2021): https://openreview.net/pdf?id=1yeYYtLqq7K

2) 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).

3) 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.
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).

2) 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. 
It 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?

3)  It would be good to have information about training costs, computational and sample efficiency of the menthod

4) 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.

Minor comments:
1) Figures 1-4 have grey lines around the border that don't nescessarily overlap with the division of subpanels
2) 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.

Taken 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.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: 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.
strengths_and_weaknesses: Strengths:
The proposed architecture is intuitive and interesting. It also addressed an important shortcoming with the previous transformer model.

The focus on interpretability is interesting and can potentially lead to important discoveries in neuroscience about the role of each brain region.

The paper is well presented and easy to follow.

Weaknesses:
Some 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?
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?

limitations: N/A
review_title: null

final_score: rating: 4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.
review_content: 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.
strengths_and_weaknesses: ## Strengths 
 + Research question is important and the presented modeling idea makes sense
 + The paper is well written and easy to understand
 + The motivation in the Introduction section as well as the comparison to related work is decent
 + Compares the model to a variety of baselines + metrics

## Weaknesses
 - New attention mechanism does not lead to measurable improvement
 - Contrastive objective does not lead to measurable improvement
 - Unclear how attention maps provide "interpretability"
 - Some of the methods not clearly described (augementation, masking, positional encodings?)

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?

 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?

 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.

 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?

 1. L 119: Please expand on how exactly you alter the spike counts for the contrastive loss and provide the dropout probability.

 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.






## Details

 - 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.

 - 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.

 - 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.

 - 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.

 - 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

 - L 107: you could make clear that \tilde{z} is the “output *log* firing rate”

 - L 131: Introduce later used abbreviation NLB

 - 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

 - 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?).

 - 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.

 - L 258: I do not understand the meaning of the part of the sentence starting “as well as the discovery…”. 

 - 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.


limitations: Yes
review_title: null

Meta-review: metareview: This paper introduces a spatiotemporal neural data transformer (STNDT) model for multineuronal spike trains. The method achieves state of the art performance on a variety of tasks within the Neural Latents Benchmark (NLB), which was introduced at NeurIPS in 2021. Technically speaking, the approach is a slight tweak on standard transformers, but it allows for trial-by-trial reweighting of values based on a neurons-by-neurons similarity matrix. It is interesting that this tweak is enough to boost performance.

While the reviewers expressed some hesitation with regard to the degree of technical novelty and the questionable interpretability of the model, I think it is important to recognize improvements in benchmarks that the field has put forward at NeurIPS. The paper is well written and the results are thorough. Overall, I think it is a valuable contribution.

In addition to addressing the reviewers' concerns, I would like the authors to address the following points in the final paper:
- I still think Figure 1 is confusing. The spike train in the top row and the firing rates in the bottom left seem to suggest that there are N=3 neurons and T=5 time steps, but the matrices and stated shapes suggest that there N=5 neurons and T=3 time steps. I believe the authors intended to have N=3 and T=5. If so, please correct the figure to help avoid any confusion!
- The introduction does a good job of summarizing related work in statistical neuroscience, but it does not, as far as I can see, discuss the enormous literature on variations of transformers. I do not know the best references, but I would be surprised if there haven't been related works in the transformer literature for other spatiotemporal tasks (e.g. applied to modeling audio or video data). The lack of discussion of related work in the broader ML community is a glaring omission. 
confidence: Less certain
Decision: Accept