Neuroformer: Multimodal and Multitask Generative Pretraining for Brain Data

Thank you for acknowledging the interesting direction of our work, and the strong results.

Architecture. Because multiple authors had similar questions. We have added a global comment explaining the overall workflow, we suggest reading that first.

Neurons as tokens. Yes, each Neuron is modeled as a token (ID). In a GPT, not all tokens (byte-pair encodings, ~words, typically around ~50,000) are included in every block. Every block (context) consists of the BPE encodings required to form that example/sentence. In our case, we populate the block with the number of Spikes in the current interval. $O(1000)$ would mean 1000 spikes in an interval of 0.05s in our case, which would require a population of tens of thousands of neurons or more (spiking activity is quite sparse). This is one of the advantages of our method. The context is not a sparse matrix of (N_population), but only consists of (N_spikes_per_interval) where typically $N_{spikesperinterval} << N_{population}$. Other methods utilizing sparse matrices attempt to compress the representation in order to accommodate larger neuronal populations, but in doing so lose neuronal resolution attention. This is precisely what we wanted to avoid.

GLM vs. NF. We have updated the manuscript to include the actual correlations: the mean correlation for GLM $r=0.232$, Neuroformer $r=0.297$, and for Real-Real $r=0.409$. The correlations were calculated between groups of 4 trials.

Table 1. Thank you for recognizing the strong performance of our model on decoding speed. We did not include any past_behavior as input to any of the models, as the models could then simply extrapolate from previous behavior. For all models the output was the speed for the next 0.05s, over the whole of the holdout dataset. For all models bar GRU and Neuroformer, the input was the current state (spikes in the corresponding 0.05 seconds). For the GRU and Neuroformer, which can contextualize over states, we included the past state (0.15s of spikes before current state). We have updated the manuscript to include this information. We actually think figure 5b) is a cooler result, since it shows generatively pretrained models can transfer to decoding and match some of the other methods with only a fraction of the supervised decoding data.

N_a and N_z matrices. Both matrices have shape N x N. N_a denotes the pairwise attention between neurons, and N_z denotes how many times two neurons occur within the same bin. We used N_z to compute the average attention that was attributed between two neurons.

Thank you for taking the time to review our work.

Neural Latents Benchmark. NLB is a great benchmark, but it consists of electrophysiology (ephys) data, typically derived from multi-electrode arrays, rather than neuronal-resolution optical physiology (ophys) data which is what our model is designed for (2-photon calcium imaging data). Please also note that these datasets contain 100-200 electrode areas, while our datasets contain ~2000 Neurons. Our model can be adapted with few modifications to consider this data, and we're excited to explore such directions. Nevertheless, the benchmarks we have used our known to perform on-par with such methods. Please see Paragraph 4 of our global comment.

Representation Learning Results. There has been an important misconception here that we want to clarify. Our comment was not intended to mean that our model does poorly at representation learning. We have included some visualizations in the appendix (J, Multimodal Latent Variable) to show the low-dimensional latents the contrastive objective of our model learns. Rather, we wanted to state that some prominent representation learning methods fared poorly on decoding speed from the large-scale calcium imaging data (see K, CEBRA Behavior Results.)

Other than all the methods we tried, including the aforementioned one, the only other prominent method is RADICaL [1], which is an AutoLFADS type method adapted for calcium imaging data, and there is currently no openly available code implementation for the model.

Modalities. The modalities used are neural data, video, eye position (latitudinal and longitudinal angle), and speed. The model can additionally incorporate an arbitrary number of modalities or prediction tasks, and can be adapted to do so by just adding a few lines to a config file.

[1] Zhu et al., A deep learning framework for inference of single-trial neural population activity from calcium imaging with sub-frame temporal resolution. (2021) bioRxiv 2021.11.21.469441;

Thank you for taking the time to review and acknowledge our work. It makes us happy to see that people understand the ambitious direction we’re pointing towards.

