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

Authors: yizi  zhang
yanchen  wang
donato m jiménez-benetó
zixuan  wang
mehdi  azabou
blake  richards
renee  tung
olivier  winter
the  international
brain  laboratory
eva  dyer
liam  paninski
cole  hurwitz
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:17:34

Abstract: Neuroscience research has made immense progress over the last decade, but our understanding of the brain remains fragmented and piecemeal: the dream of probing an arbitrary brain region and automatically reading out the information encoded in its neural activity remains out of reach. In this work, we build towards a first foundation model for neural spiking data that can solve a diverse set of tasks across multiple brain areas. We introduce a novel self-supervised modeling approach for population activity in which the model alternates between masking out and reconstructing neural activity across different time steps, neurons, and brain regions. To evaluate our approach, we design unsupervised and supervised prediction tasks using the International Brain Laboratory repeated site dataset, which is comprised of Neuropixels recordings targeting the same brain locations across 48 animals and experimental sessions. The prediction tasks include single-neuron and region-level activity prediction, forward prediction, and behavior decoding. We demonstrate that our multi-task-masking (MtM) approach significantly improves the performance of current state-of-the-art population models and enables multi-task learning. We also show that by training on multiple animals, we can improve the generalization ability of the model to unseen animals, paving the way for a foundation model of the brain at single-cell, single-spike resolution.

Main paper: 
 introduction: 
 Recent studies in experimental neuroscience suggest that neural computation is highly distributed across the brain and sparse in nature [34, 23] . Much of our current understanding of the brain, however, originates from studying small circuits of neurons in hand-selected brain areas during stereotyped behaviors. This has resulted in the development of neural population models that are often brain region-specific and narrowly crafted for particular experimental contexts, limiting their broader applicability and insights into distributed brain function [38] . The arrival of multi-animal neural datasets that span hundreds of interconnected brain regions [18, 17] necessitates the development of a more general approach for building neural population models.To address these challenges, recent efforts have been directed towards building models that can be trained on neural data collected across multiple sessions and animals [3, 42] . These models are pre-trained on large corpuses of neural population data and then fine-tuned for downstream tasks such as behavior decoding and brain-computer interface (BCI) control [42] , leading to improved performance and generalization to novel sessions and animals. While these models are a promising step in the right direction, two crucial elements are currently missing. First, the pre-training is only performed on neural data from the sensorimotor network (M1, PMd, S1) which limits the applicability of these approaches to whole-brain analyses. Second, these models do not explicitly model the underlying brain regions, instead treating the population as a homogenous set of neurons. We argue that a foundation model for neural spiking data must be able to seamlessly translate across all spatial scales, including population-level, region-level, and single-neuron-level dynamics.In this work, we build towards a first foundation model for neural spiking data which can solve a diverse set of predictive tasks across diverse brain areas. Similar to [41, 42] , we utilize masked modeling where parts of the input are masked and then reconstructed using the unmasked inputs. To explicitly capture neural dynamics across all spatial scales, we introduce a multi-task-masking (MtM) approach where the model alternates between masking then reconstructing neural activity in masked time steps, neurons, and brain regions. We learn a "prompt" token which allows the model to seamlessly switch between different masking objectives during training [35] . During evaluation, this prompt token can be utilized to perform "mode switching" where downstream tasks are associated with specific masking schemes.We evaluate our approach using the International Brain Laboratory (IBL) repeated site dataset [18] which consists of multi-region Neuropixels recordings that target the same brain regions (secondary visual areas, hippocampus, and thalamus) across multiple animals. We design a number of unsupervised and supervised prediction tasks which include single-neuron and region-level activity prediction, forward prediction, and behavior decoding. We benchmark our MtM approach against the temporal masking scheme used by Neural Data Transformer (NDT) [41] and the random token masking scheme used by Neural Data Transformer 2 (NDT2) [42] . We show that even with the same architecture, our MtM approach significantly outperforms the temporal masking baselines and enables multi-task learning. To demonstrate that our MtM approach is a viable recipe for large-scale pre-training, we train across 34 animals and fine-tune on 5 held-out animals. The performance of our MtM approach continuously scales with more training sessions, indicating its potential as a "universal translator" of neural dynamics at single-cell, single-spike resolution.• A novel multi-task-masking (MtM) approach which can be applied to multi-region datasets to successfully learn representations that lead to better downstream task performance.• A prompt-based approach for test-time adaptation which improves performance on a variety of downstream tasks during inference.• Scaling results that demonstrate that having data from more animals provides benefits on held-out animals and sessions as well as on unseen tasks.• A new multi-task, multi-region benchmark for evaluating foundation models of neural population activity.2 Related Work 
