Towards a "Universal Translator" for Neural Dynamics at Single-Cell, Single-Spike Resolution

Weaknesses:

1. We agree with the reviewer that our submission was lacking in qualitative analysis. To partially address this, we provide visualizations of the rastermap reconstructions for all predictive tasks in Figure 1 of our one page pdf. For the final version of our paper, we will include rastermap visualizations for many different animals/sessions, so readers can see differences in neural activity across different individuals. We believe there might be some confusion about the neural dynamics learned in our models. Since NDT1 is a transformer-based model, neural dynamics are learned implicitly. We are not explicitly learning neural dynamics like a state-space model would. So it is not possible to visualize the dynamics directly in a transformer. Building transformer architectures with explicit latent dynamics is an interesting future direction.

2. These are very interesting questions! (a) When training on single region data, we would recommend just using the neuron and causal masking, as region-level predictions are not informative in this setting. When training on data from multiple brain areas, we uniformly sample brain areas for inter/intra region masking and then we average the loss per-neuron and per-timestep; this makes sure that no brain region dominates the loss. (b) We are using the IBL repeated sites dataset [1] which only have repeated sessions for one individual so we cannot easily address the difference in performance when using multi-session data from a single individual versus cross-individual multi-session data. Using the full public IBL brain-wide map dataset [2] would allow us to test this phenomena and is an exciting future direction.

[1] International Brain Laboratory, et al. "Reproducibility of in-vivo electrophysiological measurements in mice." bioRxiv (2022): 2022-05.

[2] International Brain Laboratory, et al. "A brain-wide map of neural activity during complex behaviour." Biorxiv (2023): 2023-07.

Questions:

1. This is an interesting question and something we have been thinking about as well. As a first attempt to answer this question, we ran a new experiment where we try to predict region A from region B for all regions in a test session using a model trained with MtM. As this is different from our current inter-region predictive task (which predicts one region from all other regions), we fine-tuned our pretrained MtM model with a new inter-region task where we randomly dropout some of the regions when predicting one region. The prediction matrix (region A to region B for all regions) is shown in Figure 2 of the one page PDF. It can be seen that there is interesting structure in the matrix where some regions can predict each other well, for example in visual areas. We believe that exploring the "functional connectivity" of different areas using MtM is an interesting future work.

2. Good question! Models trained with MtM will be immediately useful for understanding interactions between different brain areas. As discussed above, we were able to construct a "functional connectivity" matrix using MtM that can show how predictive different regions are from each other. We believe it may also be possible to examine the attention weights when predicting region A from region B to interpret temporal delays between the two areas. Using MtM to explore common patterns among different individuals is more challenging. An interesting result of our paper is that when pretraining with more individuals, the single individual performance improves after fine-tuning. This indicates that there are common patterns among individuals that are being captured by MtM. If we were to utilize the MtM objective with a latent variable model like LFADS, it may be possible to look at how dynamics in different areas differ among different individuals. This is an exciting figure direction we hope to explore.

Weaknesses:

1. We agree with the reviewer that this is a important detail missing from our paper. To address this, we added schematics for NDT1 and NDT2 that demonstrate how tokenization and neural reconstruction works (please see Figure 3 in the 1 page PDF). We will add this figure to the supplement in the final version of the paper.

2. We agree with the reviewer that the lack of visualizations is a problem with our original submission. To partially address this, we provide visualizations of the rastermap reconstructions for all predictive tasks in Figure 1 of our one page pdf. The difference in the bps metric for temporal masking and MtM can clearly be seen in the region-level predictive tasks and can also be seen in neuron and forward prediction tasks. We will make sure to add rastermap visualizations for many different sessions in the supplement of our final paper. Hopefully, this will convince readers to try out our approach for NeurIPS 2024! :-)

Questions:

1. We define the bits per spike for a neuron with the same equation as [1] on the bottom of page 6. We then average the bits per spike for each neuron in a test session to get the final performance. We will include a version of this equation in the updated manuscript.

2. The bits per spike (bps) can be negative when the predicted rates for a neuron are less informative than just predicting the average firing rate of the neuron for the trial. This can be seen in the above equation where we subtract the likelihood of the neural activity assuming an average firing rate from the likelihood of the neural activity with the predicted rates of the model. The bps also tends to be more negative when trying to predict low-firing rate neurons.

3. "Choice" refers to the choice decoding accuracy for each method. "WME" refers to the whisker motion energy decoding $R^2$ for each method . We will improve the axis labels in the final version of the paper. We will also try fixing the axis for each session as suggested by the reviewer.

[1] Pei, Felix, et al. "Neural latents benchmark'21: evaluating latent variable models of neural population activity." arXiv preprint arXiv:2109.04463 (2021).