Architecture. We have provided a global comment to all reviewers regarding the architecture. We suggest reading that first.

T_c and T_p are the corresponding current and past states' block size. They are populated with the spikes in those context windows which are made up of two representations: ID (the spatial location, or neuron in question) and DT (the temporal location, or time at which the spike occurred). To learn DTs, we break each current window into smaller sub-intervals (these can be arbitrarily small, we typically choose the temporal resolution of our measurements as the DT). The IDs and DTs are learnt as tokens, in the same way as words are in a normal GPT model. These tokens have an embedding shape (E). I.e, we learn an embedding of size E for each ID and DT. Therefore, the output at each step of the model is not a sparse representation. Rather, it is spikes within that interval. Each individual spike is fed back into the model, and the model autoregressively generates spikes in this way, just like it would do the same using words in the normal GPT setting. $p_i$ and $p_dt$ are the probabilities over the corresponding ID and DT tokens for each generation. We sample from these to generate the spikes.

Nucleus sampling. restricts the tokens we sample from according to a specified cumulative distribution. I.e, if we choose 0.95, any token that forms part for the cumulative distribution beyond 0.95, is not considered for that timestep.

Fig 2. While in figure 2 we weren’t necessarily trying to claim that attention could outperform all other methods, we agree that it would be interesting to compare attention with other methods as well. We have conducted a more thorough comparison using partial correlations, and granger causality. To compute the similarity of the ground-truth connectivity to the different methods, we computed the average sum of squared differences between the matrices (Frobenius norm). The attention matrix was more closely related to the ground truth connectivity matrix. Please see Appendix F for the visualized results.

Fig 3. The improvement in predictive capability is bounded by the variability between the actual ground-truth responses. In the quoted paper, McIntosh, NeurIPS 2016, Figure 2, the upper-bound is quoted as ~0.8, and this can depend on the way pearson correlation is calculated, i.e bin size. In the same figure, the GLM results look quite low, while the GLM we fitted approaches the ground-truth variability very well. For example, in this [1] they found that neurons which were < 50 microns from each other and synaptically connected typically had correlation values < 0.4, and that matches other studies of correlation at longer length scales, both within and across cortical areas [2]. Trial-to-trial correlations are typically low in mouse L2/3 neurons because the neurons have low firing rates (<< 1 Hz spontaneous and ~ 5 Hz evoked; e.g., Fig. 8 in Niell & Stryker 2008 J Neurosci). In supplementary materials, “Capturing Trial-To-Trial Variability,” we have included a comparison between sets of 4 ground-truth and predicted trials.

Sensorium models are not directly compatible with the data that we used to train Neuroformer, for two reasons. 1) Until 2022, models were trained on images and not video. 2) The baseline model for 2023, while compatible with video, requires precise pupil centers/directions to train the shifter network, which we do not have. Although we have tried in the last days to satisfy this request, we were unable to get well-performing models to compare with. While we still believe that the GLM we used is a strong baseline, particularly when considering the heavy inductive biases used to get good performance, we agree that more comparisons would be desirable.

Attention. These attention maps exhibit diversity. Some of the maps seem to highlight large, moving objects, others seem to highlight other mice in the field-of-view, and still other maps seem to be specific for small features, like a mouse’s tail. These are uncharted territory for neuroscience, and we're excited to explore the potential of this method in future work. To reiterate, the attentions shown in the figure are per-token, i.e for an individual spike (neuron), which is one of the advantages of the method, in that it represents neuronal populations at the spike level.

[1] Ko, H., Hofer, S., Pichler, B. et al. Functional specificity of local synaptic connections in neocortical networks. (2011) Nature

[2] Yu et al., Mesoscale correlation structure with single cell resolution during visual coding. (2019)

Thank you for the constructive feedback. We were happy to see some of our technical contributions recognized, alas, this seems not to be reflected in the final evaluation score. We hope to alleviate your concerns regarding additional benchmarks. Please also see Paragraph 4 of our global comments.