foundation models for neuroscience: 
 The advent of large-scale, self-supervised foundation models has marked a paradigm shift across various domains of machine learning. These models, diverging from traditional annotation-reliant supervised models, exhibit an impressive ability to generalize across a spectrum of tasks. These models have transformed natural language processing [27, 28, 40] , vision [29, 16] , and robotics [30] , and are beginning to reshape the landscape of life sciences [5, 1] . The application of foundation models to neuroscience is of significant interest. While there has been considerable progress in building large-scale models for EEG, [6] , fMRI [36] , and sEMG [19] , which have ample data availability, no large-scale model exists for neural data at single-neuron, single-spike resolution. To address this gap, two new methods, a supervised method, POYO [3] , and a self-supervised method, NDT2 [42] , were trained on a large corpus of spiking data from ∼12 monkeys. While promising, the datasets used for training are from just a few animals and brain regions, and therefore lack scale and diversity, limiting their applicability to other brain areas and behavioral contexts. We alternate between different masking schemes during training along with a learnable "prompt" token which provides context to the model about the associated task [35] . During evaluation, we provide the associated prompt token for the downstream task to perform test-time adaptation of the model. Our MtM approach is architecture-agnostic as masking is performed on the input data (not the tokens). For a full discussion of MtM, see Section 3.Transformers [39] have recently been applied to neural population activity in both supervised [3, 20] and self-supervised [41, 42] settings. POYO [3] is a supervised multi-session model for predicting behavior from neural activity. The POYO architecture utilizes learnable unit embeddings for each neuron and a novel approach for tokenizing individual spike events with relative position encodings to incorporate precise spike timings [3] . For self-supervised learning, NDT1 [41] and NDT2 [3] are two existing transformer-based methods. Both utilize masked modeling and assume a Poisson emission model. For NDT1, each time bin is a token with dimensionality equal to the number of neurons. NDT1 uses a simple encoder-only transformer to reconstruct masked time bins during training. To incorporate information across many sessions spanning different sets of neurons and across individuals, NDT2 [42] was introduced. This model is ViT-style [7] time-series transformer architecture that uses spatiotemporal patches as tokens and a learned session-level context embedding. A token is defined as a single time step of neural activity for a subset of neurons and is masked and reconstructed during training. By performing patching of the neurons, NDT2 can easily be applied to multiple sessions.How information is encoded within and across different brain areas is an important question that underlies the study of brain organization [12] , the evolution of brain development [2] , and the diagnosis of different network-level brain diseases [4] . Recent advances in electrophysiological techniques now allow for recording neural activity across many interconnected regions simultaneously [13, 33, 43, 37] . This has inspired recent efforts to build more fine-grained estimates of multi-region communication [26, 32, 12] and to investigate neuron-level information processing across multiple brain areas [17] . To analyze these multi-region datasets, generative models have been developed that aim to identify low-dimensional latent variables representing shared activity among recorded areas [11, 31, 14, 8, 9] . Recently, dCSFA [11] and DLAG [9] were introduced to model temporal delays between two brain areas. mDLAG [10] further extends this approach to an arbitrary number of brain areas. While these approaches are interpretable solutions to understanding intra-and inter-region neural dynamics, they also make limiting assumptions on the structure of the communication signals and cannot be easily scaled to many brain regions and sessions. 
methods: 
 A foundation model for neural activity must be able to seamlessly "translate" across all spatial scales of the brain, including population-level, region-level, and single-neuron-level dynamics. To this end, we introduce a multi-task-masking (MtM) approach for self-supervised learning of neural activity.During training, we alternate between masking and then reconstructing neural activity across masked time steps, neurons, and brain regions. We utilize a learnable "prompt" token which provides the model with context about which masking scheme is being applied during training. This prompt token can be passed to the model at test time to adapt the model to the associated downstream task [35] .Masked modeling for neural activity is typically performed by masking and then reconstructing activity in random time steps [41, 42] . While this masking scheme allows for learning temporal dynamics, it can ignore important spatiotemporal structure present in neural activity. To address this limitation, we propose four masking schemes designed to capture diverse patterns in neural data (see Figure 1 for a visualization of these masking schemes).• Causal masking. Similar to a GPT-like model [28] , we mask future time steps and then predict them using past time steps. While we learn next time step prediction, this masking scheme can be extended to multiple future time steps prediction as well. • Neuron masking. We randomly mask neurons and reconstruct their activity using the unmasked neurons as context. This masking scheme allows the model to learn how individual neurons relate to the activity of the full population. This is conceptually similar to "coordinated dropout" introduced in [15] . • Intra-region masking. We randomly mask neurons in a randomly chosen brain region and then reconstruct them using only unmasked neurons from the same region as context. This masking scheme allows the model to learn intra-area dynamics. • Inter-region masking. We randomly mask neurons in a randomly chosen brain region and then reconstruct them using unmasked neurons in other regions as context. This masking scheme allows the model to learn about cross-region interactions.Each of these masking schemes teaches the model about different structure in neural populations. We hypothesize that a model trained with all these diverse masking schemes will be able to solve many different tasks at inference time.To train using MtM, we randomly sample a masking scheme M for each batch of neural data. We mask the input data according to the sampled masking scheme and then pass this to a tokenizer and transformer-based architecture. To preserve the temporal or spatial order of the data, we add positional embeddings P E to the neural data tokens Z. To provide context to the model about the masking scheme that was sampled, we prepend a learnable "prompt" token P to the neural data tokens. This prompt token is a D-dimensional learnable embedding which the model can use to adapt its behavior during training or evaluation. For a batch of neural data X, our training is as follows:EQUATIONwhere X is the neural data reconstructed by the model, based on the time-varying rates inferred from the transformer's output, assuming a Poisson emission model. The MtM approach is agnostic to the choice of tokenizer and transformer, allowing it to be utilized with any architecture. For this work, we utilize the same tokenization scheme and transformer-based architecture as NDT1 and NDT2 for all comparisons to these methods (see Section 2.2 for details).Depending on the downstream task, the information learned using each masking scheme might be more or less useful. We utilize a prompt-based "mode switching" [35] approach where the prompt token that is best associated with the downstream task is prepended to the neural data tokens during inference and fine-tuning. 