Figures 1 and 3 in the rebuttal PDF are very nice and informative. Thank you for providing those. I think Figure 1 should even be part of the manuscript itself (with more visualizations in the supplementary materials)

Q1: It seems to me that MtM smoothes adjacent time bins quite a lot. Ground truth often has almost empty time bins, followed by a sharp increase in the firing rate right in the next bin. But in MtM's reconstruction the change of firing rate is always gradual from bin to bin. And that worries me because sharp increases and sudden spikes or drops are very important for actual neuroscientific insight. Why dies this smoothing occur? Is adjacency in the time series plays some sort of a special role in a transformer? Or is it due to the loss function? Can you define a loss function that would strive to preserve sharp changes?

Q2: I am still somewhat unhappy with the "bits per spike" measure :), however seeing the actual definition was very helpful and for sure needs to be part of the manuscript. The issue I take with it - what is the benefit of adding this additional layer of abstraction and measure via Poisson model? Why not go super simple and just count the number of spikes the prediction got wrong? (and then average in a suitable manner to make comparable across). Spikes are binary, and this makes it easy to count them. Maybe it's just me, but when I see a metric like "bps: 0.51" I have no idea how good or bad is it. Maybe if you provide info on what the bps would be for a perfect reconstruction? I guess what I am looking for is a more intuitive number here, like "73% of spikes are where they should be, 27% are missing / extra" or "the number of spikes MtM predicts within each bin is off by 2.7 spikes per bin, or 14%".

Basically I am still not able to understand when I am looking at "Baseline bps: 0.03" vs. "MtM bps: 0.51" should I be impressed? How impressed should I be? Reconstruction-wise both Baseline and MtM look pretty far away from the ground truth.

Q3 (exploratory, just for fun): Do I understand correctly that currently tokens are based on firing rates? Is yes, then my question is would it be possible to use analgues of "spiking words" or "spiking syllables" of fixed length instead? Something like

||||__|||| -> token 1
|||||||___ -> token 2
___||||___ -> token 3
... etc

Where "|" are spikes and "_" are non-spikes. To basically build a vocabulary of all possible 10ms spiking patterns and use those as tokens?

Or, actually, why use bins at all? What would happen if you use spiking sequences as-is, so that the time step is small enough to just be either | or _ ?

Weaknesses:

1. We thank the reviewer for this suggestion to compare to a wider range of neural population models. We want to clarify that MtM can be used as a learning objective for any architecture (not just transformers). To demonstrate this, we ran a new experiment using the LFADS architecture (a sequential VAE) trained with both temporal masking and MtM. The results can be seen in Table 1 of the Global Response. MtM training improves all activity prediction metrics for LFADS especially intra-region masking and inter-region masking. The best performing model is still NDT1-MtM, but the LFADS-MtM model is quite competitive. This experiment provides evidence that MtM is a model-agnostic objective that can teach region-level population structure to both transformer and RNN architectures. We also attempted to compare MtM to mDLAG on three regions (VISa2/3, VISa4, VISa5) of a single test session. For these regions, there were a total of 42 neurons, 828 trials, and 100 time bins (20 ms bin size). We trained mDLAG for 5.5 hours (5000 EM iterations) and the model did not converge. A visual inspection shows that the learned latent factors for each region were nearly identical and that activity reconstruction from the learned latents consistently led to low co-bps (negative to zero). This experiment provides evidence that scaling mDLAG to full IBL datasets with tens of regions and thousands of neurons is not feasible. We will provide all details for this experiment in the updated manuscript.

2. We agree with the reviewer that this is an important detail missing from the initial submission. We will make sure to add a discussion of both temporal masking for NDT1 and random token masking for NDT2 in the final submission. We also plan to add schematics for the NDT1 and NDT2 architectures in the updated manuscript (Figure 3 of the one page PDF) to demonstrate how they tokenize and reconstruct neural data.

Questions:

1. This is a great question and something we are also thinking hard about! We believe that by incorporating data from many different sources, a neural foundation model can help us answer important questions about cross-individual (and even cross-species) neural variability. We also believe that neural foundation models, which are trained on many different combinations of brain regions, can eventually be used to predict missing brain regions in new recordings. We also believe that by training neural foundation models on healthy animals, we can test hypotheses about neural activity in diseased animals. This advancement could potentially lead to new treatments and interventions for neurological disorders. Finally, it will be possible to fine-tune pretrained foundation models on new tasks (brain computer interfaces, etc.) to achieve state-of-the-art performance [1]. We believe our work is a step towards realizing this vision and we are hopeful that the community can start working with and building similar models. We will improve our discussion of foundation models of neural data in the updated manuscript.

[1] Azabou, Mehdi, et al. "A unified, scalable framework for neural population decoding." Advances in Neural Information Processing Systems 36 (2024).