F1 scores. In study [1], neurons within 50 microns and synaptically connected showed correlations below 0.4, consistent with other research [2]. Low trial-to-trial correlations in mouse L2/3 neurons are due to their low firing rates, as detailed in Niell & Stryker 2008 J Neurosci. In supplementary materials G, we include a comparison of four sets of ground-truth and predicted trials, showing that our model’s predictions approach the ground-truth variability in responses. Furthermore, please note that in the visnav datasets, each prediction is conditioned on never-before-seen stimulus, as the mouse is freely navigating within a virtual environment setup (please see Figure 7 for experimental setup).

[1] Ko, H., Hofer, S., Pichler, B. et al. Functional specificity of local synaptic connections in neocortical networks. (2011) Nature

[2] Yu et al., Mesoscale correlation structure with single cell resolution during visual coding. (2019)

Benchmarks. Please see our comment (Paragraph 4). The benchmarks we used for comparison are common, and quite competitive even in the NLB settings, where they reach performance close to that of the most performant models on the benchmark, like AutoLFADS. In particular, we tried our best to make a performant GRU by tuning hyperparameters, using tapering, and bidirectionality (for more details see Appendix D).

Representation Learning for Behavior. This is an interesting thought. Following your request we have extracted the representations learnt by our contrastive module and provide our results in Figure 14 (Appendix J). The neural features separate according to speed. Please note that no dimensionality reduction techniques were used, these are the raw features from our model.

Long-Term Inference. Yes, Neuroformer is able to simulate neuronal responses over very large time horizons. The provided correlation results, and generated spike trains in figure 3b), are all generated using no ground-truth data, over 96 seconds. (w_p) the temporal window for spikes provided to (T_p), is only 0.15s (provided in Appendix C), and other than the first step in the generation of the 96 second trial, all subsequent steps are conditioned on the spikes previously generated. Appendix G shows that these simulations capture the variability of the ground-truth responses.

I appreciate the authors' responses to my questions. I read through the author's responses to all reviewers. This paper, in my opinion, is, at best a borderline case. 

As the primary demonstration of this work pertains to the generation of neural spikes, I proposed a comparison with alternative transformer-based architectures without specifically referencing the Neural Latent Benchmark. I recognize that the author may find it difficult to modify those alternative architectures to fit their data within the time constraint of the rebuttal period. I then proposed a discussion of cross-modality versus single-modality. In the absence of both comparisons, it is difficult for a system neuroscientist to migrate their modeling to cross-modality, given the potential difficulty of training (again, this is merely my conjecture; the author's response implied that they are unable to get the architecture to function on the neural latent benchmark). Its potential impediments could impede its applicability within the system neuroscience community.

A well-fitted GRU should not replace the comparison of alternative transformer architectures. The GRU fitted to the author's datasets would differ from the GRU fitted to NLB. To claim that comparing to GRU is sufficient is akin to comparing pears and oranges. Critically, I asked about the comparison with alternative transformer architectures due to the potential for the attention head to contain interesting information about stimulus encoding. 

I also find the speed decoding shown in Fig 14 not convincing. There is clear overlap between points of different colors (or maybe it is easier to see separation if the opacity level is less than 1?). I’d also appreciate some clarification of having the dimension of neural features as 3. I could not find such architecture details in Appendix C.

I echo with the Reviewer pMwz that the architecture details are unclear.

I will raise the score to 5, and this is my final score for this submission.

Minor

I also find the updated paper a bit sloppy in its presentation. 3 subfigures in Fig 2 are missing. 9/10 subfigures in Fig 12 are missing.

We are sharing an anonymous link with an animation showing how spikes are auto-regressively generated, a video showing average attention per-time step from the point-of-view of the mouse, and our current codebase.

Note: To avoid any issues we have uploaded the aforementioned materials in the supplementary section.