dataset: 
 For training and evaluating our models, we use the International Brain Laboratory repeated sites dataset [18] . This dataset consists of Neuropixels recordings collected from nine labs which utilize a standardized experimental pipeline. The recordings target the same five brain regions across 48 adult mice performing a complex decision-making task. The probe was localized after the experiments using reconstructed histology and the brain regions were annotated. We utilize trial-aligned, spikesorted data from 39/48 mice for our analyses. From these recordings, we have a total of 26,376 neurons for training and evaluation (∼676 neurons per session on average). We bin the neural activity using 20ms windows and we fix the trial-length to 2 seconds (200 time bins). For behavior decoding, we exclude trials based on reaction time outliers as defined by the IBL brain-wide map [17] .We utilize a number of unsupervised and supervised metrics to evaluate how well a neural population model generalizes to different downstream predictive tasks.• Co-smoothing. Predicting the activity of a held-out neuron using all other neurons [25] .• Forward prediction. Predicting future activity from past activity. We predict the last 10% (200 ms) of the trial-aligned activity (2 seconds) for this metric.• Intra-region co-smoothing. Predicting the activity of a held-out neuron using neurons in the same brain region.• Inter-region co-smoothing. Predicting the activity of a held-out neuron using neurons in other brain regions. This is similar to the leave-one-out region validation from [10] .• Choice decoding. Predicting the choice the mouse makes using trial-aligned neural activity.• Motion energy decoding. Predicting motion energy of the mouse's whiskers using trialaligned neural activity. The motion energy is extracted from simultaneous video data.We evaluate each unsupervised activity prediction metric for all neurons in a session. To evaluate activity prediction, we utilize the co-bps metric introduced in [25] which measures the performance of a model using bits per spike; see the Appendix for a detailed definition. For behavior prediction, we train a linear classifier on the output firing rates of each model. To predict choice, we use the spiking activity of all neurons across all timesteps. To predict motion energy at each timestep, we again use the neural activity across all timesteps. We use accuracy as the metric for choice decoding and the R-squared metric to quantify the proportion of variance explained for decoding motion energy.By design, there is a correspondence between the metrics introduced in 4.2 and the novel masking schemes introduced in 3.1. The correspondences are as follows: neuron masking and co-smoothing, causal masking and forward prediction, intra-region masking and intra region co-smoothing, and inter-region masking and inter-region co-smoothing. For each evaluation metric, we propose a masking scheme that should lead to good downstream performance. By alternating between these masking schemes during training and utilizing a learnable prompt token at inference time, MtM should be able to generalize well to these different evaluation tasks.For choice decoding and motion energy decoding, the correct prompt token at inference time is unknown. For our experiments, we utilize the neuron masking prompt token for choice decoding and causal masking for motion energy prediction. We found, however, that the choice of prompt token for behavior decoding is not that important for good downstream performance (see Appendix C).Activity prediction benchmarking is traditionally performed using a pre-fixed set of held-out neurons [25] . Models learn to predict the same held-out neurons using the held-in neurons during training. While this evaluation scheme works for a fixed held-out dataset, evaluating activity prediction on all neurons and on all brain regions would require training hundreds to thousands of individual models.Since the goal of a foundation model is to have a single model perform well across all metrics, we argue this evaluation method must be changed.Therefore, we introduce a novel test-time evaluation scheme for benchmarking foundation models of neural spiking activity. During training, all models are trained on the full neural population activity with model-specific learning schemes. During evaluation, we zero mask different parts of the input data to construct held-out subsets of data to evaluate the model. We utilize this test-time masking scheme to compute all the activity prediction metrics introduced in Section 4.2. Due to the masking of sub-token channels during co-smoothing, inter-region and intra-region tasks, we attempted to concatenate a binary mask to the model input to indicate which channels were masked. However, this approach did not yield any advantages for NDT1 compared to zero masking. We provide further details on this in the Appendix.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We include the training details in Appendix D and include the seed, model architecture, and hyperparameter choices in the attached code.Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
architectures: 
 For all of our experiments, we implemented two existing transformer architectures designed for neural population activity. For single-session evaluation, we re-implemented NDT1 [41] and NDT2 [42] . For multi-session evaluation, we re-implemented NDT1-stitch and NDT2 (described in Section 2.2). NDT1-stitch learns, for each session, a linear projection layer for embedding the neural activity vector. We also include a session-level context embedding for NDT1-stitch as we found it improved multi-session training. NDT2 utilizes a ViT-style [7] transformer architecture and a learned sessionlevel context embedding. We set the patch size of NDT2 to be 128 neurons as, with ∼676 neurons on average per session, smaller patch sizes were prohibitively slow to train. For more details about the architectures and re-implementation details, see Appendix D. Masking scheme ablation. To understand the importance of each masking scheme, we first train and evaluate all architectures on a single session using each individual masking scheme introduced in Section 3.1. We also train and evaluate a MtM model with and without prompting. As discussed in Section 4.3, each masking scheme should perform best on the associated metric and MtM should perform well across all metrics.MtM vs. temporal masking. For a comprehensive evaluation of the MtM approach in comparison to temporal masking, we then train the single-session architectures with temporal masking and MtM across all 39 sessions (animals). Each model is trained for 1000 epochs, with the best checkpoint selected for evaluation based on the highest average single-neuron reconstruction R 2 across the 50 most active neurons in a session. We fix the learning rate and model architectures for all experiments. We evaluate the trained models on all metrics introduced in Section 4.2. For additional details about the hyperparameters and baselines, see Appendix D. For activity reconstruction, we report the average bps. For choice and whisker motion energy decoding, we report the average accuracy and R 2 , respectively, across all test trials. We use the NDT1 architecture for all comparisons.MtM vs. temporal masking pre-training. We are also interested in evaluating the performance of MtM, in comparison to temporal masking, for multi-session pretraining. Although the identity of the neurons is changing across sessions, we hypothesize that by training across multiple IBL repeated site datasets, which share anatomical structure, our MtM-based approach should generalize better to unseen IBL repeated site sessions. To this end, we pretrain all multi-session architectures using MtM and temporal masking on 10 and 34 sessions of data. We then evaluate these multi-session models on 5 held-out sessions by fine-tuning their self-supervised loss (MtM or temporal) on each held-out session. This allows the models to learn session-specific information such as the session embeddings. We report all metrics across these 5 held-out sessions which include 397-579 trials in the training split, 57-83 trials in the validation split, and 114-166 trials in the test split.Behavior decoding from individual brain regions. To evaluate how well our fine-tuned, 34session pre-trained MtM can generalize to unseen tasks, we also perform behavior decoding using single brain regions from the 5 held-out sessions. To perform behavior decoding using a single brain region, we mask out all other regions and predict the rates of the neurons in the specified region. Then, we train a linear model on these output rates to predict choice and whisker motion energy. For MtM, we prepend the intra-region prompt token as it is associated with this downstream task. We compare behavior decoding results for our 34-session pre-trained MtM to the 34-session pre-trained temporal masking model across all brain regions.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We use public datasets for our experiments. Our code is attached for reproducibility of our results.• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