Weaknesses:

1. We thank the reviewer for the suggestion to run more baseline architectures with the MtM objective. Our goal was to compare masked modeling approaches for neural population modeling which just includes temporal masking as a baseline. Based on this feedback, however, we ran a new experiment where we trained the LFADS architecture (a sequential VAE) with both temporal masking and MtM. The results can be seen in Table 1 of the Global Response. MtM training improves all activity prediction metrics for LFADS especially intra-region masking and inter-region masking. The best performing model is still NDT1-MtM, but the LFADS-MtM model is quite competitive. This experiment provides evidence that MtM is a model-agnostic objective that can teach region-level population structure to both transformer and RNN architectures.
2. We agree with the reviewer that there is a lack of task and species diversity in our original submission. To partially address these concerns, we fine-tuned our 34-session pretrained NDT1-MtM on the Neural Latents Benchmark (NLB) MC_RTT dataset. For details on this experiment, please see the Global Response. With a single private submission to the leader board, the MtM performance (.19 bps) significantly improves upon NDT's previous best performance (.16 bps). These results show promising transfer of our pretrained MtM models to a new task and species. We also want to point out that while the IBL datasets lack species and task diversity, they have high diversity of recorded brain areas. Brain region diversity is an unexplored axis for large-scale pretraining as the POYO and NDT2 papers only train on data from motor areas. We wanted to tackle this challenge first before moving on to cross-species transfer. However, we believe that these approaches are complementary and that building large-scale models that generalize across different tasks, species, and brain areas will be an exciting future direction.

Questions:

1. We agree with the reviewer that our paper lacks a simpler NLB baseline and have implemented an LFADS baseline model trained with temporal and MtM masking (results in Table 1 of the Global Response). NDT1 is still the top performing model when trained with the MtM objective although LFADS-MtM is competitive. While we agree that the LFADS RNN architecture is simpler, transformers have been shown to scale extremely well to massive datasets and also have benefits for computational efficiency (please see Figure 6 of [1]). We will clarify these reasons for using transformers in the updated manuscript. We did not include an additional MLP baseline for behavior decoding because we are decoding from the inferred rates and, therefore, would not expect to do better than an MLP. This is the strategy used in the Neural Latents Benchmark [2] to evaluate behavior prediction.

2. This is a great question and something we are interested in exploring as future research. To attempt to answer this question, we fine-tuned our 34-session pretrained NDT1-MtM on the Neural Latents Benchmark (NLB) MC_RTT dataset. All details for this experiment can be found in the Global Response.

3. Yes, there is some overloading of terminology in our paper between behavioral "tasks" in neuroscience and the predictive "tasks" of MtM. In MtM, there are multiple neuron-level and region- level predictive tasks. We will make this distinction clear in the final version of the paper.

4. Yes, both NDT1 and NDT2 use a positional embedding layer for each token. NDT1 uses the timestep for the positional encoding. NDT2 uses 2 types of position encoding: (1) timesteps and (2) neuron group identity. We did not include the positional embedding step in Equation (1), but we can correct this omission in the final version. We also created schematic diagrams for both NDT1 and NDT2 (see Figure 3 in the one page PDF) which we will add to the supplement to improve our description of the underlying architectures.

5. In the held-out sessions, there are 397 ~ 579 trials in the train split, 57 ~ 83 trials in the validation split, and 114 ~ 166 trials in the test split.

6. Yes, good catch. This should read "10% of the trial (200ms)". We will correct this in the final version.

7. For each trial, we predict the binary choice using the entire $N \times T$ matrix $X$, which represents the spiking activity of all neurons across all timesteps.

8. In each trial, we use the entire $N \times T$ matrix $X$ to predict the motion energy at every timestep $t = 1, 2, …, T$, where $T$ is the total number of timesteps.

9. For binary choice, we report the test accuracy that will be reached when constantly predicting the train set's majority class. For whisker motion energy, we report the test $R^2$ that will be reached when consistently using the train set’s trial average. Despite imbalanced datasets (e.g., c7bf2d49), Figure 5 shows that decoding accuracy and R2 from selected regions are generally larger than these chance levels.

   Chance Level	5dcee0eb	c7bf2d49	3d9a098	d57df551	824cf03d
   Choice	0.58	0.91	0.52	0.49	0.62
   Whisker Motion Energy	0.44	0.30	0.59	0.43	0.39

[1] Ye, Joel, and Chethan Pandarinath. "Representation learning for neural population activity with Neural Data Transformers." arXiv, 2021.

[2] Pei, Felix, et al. "Neural latents benchmark'21: evaluating latent variable models of neural population activity." arXiv preprint arXiv:2109.04463 (2021).