Paper: Neuroformer: Multimodal and Multitask Generative Pretraining for Brain Data (ICLR-2024)

Reviewer comment (Reproducibility):
Although the paper is generally well-written, many model details are not sufficiently clear.

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

Reviewer comment (Writing):
The architecture description may be confusing because one might naively expect time bins, rather than neurons, to be tokens in a time-series modeling approach similar to GPT. This setup requires a clearer and more detailed explanation.

Author response:
**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.

Reviewer comment (Evaluation):
In Section 4.3.1, the sentence "our model's predictions w[h]ere more closerly corelated with the ground-truth" should include quantification, such as the fraction of neurons (or clarification on whether the dots in figure 3c are neurons) where Neuroformer outperforms GLM, and the p-value (0.02), which is currently mentioned only in the figure caption. The evidence provided is not very strong, so the authors should avoid overselling.

Author response:
**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.

Reviewer comment (Experiments):
Table 1 presents the strongest result in the paper, showing that Neuroformer strongly outperforms all other models. However, neither the comparison models (Lasso, GLM, etc.) nor the prediction task are described in sufficient detail. It is unclear what "behaviour" refers to, and what the input (how many time steps) and output (how many time steps) of the prediction problem are. The authors should present their experiment so that it is clear the comparison models in Table 1 are not “strawman” baselines.

Author response:
**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.

Reviewer comment (Writing):
On page 6, the meanings of the N_a and N_z matrices are unclear.

Author response:
**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.

Reviewer comment (Evaluation):
As the authors mentioned in Appendix A, the Neuroformer has poor results (Figure 10) in low-dimensional latent space learning. One possible reason is that the transformer-related neural networks are too expressive so that good dynamics in latent space are no longer necessary. But in terms of interpretability, neuroscientists prefer to observe meaningful latent space in many experimental scenarios, which may be more important than a good behavior prediction performance.

Author response:
**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.

**[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;

Reviewer comment (Writing):
What kind of modalities are used in section 4.3? Although we may infer these modalities in section 5, it seems there are no clear describing sentences about them in the whole section 4.3.

Author response:
**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.

Reviewer comment (Experiments):
Weak baselines in Figs. 2+3.

Author response:
**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.

Reviewer comment (Writing):
Architecture (especially decoder) not entirely clear from paper.

Author response:
`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.

Reviewer comment (Experiments):
In Fig. 2, inferring connectivity from correlations is not a good approach; the more appropriate method would be using partial correlations, which might correspond more closely to the ground truth than the Pearson correlations shown in Fig. 2d. Question: Does the proposed approach perform on par with or better than partial correlations?

Author response:
**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.  

| Metric              | Value  ↓ (lower is better)  |
|---------------------|----------|
| Attention           | 140.4    |
| Granger Causality   | 138.1    |
| Partial Correlations| 156.9    |

Reviewer comment (Writing):
In the architecture description (pages 3+4), the sentence on p.3 about projecting neuron spikes into an embedding space E ("vectors (T_c,E)") is unclear. The meaning of the parentheses around T_c,E is confusing—T_c appears to be a scalar (sequence length). Is it concatenated with E? If so, why? If not, what is happening?

Author response:
`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.

Reviewer comment (Writing):
It is unclear what the decoder outputs. The interpretation from the text and figures is that it might be a sparse representation of future spikes, with ID and dt indexing a binary matrix (#neurons × #timesteps) of spikes. However, this cannot be clearly mapped onto the losses in Eqs. 3+4, partly because p_i and p_dt are undefined. The reason for having two losses instead of outputting a zero/one vector for each neuron is not explained. The meaning of dt is also unclear—if predicting the next frame, why is dt needed?

Author response:
`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.

Reviewer comment (Writing):
The last paragraph on p.4 is opaque. The sentence involving nucleus sampling is unclear and the meaning of sub-intervals is also not clear.

Author response:
**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.

Reviewer comment (Evaluation):
The F1 scores in Figure 5 are rather low at their absolute values. It would be helpful if the authors put the F1 score in perspective (why is this F1 score indicating good performance?). Is it possible for the authors to comment on the dip of the F1 score after adding the video modality in the Visnav, Lateral dataset?

Author response:
**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)

Reviewer comment (Experiments):
What is the range of T_p in those calcium imaging datasets? Is it possible to demonstrate long-term inference with a neuroformer?

Author response:
**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.

Reviewer comment (Experiments):
In Fig. 3d, the attention maps seem to suggest that the intermediate blocks of the neuroformer contain interesting features. It is common for the community that pretrains transformers for time series (like HuBert or Whisper for speech) to use features from intermediate blocks for decoding. Is it possible to show how well these intermediate blocks can be used to predict behavior?

Author response:
**Intermediate features** The features in Fig. 14 came from the contrastive alignment module. For the experiments to extract the 3D latents, we set the $E_{clip}$ (see Hyperparameters, Appendix C) to 3, projecting all the features to 3 dimensions for alignment during training (our hypothesis was 2 dimensions for direction, since the mouse is moving on a 2D plane, and 1 dimension for magnitude). We wouldn't expect intermediate representations to completely separate into speed without any explicit supervision. Your suggestion is nevertheless a good analysis case. Figure 2 (toy model experiment) is pointing towards the same direction. We have further extracted intermediate representations from the transformer blocks of a model not trained on speed, and fit a linear regression line mapping them to speed. We also use random arrays of the same shape, as a control. Note that the mapping is from a 256-dimensional embedding. We have included the results as a jupyter notebook (`neuroformer_extract_latents.ipynb`) in the supplement. Interestingly, *as we fuse more modalities, the representation becomes more linearly separable in terms of speed.*

  ###    Linearly Probing Speed
| Feature                     | R2               |
|-----------------------------|------------------|
| random                      |  0.0521 |
| past state                     | 0.2299 |
| past state + stimulus          | 0.5341 |
| past state + stimulus + current state  | 0.6357 |