results: 
 For all results in the main text, we utilize NDT1 and NDT1-stitch architectures. We found that the NDT1 architecture outperformed the NDT2 architecture across all metrics and sessions for both temporal masking and MtM. We hypothesize that the patching scheme of NDT2 does not generalize well to spike-sorted, multi-region Neuropixels recordings which have up to ∼676 neurons on average per session. All NDT2 results are reported in Appendix B and C. These results have similar trends (with overall lower metric scores) as the NDT1 results.Masking scheme ablation. The results for the masking scheme ablation on a single session are shown in Table 1 . We report the neuron-averaged activity metrics for all masking schemes, including the co-smoothing, forward prediction, intra-region, and inter-region activity prediction. As shown in the table, each masking scheme leads to an improvement on its associated metric (see Section 4.3) over the temporal masking baseline. This is a promising result as it illustrates how each masking scheme can teach the model about a different aspect of the neural population. Our MtM method also shows strong improvements in activity prediction across all 4 metrics in comparison to the temporal baseline. Although the activity prediction results of the single masking schemes can sometimes outperform MtM on the associated metric, the overall performance of MtM across all metrics is high. This demonstrates how training with diverse masking schemes can lead to a more structured understanding of neural activity. Finally, we show that MtM with prompting is a modest (in 3/4 metrics) improvement over MtM without prompting. Overall, this masking scheme ablation demonstrates the strength of our MtM approach for structured learning of neural data.MtM vs. temporal masking. We show results for our comparison of MtM to temporal masking across all 39 sessions in Figure 2 . As seen in Figure 2 , our MtM training approach leads to significant improvements across all 4 unsupervised activity prediction tasks. The largest improvements of MtM over temporal masking are for intra and inter-region activity prediction, as temporal masking is unable to learn this structure. For behavior decoding, we find that MtM and temporal masking have comparable results for choice decoding and MtM slightly outperforms temporal masking on whisker motion energy prediction. As we are using the full population of neurons for behavior decoding in these analyses, the similarity in results between MtM and temporal masking is unsurprising given that the temporal masking is performed on the full population for each time step. We hypothesize that when decoding behavior from individual brain regions, the MtM approach should outperform temporal masking as it learns brain region-specific structure (see Figure 5 ). These single session results demonstrate that MtM is a promising method for learning population-level, region-level, and single-neuron-level dynamics from neural population data. MtM vs. temporal masking pre-training.We report results for multi-session pretraining using MtM and temporal masking on 34 sessions of data in Table 2 and Figure 3 . In Table 2 , we show session-averaged results on the 5 held-out sessions for both the single-session and 34-session pre-trained MtM and temporal masking. For both single-session and multi-session, MtM outperforms temporal masking across all metrics except choice decoding (where the results are quite similar). Both methods benefit from multi-session pre-training as all unsupervised and supervised metrics improve for the 34-session pretrained models. Similar to the single-session results, the biggest improvements for MtM are for the unsupervised activity metrics especially inter-and intra-region prediction. In Figure 3 , we show a scatter plot of all metrics for the 5 held-out datasets for MtM and the temporal baseline. MtM shows improvement over the temporal baseline for all activity metrics. The results of our scaling analysis (1, 10, and 34 session training), can be seen in Figure 4 . For MtM, scaling the number of pretraining sessions leads to improved performance on all downstream metrics except for choice decoding which saturates at 10 sessions. For temporal masking, the performance of co-smoothing, intra-region, inter-region, choice decoding, and motion energy prediction saturates at 10 session pretraining. These results demonstrate that the MtM is a more promising approach than temporal masking for scaling neural population models to multi-animal, multi-regional datasets.Behavior decoding from individual brain regions. The results for behavior decoding using single brain regions on the 5 held-out sessions are shown in Figure 5 . We provide the chance level for each session in the Appendix. From these results, one can see that MtM provides a significant improvement on brain region behavior decoding in comparison to temporal masking across both choice decoding and whisker motion energy prediction. This is especially apparent in hippocampus brain areas (orange) where for session 5dee0eb, the temporal masking baseline is below chance-level on choice decoding for many sub-areas, but MtM is above chance level on all sub-areas. These results illustrate the ability of MtM to extract region-specific information from multi-region neural populations. Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We specify all the training and hyperparameters details in Appendidx D and D.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
discussion: 
 In this work, we take an important step toward a foundation model of the brain at single-cell, singlespike resolution. We introduce a novel approach to self-supervised learning, multi-task-masking (MtM), which is able to learn population-level, region-level, and single-neuron level structure from neural population activity. We validate this approach on neural recordings taken from a large, diverse dataset (39 mice across multiple brain regions). We find that MtM significantly outperforms current state-of-the-art masking modeling schemes for neural data prediction, and enables multi-task generalization. We also provide scaling results demonstrating that MtM is a promising approach for large-scale pre-training on neural population data across animals and sessions.A number of limitations remain. First, the data diversity, while higher than other neural pre-training datasets [3, 42] , is still significantly lower than a full brain-wide map of neural population activity [17] . Training MtM on the IBL brain-wide map dataset, which contains 300 brain regions and hundreds of animals, is an exciting future direction. Second, our current architectures, NDT1 and NDT2, are simple time-series transformers that utilize basic tokenization schemes, i.e. temporal tokens. Recent work has demonstrated that learning global dependencies across temporal tokens leads to poor forecasting results on multivariate time-series datasets [21] and that more sophisticated architectures can outperform the NDT architectures on behavior decoding from neural data [3] . Therefore, improving the underlying architecture of MtM is another direction that should be explored. region, inter-region, as well as the proposed MtM method with and without the prompt token on a selected single session. For both NDT1 and NDT2, Tables 4 and 5 show that the prompted MtM model outperformed the baseline temporal (random token) masking model on most evaluation tasks, except for choice decoding. In addition, each masking scheme performed well on its corresponding task. Figure 9 compares the baseline random token masking NDT2 to the prompted MtM NDT2 across 39 single sessions. The prompted MtM NDT2 consistently outperformed the baseline random token masking NDT2 on all activity reconstruction tasks, while performing similarly on behavior decoding tasks. Binary mask token ablation for NDT2. To optimize NDT1 and NDT2 for predicting held-out neurons, sub-token channels are masked during co-smoothing, inter-region and intra-region tasks, except for forward prediction. Initially, a binary mask was concatenated to the input of NDT1 to indicate which channels were masked, but this approach did not show benefits, leading to the use of zero masking instead. We also tested concatenating a binary mask with NDT2, which showed a slight improvement in Table 6 , though overall performance metrics for random token masking remained negative. Single neuron evaluation. To better understand neural activity prediction performance at a singleneuron level, we conducted an evaluation of single-session NDT1 on each neuron using bits per spike (bps), R 2 between the ground truth and predicted peristimulus time histogram (PSTH R 2 ), and the single-trial R 2 averaged across all trials (trial average R 2 ), on one session, in Figure 10 . We find that MtM outperforms the temporal baseline across most neurons regardless of firing rate.We did find a strong correlation between the performance evaluated on each metric and the mean firing rates of each single neuron. For the bps (bits per spike) metric, scores for active neurons tend to be more concentrated, while scores for inactive neurons are relatively dispersed, exhibiting both extremely low and high values. For both R 2 metrics, the performance shows a positive correlation with the mean firing rates. In particular, those neurons with extremely low mean firing rates typically exhibited R 2 scores that were extremely low (approaching zero). This observation might be related to the inherent difficulty in predicting the behavior of neurons with low mean firing rates, or the property of metrics themselves. For instance, when applied to neurons with low mean firing rates, the R 2 metric might tend to yield values closer to zero. Across all three metrics (Bps/Trial average R 2 /PSTH R 2 ), our model demonstrated substantial improvements for neurons with relatively higher mean firing rates. However, for neurons with lower mean firing rates, notable improvements were only observed in the bps metric. For NDT2, we report results for pretraining using MtM and the baseline random token masking on 34 sessions of data in Table 7 and Figure 11 . In Table 7 , we show session-averaged results on the 5 held-out sessions for both the single-session and 34-session pre-trained MtM and random token masking. For both single-session and multi-session, MtM outperforms random token masking across all metrics except choice decoding (where the results are quite similar). Both masking schemes benefit from multi-session pre-training as all unsupervised and supervised metrics improve for the 34-session pretrained models. Similar to the single-session results, the biggest improvements for MtM are for the unsupervised activity metrics, especially inter-and intra-region prediction. In Figure 11 , we show a scatter plot of all metrics for the 5 held-out datasets for MtM and the random token baseline.Table 7 : Fine-tuning performance of NDT2 pretrained with MtM vs. random token masking on activity reconstruction and behavior decoding. Activity reconstruction performance is reported in neuron-averaged bps. For behavior decoding, trial-averaged accuracy (R 2 ) for choice (whisker motion energy) decoding is shown. All metrics are averaged across 5 held-out sessions, and a higher value indicates better performance. Scaling analysis. To examine NDT2's ability of scaling data, Figure 12 shows that NDT2 multisession pre-training also benefits from scaling from 1 to 10 sessions, but we only observe marginal gains or no improvement going from 10 to 34 sessions. NDT2 (Figure 12 ) benefits less from multisession IBL pre-training compared to NDT1-stitch (Figure 4 ), likely due to the inability of NDT2 with a fixed patch size to handle the large neuron count variations (200 to 1000 neurons) across IBL sessions. Another reason is the NDT2 takes too many neuron patches at the same time, and it's very challenging to deal with this long sequences when data is scaling.Prompting mode ablation. We also conducted ablation studies on NDT1-stitch of different prompt mode during behavior decoding. We only apply prompt during the model inference, and observe the different prompt effects to our behavior results. As shown in the Table 8 , we ablate four prompt modes.Table 8 : Evaluation of NDT1's behavior decoding performance when pretrained and fine-tuned using the MtM approach, tested with different prompt tokens at inference time. Behavior decoding is assessed using trial-averaged accuracy for choice and trial-averaged R 2 for whisker motion energy, with the reported metrics averaged over 5 held-out sessions. For both metrics, a higher value indicates better performance. Architecture details. We created a schematic for the NDT1 and NDT2 architectures in Figure 13 to clarify how these architectures tokenize and reconstruct neural data. NDT1 uses the timestep for the positional encoding. NDT2 uses 2 types of position encoding: (1) timesteps and (2) neuron group identity. Training details. The 34-session NDT1 model has 25.55 million parameters, many of which come from session-specific stitchers. However, these stitchers use only one linear layer per session, making the model capacity equivalent to that of a single-session NDT1 model. We do not provide model sizes for the single-session NDT1, as their parameter counts depend on the number of neurons in each session. Both the single-session and 34-session NDT2 models do not require session-specific stitchers and consist of 1.09 million parameters. We trained our model using AdamW optimizer [22] for 1000 epochs with a learning rate of 1e -4 using a cosine scheduler. We put a weight decay 0.01 to avoid overfitting. We utilized a relatively small batch size of 16 during the training. We split our dataset based on the session to training, validation, and test set with a proportion of 70%, 10%, and 20%. We saved the model checkpoint based on a trial-average R 2 of top-50 active neurons, which we selected top-50 active neurons and calculated each neuron's R 2 through averaging the trials.Compute. NDT1-stitch was trained on a machine with a single RTX8000 GPU. NDT2 was trained using Tesla V100 GPUs with 32Gb memory. The 10-session and 34-session NDT1-stitch models were trained for 1 and 3 days, respectively, while the 10-session and 34-session NDT2 models took 2 and 5 days, respectively. Single-session NDT1 and NDT2 models, as well as finetuning, were trained on a single Tesla V100 GPU, requiring 3 to 6 hours. We make sure our experiments are reproducible by seeding.Hyperparameters details. The masking ratio is an important model hyperparameter for NDT1 and NDT2. For neuron, intra-region, temporal, and our proposed MtM masking schemes, the masking ratio is fixed at 30%, favoring the performance of the baseline temporal masking method across the activity reconstruction tasks. The causal (next timestep prediction) and inter-region (mask whole region) schemes do not have this hyperparameter, making their performance invariant to the selected mask ratio.For NDT2, the spatiotemporal patch size is another important hyperparameter. Due to computational constraints, we set it to 128 neurons (≈ 1000 tokens). Future work should analyze how varying the patch size impacts NDT2's performance on neural activity reconstruction and behavior decoding tasks. Table 9 : Effects of masking ratio on NDT1 performance in neural activity reconstruction. The reported metrics quantify performance in terms of average bits per spike (bps) across all neurons from a selected session. A higher bps value indicates better performance. Ablation study on positional embeddings. We perform an additional experiment to evaluate the impact of the positional embedding on MtM's performance for different downstream tasks. In Table 10 , we compare the effects of RoPE and learnable positional embeddings on the NDT1 architecture pretrained using the MtM method. The results demonstrate that RoPE consistently outperforms learnable positional embeddings across all downstream tasks. Co-smoothing evaluation metrics. We use bits per spike as a metric to evaluate model performance across co-smoothing, forward prediction, and both inter-region and intra-region tasks. The bits per spike for a neuron is defined as follows:bits/spike = 1 n sp log 2 (L(r; Xn,t) -L(rn,:; Xn,t)),where Xn,t represents the neural activity of the target neuron n at time t, r is the inferred firing rate from the model, rn,: is the mean firing rate for neuron n, and nsp is its total number of spikes. We then average the bits per spike across all neurons in a test session to get the final performance.The bits per spike metric is closely related to the deviance, a standard metric generally used in Statistics, e.g., generalized linear models (GLMs) use deviance as a goodness-of-fit measure for a statistical model. The deviance compares the goodness-of-fit of the model of interest, e.g., MtM, to a baseline null model, where the goodness-of-fit is measured by the model log likelihood. The bps is a normalized version of the deviance metric, which compares the model predictions to the average firing rate of the neuron for the trial. The bps further normalizes the deviance by the spike count so that the metric can be comparable across neurons regardless of whether the neurons are active or inactive.Chance level for behavior decoding. For binary choice, we report the test accuracy that will be reached when constantly predicting the train set's majority class in Table 11 . 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 R 2 from selected regions are generally larger than these chance levels. • The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: We didn't report error bars, but our experimental results are averaged across a large number of sessions (datasets), which ensures the reproducibility.• The answer NA means that the paper does not include experiments. • The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We document the computing resources used in Appendix D Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We clearly indicate our cited existing assets by citing the work.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA] Justification: The paper does not release new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA] Justification: Our paper does not involve any human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?38th Conference on Neural Information Processing Systems (NeurIPS 2024). 
 
 Appendix:  


Decision: This is a solid paper building towards a foundation model for neural recordings at high resolution (single spikes of single cells). Trained on neuropixel data from many animals and sessions, it enables the prediction of neural activity and behavior, and the generalization to new animals. The reviewers asked for many additional analyses and clarifications that the authors provided and that should be included in the final manuscript.

Reviews: review: paper_summary: paper_summary: The paper presents a large-scale model pretrained on the International Brain Laboratory (IBL) corpus containing multi-region, multi-animal spiking activity of mice during a decision-making task. It introduces a self-supervised learning approach with a novel masking scheme called MtM which alternates between masking and reconstructing activity in timesteps, neurons and brain regions. The masking strategy was shown to be helpful in learning different structures in neural population activity and outperforms temporal masking scheme of previous models. The paper shows positive scaling results on held-out animals as the number of training animals increases.
strength_weakness: strength_weakness: 1.	The number of baselines is quite limited. Only two transformer architectures (NDT1 and NDT2) were benchmarked, and the proposed MtM method is mainly compared with one baseline of temporal masking. It appears that MtM could be used as the training objective for similar existing models that also assume a Poisson emission model, e.g. RNN-based LFADS (Pandarinath et al. 2018) and other models in the NLB benchmark (Pei et al. 2021). Evaluating MtM on more architectures could be helpful to gauge if MtM effectiveness is truly model-agnostic as the authors noted or is biased more toward certain type of architecture. 
2.	The dataset only consists of one behavior task (decision making), while existing pretrained models POYO and NDT2 referenced in the paper were trained/evaluated on multiple behavior (motor) tasks, even across species (monkey vs. human). It is unclear if and how the community could benefit from finetuning the proposed pretrained model on other datasets with a different behavior paradigm.
3.	Some experiment details are missing or not presented clearly (elaborated in Questions section below).
questions: questions: 1.	It would be helpful to compare the proposed model with more NLB baselines to justify the architecture choice. Does NDT1 with MtM significantly outperform simpler baselines like RNN with MtM to justify the computation cost? Even knowing how well baselines like MLP at directly decoding the behavior could help readers better evaluate the strength of the proposed model.
2.	It is not clear if the pretrained model would provide some leverage on other datasets with different behavior structures. Would it outperform a model trained from scratch on the new dataset?
3.	It might be worth making it clear that “task” in “multi-task” used in the paper refers to the machine predictive task (decoding choice and motion energy), rather than animal behavior task (decision making – the only behavior investigated). 
4.	Equation (1): Does the transformer use any kind of positional encoding? If so, what type of positional encoding? 
5.	For finetuning on a held-out session, I assume the reported metrics were calculated on the test trials within the session. How many trials are there in each train/val/test split of each session? 
6.	Line 174: possibly a typo, should be 1% of the trial (20ms).
7.	How was choice decoded from $\hat{X}$? I assume from Figure 1 that $\hat{X}$ is a N x T matrix, where N is number of neurons and T is number of timestep. If so, at which timestep of $\hat{X}$ that the linear layer was applied to decode the choice?
8.	Similarly, was the motion energy decoded at each timestep of $\hat{X}$ or at some particular point during the trial?
9.	Figure 5: what is the chance level for each session? For imbalance datasets, the chance level might not be 0.5, especially if the data was preprocessed for quality control in some ways. If so, it’s worth plotting the chance level in the figure.
limitations: limitations: The authors have addressed the limitations in the manuscript.
paper_summary: paper_summary: The authors propose a self-supervised approach to building a foundation model of single-trial neural population dynamics. The approach utilizes existing transformer architectures, which are trained using “multi-task-masking” (MTM), which alternates between several scales of prediction tasks, including co-smoothing, causal prediction, inter-region prediction, and intra-region prediction. The approach is demonstrated on the IBL repeated site dataset, where activity prediction, forward prediction, and behavior decoding using a fixed architecture are improved using MTM training relative to a baseline where training only utilizes temporal masking.
strength_weakness: strength_weakness: The paper would benefit from:
- Comparing to a wider range of models of neural population dynamics (e.g., LFADS, GPFA, mDLAG, etc). Many previous approaches have reported co-smoothing performance (e.g., Neural Latents Benchmark). How does MTM-NDT1 compare against those previous approaches on co-smoothing? How does the proposed approach compare to multi-area models (e.g., mDLAG) for on the inter-region co-smoothing task?
- More fully describing the details of the “temporal masking” baseline. The details are referenced out to a citation rather than described within the text. Since this baseline appears throughout all of the results, it seems important to describe it within this paper.
questions: questions: - What is the motivation for building a foundation model for neural data? What are potential applications for such a model? Can such a model teach us anything about the brain?
limitations: limitations: The authors adequately state the limitations of their modeling approach, training paradigm, and datasets.
paper_summary: paper_summary: This work proposes a transformer architecture (based on previously existing ones) and, more interestingly, a training procedure that, when applied to spiking neural data, should result in a foundation model for spiking neural data. The clever bit about the procedure is that the learning model is asked to reconstruct differently-masked portions of the data: activity of one held-out neuron, future activity of a neuron, activity of a neuron based on neurons in the same region, activity of a neuron based on activity of neurons in other regions. During training these tasks are presented to the model interchangeably, so it has to learn to do all of them (although with a help of an embedding vector that can decode for the type of the task).

The results show that the proposed model achieves better numbers in a variety of tests, however I found it hard to evaluate the significance of the result using the reported metric of bits per spike and how it translates to real-world performance. Please see my comments below and I hope to learn more about this from authors' rebuttal.
strength_weakness: strength_weakness: * I would like to see more about the architecture of the transformer (even though it is based on existing model), ideally a more detailed visualisation, this would help to imagine how exactly the data is going in, how the masking is applied, etc.

* I am not too happy about the bit-per-spike metric being the only one reported. It is fine to have one number to be able to characterize and compare the performance succinctly, but it's awkward to translate back to "how well does it work, really?". I think it would be great if the paper would include actual raster plots of good / average / worst reconstructions of the masked neuron / future / intra-region / inter-region regions -- this would allow a neuroscientist to decide with a quick glance whether "wow I should start using this" or "nah I am gonna wait till NeurIPS 2025"
questions: questions: 184: Please include a formal definition of "bits per spike"

Table 1: How a "bits per spike" measure can be negative? What does it mean?

Figure 5: Please explain "Choice" and "WME" and axis units in the caption and/or on the figure. Perhaps it would be better to fix the range on the axis, so that one could easily estimate also the variability between the sessions and MtM improvement in the sessions.
limitations: limitations: The authors have adequately highlighted some limitations, but I would like to hear a bit more about the limitations of the reconstructed signal from neuroscientific perspective: are they usable? Do they brake or maintain some important characteristics of spiking signal?
paper_summary: paper_summary: This work develops a multi-task-masking (MtM) approach, based on a self-supervised Transformer, that masks and reconstructs activity across different dimensions for neural spiking data learning. Evaluated on the International Brain Laboratory dataset, the model improves tasks such as single-neuron and region-level activity prediction, forward prediction, and behavior decoding. This method enhances performance and generalization, advancing towards a comprehensive foundation model for the brain.
strength_weakness: strength_weakness: 1. Lack of Qualitative Analysis:

There is a lack of qualitative analysis regarding the neural dynamics obtained from population neural activity across different individuals, such as visualizing features output by NDT. Are there similar patterns? Are there similarities between single-neuron and population-level activities?

2. Insufficient Ablation Studies:

a) What are the differences between training on multi-region data versus single-region data? When training with data from three brain areas, is there a phenomenon where the activity of one brain region dominates?

b) What are the performance differences when using multi-session data from a single individual versus cross-individual multi-session data?
questions: questions: 1. Typically, cross-individual modeling aims to explore common patterns among different individuals, while cross-area modeling aims to understand information interactions between different brain areas. How much do you think the current foundation model can contribute to these explorations?

2. Generally, in neuroscience, there is a greater interest in understanding the interactions with other internal brain regions, which would be helpful for functional modeling in network level. What insights can the current foundation model provide for research in these areas?
limitations: limitations: see Weaknesses.

Comments: title: Response to Reviewer BBJQ
Response to Reviewer JcgB
Response to additional questions from Reviewer xe9o - Part 2
Response to additional questions from Reviewer xe9o - Part 1
Applying MtM to interpretable models of neural activity (in Global Response)
comment: Thank you for the explanation. The extra experiments alleviated my concerns. I've increased my rating accordingly.
We hope the reviewer found our rebuttal helpful. If there is anything we can further clarify, we are happy to do so!
We thank the reviewer for their positive response to our rebuttal. If there is anything we can further clarify or do to improve the reviewer's overall assessment of our work, please let us know.
Thank you for the thoughtful rebuttal. I believe the final product will indeed be improved by incorporating the ideas from these reviews and rebuttals. I stand by my score.
Thank you for explanations and clarifications. They were helpful to alleviate some of my concerns and I will raise my overall score a bit. Good luck!
**Q3**


Before answer the question, we want to clarify that MtM is agnostic to the underlying transformer architecture and that the tokenization scheme will depend on the architecture that is used.

To answer the question, a token for NDT is the observed spike counts for a single 20ms timestep of the population (please see Figure 3 in [5]). This means that NDT is using the exact spike times (up to 20 ms resolution) to fit the model. The reviewer's suggested tokenization scheme would be comparable to NDT with more fine binning (1 ms bins). While this is feasible, given the sparsity of neural data, this would give rise to many more empty bins and would also be computationally prohibitive (i.e., self-attention scales quadratically with the number of tokens). As the reviewer mentioned, "why use bins at all?". Treating each spike as a token, as done in the POYO paper [6], is also a potential solution. As this is computationally prohibitive, they adopted a perceiver IO architecture [7] so the attention is only applied to a set of latent tokens, which is much lower dimensionality. While this architecture is a promising step in the right direction, it has only ever been used as a supervised model to decode behavior from spikes. Adapting POYO and more interesting transformer architectures to be self-supervised such that they can be used with MtM is an exciting future direction.


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

[6] Azabou, Mehdi, et al. "A unified, scalable framework for neural population decoding." NeurIPS 2024.

[7] Jaegle, Andrew, et al. "Perceiver io: A general architecture for structured inputs & outputs." arXiv, (2021).
We thank the reviewer for getting back to us so promptly! We agree that Figure 1 should be a part of the manuscript and we will include a version of this figure in the main text of the updated manuscript. We believe there might be a misunderstanding about the goals of statistical models of neural data that we hope to clarify in the following response. Our response will be in two parts due to character constraints (we will answer Q2 and Q1 in the first comment). If we answer all questions satisfactorily, we hope the reviewer would consider raising their score.

**Q2**

We are glad the reviewer found our definition of bps helpful and we will make sure to add this in the updated manuscript. We suspect there might be a misunderstanding here about prior literature on neural population models. As described in [1], there is a rich literature of population models that utilize either Gaussian or, more recently, Poisson observation models (please see Table 2 of [1]) to predict neural spiking activity. The goal of neural population models is to predict the underlying firing rates of each neuron because the exact spike times are believed to be observations from a stochastic process. Similar to this seminal paper [2], we choose to model the observed spikes as samples from an inhomogeneous Poisson process whose rate corresponds to the inferred firing rate for the given neuron. To measure the quality of fit of our Poisson process, it is most natural to use a version of the log-likelihood such as the bps metric utilized in the Neural Latents Benchmark [3]. Using this commonly agreed upon metric makes sure that our work is consistent and comparable with existing research on neural population models. 

In addition, the "bits per spike" metric is closely related to the "deviance", a standard metric generally used in Statistics, e.g., generalized linear models (GLMs) use deviance as a goodness-of-fit measure for a statistical model. The deviance compares the goodness-of-fit of the model of interest, e.g., MtM, to a baseline null model, where the goodness-of-fit is measured by the model log likelihood. The bps is simply a normalized version of the deviance metric, which compares the model predictions to the average firing rate of the neuron for the trial. The bps further normalizes the deviance by the spike count so that the metric can be comparable across neurons regardless of whether the neurons are active or inactive.  

Along with the bps, we would be happy to report the R-squared in the revision to provide further intuition, but note that the R-squared depends strongly on bin size (as bin size gets small so will R-squared, since the precise timing of spikes is likely not predictable in this setting). 

**Q1**

In Figure 1 of the one page PDF, the ground-truth rastermaps are the spike counts of the trial while the MtM/NDT rastermaps are the inferred firing rates. These quantities are not directly comparable because we expect the underlying firing rate of a neuron to vary smoothly while its observed spikes (which are assumed to be stochastic) may have sharp increases. A great example of this phenomena is shown in [4] where they demonstrate that when the inferred firing rates exactly fit each spike time, this is actually overfitting to the neural data. They introduce a novel dropout strategy, similar to MtM's neuron masking, that reduces this overfitting. 
In fact, it has been shown by [2] (in Figure 2e) that the smoothed firing rates from a Poisson model actually allow for decoding animal behavior more accurately than the raw spike counts. Whether to smooth or not involves a bias-variance tradeoff. More smoothing means more bias and less variance. When we optimize for prediction accuracy, we are implicitly optimizing this tradeoff, and empirically we (and others) find that some smoothing (bias) leads to better predictions. 

[1] Hurwitz, Cole, et al. "Building population models for large-scale neural recordings: Opportunities and pitfalls." Current opinion in neurobiology, 2021.

[2] Pandarinath, Chethan, et al. "Inferring single-trial neural population dynamics using sequential auto-encoders." Nature methods, 2018.

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

[4] Keshtkaran, Mohammad Reza, et al. "A large-scale neural network training framework for generalized estimation of single-trial population dynamics." Nature Methods, 2022.
We thank the reviewer for responding back to us so quickly. We believe that MtM can easily be combined with more interpretable models of neural populations. For example, as part of our Global Response, we combined our MtM learning objective with LFADS, an interpretable, explicit dynamics model. We were able to improve the performance of LFADS on all predictive tasks introduced in the paper. Although we have no more space for figures, interpreting the latent factors learned by LFADS when trained with MtM is an exciting future direction and something we can do in the updated manuscript. We hope the ability of MtM to be applied to interpretable latent factor models like LFADS can improve the reviewer's score of our manuscript.

Although explicit dynamics models are very important for neuroscience, we still believe that purely predictive models like transformers are also useful. Neural decoding/encoding algorithms have been developed for many years to determine how much information about a stimulus or behavior is encoded in a neural population [1]. Much like these algorithms, we are interested in determining how much information each brain area has about other brain areas. Understanding which brain areas predict each other best can allow us to make hypotheses about distributed neural computation.

[1]​​ Paninski, Liam, Jonathan Pillow, and Jeremy Lewi. "Statistical models for neural encoding, decoding, and optimal stimulus design." Progress in brain research 165 (2007): 493-507.
The rebuttal PDF meets my expectations. Although the learning of neural dynamics is implicit, some patterns appear to be discernible. Please consider further joint analysis with interpretable methods in the future work, otherwise this kind of work is useless for neuroscience. I would like to maintain my score. Good Luck!
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 _ ?