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 Aaron Richards
Renee Tung
Olivier Winter
International Brain Laboratory
Eva L Dyer
Liam Paninski
Cole Lincoln Hurwitz
Venue: NeurIPS 2024 poster (2024)
Keywords: neural dynamics
transformer
self-supervised learning
pretraining
multi-task learning
masked modeling
brain-computer interfaces
Primary area: neuroscience_and_cognitive_science

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.

Towards a "universal translator" for neural dynamics at single-cell, single-spike resolution
Yizi Zhang'
Yanchen Wang'
Donato M. Jenéne
Zixuan Wang!
Blake Richards The International Brain Laboratory
Mehdi Azabou'
Renee Tung'
Olivier Winter
Eva Dyer" Liam Paninski
Cole Hurwitz
Columbia University 2 Universitat Politècnica de Catalunya
*Zhejiang University *Georgia Institute of Technology
MMila, McGill University
CChampalimaud Foundation
"The International Brain Laboratory
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. 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 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, 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-ant 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. Project page and code: https: // //tbl -mtm. gi gitthu..2.. io/
Introduction
Recent studies in experimental neuroscience suggest that neural computation is highly distributed across the brain and sparse 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-specine 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 [7] 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 [ 42 These models pre-trained on large corpuses of neural population data and then fine-tuned for downstream tasks such behavior decoding and brain-computer interface (BCI) control 42, leading improved performance and generalization to novel sessions and animals. While these models are a promising
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

step in the right direction, two crucial elements are currently missing, First. the vee-triiing is only performed on neural data from the sensorimotor network (MI, PMd, S1) which limits the applicabilitt of these approaches to whole-brain analyses. Second, these models do not explicitly model the underlying brain regions, instead treating the population as homogenous set of neurons. We argue that 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 diverse set of predictive tasks across diverse brain areas. Similar to 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 multi-task tastttasskknnn (MrM) 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 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 unsuper- vised 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) [4]] and the random token masking scheme used by Neural Data Transformer (NDT2) (4). We show that even with the same architecture, our MrM 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- trainininnn 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.
[8
The contributions of this work include:
A novel multi-task- mask--n (MrM) 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
2.1 Foundation models for neuroscience
The advent of large-scale, self-supervised foundation models has marked paradigm shift across various domains of machine learning. These models, diverging from traditional annotation-reliant supervised models, exhibit an impressive ability to generalize across spectrum of tasks. These models have transformed natural language processing [27 28 40), vision [29] [6], and robotics [30), and are beginning reshape the landscape of life sciences M. The application of foundation models to neuroscience is of significant interest. While there has been considerable progress in building large-scale scale models for EEG, [6, fMRI (36), and sEMG [9). 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 R and a self-supervised method, large corpus of spiking data from 12 monkeys. While promising, the
NDT2 [22, were trained 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.

---

baan
bonn
neurons
brin region
Task 1: Co-smoothing
Task 2: Causal Prediction
3: Inter-rgiion Prediction
4: Intra-region Prediction
Trakmus Hepocumpo Caree
F prompt
tokenization
Transformer model
Reconstructed
Figure 1: Schematic illustration of our Multi-task- Maali-task-Masking (MrM) approack: (A) introduce four metries for evaluating foundation models of neural population activity: ncuron co-smoothing, causal prediction, inter-region prediction, and intra-region prediction. For cach masking scheme, the colored area indicates what is masked and then reconstructed for evaluation. For intra-region prediction, the colored areas with hatched lines indicate areas which are masked, but not reconstructed for evaluation. Each metric be associated with specific masking scheme during training (TI, T2, etc.). (B) We alternate hetween different masking schemes during training along with a learnable "prompt" token which provides context to the model about the associated task [5). 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-agnie agnoosiic masking is performed on the input data (not the tokens). For a full discussion of MIM, see Section
2.2 Transformer architectures for neural population activity
Transformers 339 have recently been applied to neural population activity in both supervised 3 20 and supersueee 41 42 settings. POYO [3] a supervised multi-session model for predicting
behavior from neural activity. The POYO architecture utilizes learnable unit embeddings for each neuron and novel approach for tokenizing individual spike events with relative position encodings
to incorporate precise spike timings 3 For self-supervised learning, NDTI [41] and NDT2 3 are two existing transformer-basee methods. Both utilize masked modeling and assume Poisson emission model. For NDTI, each time bin is token with dimensionality equal to the number of neurons. NDTI 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 77] time-series transformer architecture that uses spatiotemporal patches tokens and learned session- context embedding. token defined as a single time step of neural activity for a subset neurons and is masked and reconstructed during training. By performing patching of the neurons, NDT2 can easily be applied to multiple sessions.
2.3 Multi-region models
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 electrophysiologica techniques now allow for recording neural activity across many interconnected regions simultancously This has inspired recent efforts to build more fine- ime-igrined estimates of multi-region and to investigate neuron-level information processing across multiple 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 Recently, dCSFA [1] and DLAG D 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
13 33 43 37
communication 26 32 [2]
brain areas 17.
8 g
3

---

neural dynamics, they also make limiting assumptions on the structure of the communication signals and cannot be easily scaled many brain regions and sessions.
3 Methods
foundation model for neural activity must be able seamlessly "translate" across all spatial scales of the brain, including population-level. region-level, and single-neuron-level dynamics. To this end, introduce 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 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]
3.1 Masking schemes
Masked modeling for neural activity is typically performed by masking and then reconstructing activity in random time steps 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 D for visualization of these masking schemes).
While this masking scheme allows for learning temporal
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 indi- vidual neurons relate the activity of the full population. This is conceptually similar to "coordinated dropout" introduced in 15]
Intra-region masking. We randomly mask neurons in 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 randomly chosen brain region and then reconstruct them using unmasked neurons in other regions as context. This masking scheme allows the model 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.
3.2 Multi-task-masking (MtM)
To train using MtM, we randomly sample masking scheme M for each batch of neural data. We mask the input data according to the sampled masking scheme and then pass this tokenizer and transformer-basee
architecture. To preserve the temporal or spatial order of the data, we add positional embeddings PE to the neural data tokens Z. To provide context the model about the masking scheme that was sampled, we prepend learnable "prompt" token P to the neural data tokens. This prompt token a D-dimensional learnable embedding which the model can use adapt its behavior during training or evaluation. For batch of neural data X, our training is as follows:
M ~ [ccuuall neuron, intra-region, inter-region)
Z Tokenizer(M C X)
Zpss = 2 + PE
(I)
Z prompp [P, Zpos]
r Transformer( (acineerrZpennn Zprompe)
~ Poisson( X |r)
where X the neural data reconstructed by the model, based on the time-varying rates inferred from the transformer's output, assuming 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 NDTI and NDT2 for all comparisons to these methods (see Section 2.2 for details).
3.3 Prompt-based test-time adaptation
Depending on the downstream task, the information learned using each masking scheme might be more less useful. We utilize prompt-based 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.
Evaluation
4.1
Dataset
For training and evaluating models, we use the International Brain Laboratory repeated sites dataset [18). This dataset consists 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- makinng task. The probe was localized after the experiments using reconstructed histology and the brain regions were annotated. We utilize trial-aligned, spike- sorted data from 39/48 mice for our analyses. From these recordings, we have a total of 26,376 neurons for training and evaluation (~-776 neurons per session average). We bin the neural activity using 20ms windows and we fix the trial-length to seconds (200 time bins). For behavior decoding, exclude trials based reaction time outliers defined by the IBL brain-wick map [17
4.2 Metrics
We utilize 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 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 tral-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 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 trial- aligned neural activity. The motion energy extracted from simultaneous video data.
We evaluate each unsupervised activity prediction metric for all neurons in session. To evaluate activity prediction, we utilize the co-bps metric introduced in [25) which measures the performance of model using bits per spike; see the Appendix for detailed definition. For behavior prediction, 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 each timestep, the neural activity across all timesteps. We use accuracy as the metric for choice decoding and the R-squared metric quantify the proportion of variance explained for decoding motion energy.
again
4.3 Metrics and masking
By design, there is 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 inference time, MtM should be able to generalize well to these different evaluation tasks.
5

---

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 Cp.
4.4 Test-time neural activity masking
Activity prediction benchmarking is traditionally performed using 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 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 single model perform well across all metrics, we argue this evaluation method must be changed.
Therefore, we introduce 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 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 binary mask to the model input to indicate which channels were masked. However, this approach did not yield any advantages for NDT1 compared zero masking. We provide further details on this in the Appendix.
5 Experiments
5.1
Architectures
For all of our experiments. we implemented two existing transformer architectures designed for neural population activity. For megle-session evaluation, we re-implemented NDTI [4] and NDT2 [2). For multi-session evaluation, we re-implemented NDT1-stitch and NDT2 (described in Section
2.2 NDT1 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 session- level context embedding. We set the patch size of NDT2 to be 128 neurons 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
5.2 Single-session
Table 1: The performance of single-session NDTI trained with various masking schemes on neural activity reconstruction tasks. The metrics are in units of bits per spike (bps), averaged across all neurons in one session. A higher bps value indicates better performance.
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 MtM model with and without prompt- ing. As discussed in Section[4.]] each masking scheme should perform best on the associated metric and MtM should perform well across all metrics.
Activity Reconstruction
Masking
Co-
Forward Prediction Region Region
Intra-
Inter-
Smooth
Temporal (Baseline)
0.84
0.42
-0.20
0.57
Neuron
1.04
-0.21
-0.22
0.78
Causal
0.44
0.48
-0.36
0.23
MtM vs. temporal masking. For a compre- hensive evaluation of the MIM approach in com-
Intra-Region
-9.86
-2.97
0.32
-9.06
Inter-Region
0.92
0.01
-0.58
0.90
MtM (Not Prompted)
0.99
0.54
0.42
0.83
parison to temporal masking, we then train the single-session session architectures with temporal mask- ing 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 π?2 across the 50 most active neurons in 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
MtM (Prompted) 0.98
0.57
0.43
0.84
details about the hyperparameters and basclines, Appendix D

---

Presicion
Orann Tudh
Ning Pute
Ca Sesceth
Seplet
6.
Forvand Prndition
Interdeecir
14
15
Tirest
Threal
Tie Royse Proceiin
Tasi 6.00
6:22
65
Enaetine
Clokbe
Vedinn Fregy.
1.0
6.3
6.0
-22 0.5 1.
15
0.5 3.5
Sasaliie
Deautine
Figure 2: Comparison of the temporal masking baseline and our proposed MrM method on single-session data. (A) and (B) show trial-averaged raster maps of CAI for ground -ouud data, MIM, and the temporal baseline. (A) The predictions from MtM and the temporal baseline are after inter-region masking where neurons in CA1 are predicted from all other brain regions. We highlight two areas (1 and 2) where MtM shows qualitatively better predictions of activity. (B) The predictions from MtM and the temporal baseline after intra-region masking where all neurons in CAI are predicted from other neurons the same brain region. We again highlight two areas (1 and 2) where MtM shows qualitatively better predictions of activity. (C) Activity reconstruction and behavior decoding across 39 sessions for MIM and temporal masking. Each point represents one session. For activity reconstruction, we report the average bps. For choice and whisker motion energy decoding, we report the average accuracy and R2" respectively, across all trials. We use the NDTI architecture for all comparisons.
5.3 Multi-session
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 the neurons is changing across sessions, we hypothesize that by training across multiple IBI. repeated site datasets, which share anatomical structure, our 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 temporal) each held-out session. This allows the models to learn session-specifie 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.
approach should generalize better to
Behavior decoding from individual brain regions. To evaluate how well our fine-tuned, 34- session 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, 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 MrM, 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- pre-trained trained MtM to the 34-session pre- pre-trained temporal masking model across all brain regions.
6 Results
For all results in the main text, we utilize NDTI and NDT1-stitch architectures. We found that the NDTI 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 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 aadd These results have similar trends (with overall lower metric scores) as the NDTI results.
6.1 Single-session
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 (sce Section 4.3
7

---

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 MIM with prompting is modest (in 3/4 metrics) improvement over MIM 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 As seen in F 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 MIM 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- region-level, and single-neuron-leve dynamics from neural population data.
6.2
Multi-session
Co-Smooth
Intra Region
Choice
1.5
1.00
MtM vs. temporal masking pre-training. We report results for multi-session pretraining using MtM and temporal masking on 34 ses- sions of data in Table 2 and Figure 3] In Table 2. we show session- -averan- results on the 5 held-out sessions for both the single-session 34-session pre-trained MtM and temporal mask- ing. For both single-session and multi-session, MtM outperforms temporal masking across all metrics except choice decoding (where the re- sults quite similar). Both methods benefit from multi-session pre-training as all unsuper- vised 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 met- rics especially inter- and intra-region prediction. for the 5 held-out datasets for MtM and the temporal baseline. MtM shows improvement over the temporal baseline for all activity metries.
Mew
1.5
1.0
0.5
0.75
0.5
-..50.0 0.5
0.75
1 1000
Whisker
Forward Prediction
Inter Region
Mation Energy
session and
MEN
1.5
1.0
0.5
0.50
0.0
0.50
0.75
Baseline
Baseline
Baseline
Figure 3: Fine-tuning performance comparison of NDT1- pretrained with MIM vs. temporal masking for activity reconstruction and behavior decoding across S held-out sessions. For activity reconstruction, each point shows the average bps across all neurons in held- out session. For behavior decoding, each point shows the tvial-averaged accuracy (choice) and Ie (WME). In Figure 3 we show scatter plot of all metrics
Table 2: Fine-tuning performance of NDTI-stitch pretrained with MIM vs. temporal masking on activity reconstruction and behavior decoding. Activity reconstruction performance is reported in averrnna bps. For behavior decoding, trial-averaged accuracy and R² are shown for choice and whisker motion energy, respectively. All metrics are averaged across 5 held-out sessions, and higher value indicates better performance.
Activity Reconstruction
Behavior Decoding
Training
Masking
Co-Smooth Forward Pred Intra-Region Inter-Region Choice Whisker Motion Energy
Temporal (Baseline)
0.55
0.17
-0.49
0.18
0.87
0.65
Single-Session
MtM (Prompted)
1.00
0.28
0.70
0.83
0.85
0.67
Temporal (Baseline)
0.79
0.43
-0.08
0.34
0.89
0.67
Multi-Session
MtM (Prompted)
1.11
0.46
0.85
0.96
0.88
0.68
8

---

Whrisker
Co Smooth
Forward Prediction
Intra Region
Inter Region
Choice
Metion Energy
Average Bits per Splke
0
1.0
0.70
0.99
0.4
0.8
RF
1.05
0.90
6.65
Averege
0.85
1.00
0.7
NET: Sasch
0.80
0.95
10
34
1
10
34
10
34
10
34
10
34
10
34
Avorogn Bits por Spite
0.4
0.4
0.1
Rccuracy
-0.2
Avarage
0.3
-0.3
-044
0.60
Stnch
-05
10
10 Sessions
30
10 34
36
Narm Secsions
Sessins
Sessions
Sessions
Sessions
Figure 4: Comparison of scaling curves between NDT1-stiteh pretrained with the MIM method the temporal masking baseline. The reported metrics avverne bits per spike (bps), choice decoding accuracy. and whisker motion energy decoding R² are averaged over held-out sessions. fine-tune each pretrained model with its supfrsuue loss (MtM or temporal) on the 5-heldout sessions and then evaluate with of our metrics. "Num of Sessions" denotes the number of sessions used for pretraining.
The results our scaling analysis (1, 10, and 34 session training), can be seen in Figure 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 a more promising approach than temporal masking for scaling neural population models to multi-an aniti- multi multi-regional datasets.
Behavior decoding from individual brain regions. The results for behavior decoding using single brain regions 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 brain region behavior decoding in comparison choice decoding and whisker motion energy prediction. This is especially apparent in hippocampus brain areas (orange) where for session 5deeOeb, 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.
temporal masking across both
Choice
Sdcee0eb
d7bf2d49
03d9a098
d57df551
824cf03d
Photi
0.8
0.8
509
3UB
0.4
0.9
APH
0A1
0.4
DG
0.25
0.50
0.75
0.9
1.0
0.50
0.75
0.50
0.75
0.7
0.8
WIME
0.6
0.8
kPPpp
ViSa
MtM
0.6
0.6
0.4
Pol
0.4
0.4
0.3,
0.25
0.50 Baseline
0.75
0.4
0.6
0.4
0.6
0.50
0.75
0.3
0.4
0.5
Baseline
Baseline
Baseline
Baseline
Figure 5: Comparison NDT1-stitch pretrained with the MtM method vs. the baseline temporal masking on behavior decoding from individual brain regions. The rows display choice decoding accuracy and whisker motion energy decoding H² Columns represent individual held-out sessions. Each point shows the behavior decoding performance when using neural activity from specific brain region, with colors denoting different brain regions.

---

7 Discussion
In this work, we take an important step toward foundation model of the brain at single-cell, single- spike resolution. We introduce a novel approach to self-supervised learning, multi-task-maat mask-na (MiM), 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 large, diverse dataset (39 mice across multiple brain regions). We find that MtM significantly outperforms current state-of-the-an art masking modeling schemes for neural data prediction, and enables multi-task generalization. We also provide scaling results demonstrating that MtM is 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 IBI. brain-wide wide map dataset, which contains 300 brain regions and hundreds of animals, is an exciting future direction. Second, our current architectures, NDTI and NDT2, are simple time-series transformers that utilize basic tokenization schemes, i.c. temporal tokens. Recent work has demonstrated that learning global dependencies across temporal tokens leads to poor forecasting results on multivariate time-series datasets [2]] 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.
Acknowledgments and Disclosure of Funding
This project was supported by the Wellcome Trust (209558 and 216324), National Institutes of Health (IU19NS123716), the Simons Foundation, the National Science Foundation (NSF award CIF:R1:2212182, NSF CAREER awards IIS IIS-2146072), NSERC (Discovery Grant: RGPIN-2020. 05105; Discovery Accelerator Supplement: 566355-2022) and CIFAR (Canada AI Chair; Learning in Machine and Brains Fellowship), and by DoD OUSD (R&E) under Cooperative Agreement PHY-2229929 (The NSF Al Institute for Artificial and Natural Intelligence), as well as generous gifts from the Alfred Sloan Foundation, the McKnight Foundation, and the CIFAR Azrieli Global Scholars Program.
Arthur B. MeDonald Fellowship:
References
[1] Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pages 1-3, 2024.
[2] Chris G Antonopoulos, Shambhavi Srivastava, Sandro E de S Pinto, and Murilo S Baptista. Do brain networks evolve by maximizing their information flow capacity? PLOS computational biology, 11(8)= 2015.
[3] Mehdi Azabou, Vinam Arora, Venkataramana Ganesh, Ximeng Mao, Santosh Nachimuthu, Michael Mendelson, Blake Richards, Matthew Perich, Guillaume Lajoie, and Eva Dyer. A unified, scalable framework for neural population decoding. Advances in Neural Information Processing Systems, 36, 2024.
[4] Matthew R Brier, Jewell B Thomas, and Beau M Ances. Network dysfunction in alzheimer's disease: refining the disconnection hypothesis. Brain connectivity, 4(5):299-311, 2014.
[5] Haotian Cui, Chloe Wang, Hassaan Maan, Kuan Pang, Fengning Luo, Nan Duan, and Bo Wang. segpt: toward building a foundation model for single-cell multi-omics using generative ai. Nature Methods, pages 1-11, 2024.
[6] Wenhui Cui, Woojae Jeong, Philipp Thölke, Takfarinas Medani, Karim Jerbi, Anand A Joshi, and Richard M Leahy. Neuro-gpt: Developing foundation model for ceg. arXiv preprint S3Ki::2311.07666 2023.
[7] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale, arXiv preprint 2010.1010.1190000 2020.
10

---

[8] Joshua Glaser, Matthew Whiteway, John P Cunningham, Liam Paninski, and Scott Linderman. Recurrent switching dynamical systems models for multiple interacting neural populations. Advances in neural information processing systems, 33:14867-14878 2020.
[9] Evren Gokcen, Anna I Jasper, João Semedo, Amin Zandvakili, Adam Kohn, Christian Machens, and Byron M Yu. Disentangling the flow of signals between populations of neurons. Nature Computational Science, 2(8):512-555) 225,, 2022.
[10] Evren Gokcen, Anna Jasper, Alison Xu, Adam Kohn, Christian K Machens, and Byron M Yu. Uncovering motifs of concurrent signaling across multiple neuronal populations. Advances in Neural Information Processing Systems, 36, 2024.
[11] Rainbo Hultman, Kyle Ulrich, Benjamin D Sachs, Cameron Blount, David E Carlson, Nkemdilim Ndubuizu, Rosemary C Bagot, Eric M Parise. Mai-Anh T Vu, Neil M Gallagher, al. Brain-wide electrical spatiotemporal dynamics encode depression vulnerability. Cell. 173 (1):166-180, 2018.
[12] Xiaoxuan Jia, Joshua H Siegle, Séverine Durand, Greggory Heller, Tamina K Ramirez, Christof Koch, and Shawn R Olsen. Multi-regional module-based signal transmission in mouse visual cortex. Neuron, 1000))1555-1159 2022.
[13] James Jun, Nicholas A Steinmetz, Joshua Siegle, Daniel J Denman, Marius Bauza, Brian Barbarits, Albert K Lee, Costas A Anastassiou, Alexandru Andrei, Çağatay Aydın, al. Fully integrated silicon probes for high-density recording of neural activity, Nature, 551(7679): 232-236, 2017.
[14] Stephen L. Keeley, David M Zoltowski, Mikio C Aoi, and Jonathan W Pillow. Modeling statistical dependencies in multi-region spike train data. Current Opinion Neurobiology, 65: 194-202, 2020.
[15] Mohammad Reza Keshtkaran and Chethan Pandarinath. Enabling hyperparameter optimization in sequential autoencoders for spiking neural data. Advances neural information processing systems, 32, 2019.
[16] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yev Lo, et al. Segment anything. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4015-4026, 2023.
[17] International Brain Lab, Brandon Benson, Julius Benson, Daniel Birman, Niccolo Bonacchi, Matteo Carandini, Joana A Catarino, Gaelle A Chapuis, Anne K Churchland, Yang Dan, et al. brain-wide map of neural activity during complex behaviour. bioRxiv, pages 2023-07, 2023.
[18] International Brain Laboratory, Kush Banga, Julius Benson, Niccolò Bonacchi, Sebastian A Bruijns, Rob Campbell, Gaëlle A Chapuis, Anne K Churchland, M Felicia Davatolhagh, Hyun Dong Lee, et al. Reproducibility of in-vivo clectrophysiological measurements in mice. bioRxiv, pages 2022-05, 2022.
[19] Ctrl labs at Reality Labs, David Sussillo, Patrick Kaifosh, and Thomas Reardon. A generic noninvasive ncuromotor interface for human- numan-computer interaction. bioRxiv, pages 2024-02, 2024.
[20] Ran Liu, Mehdi Azabou, Max Dabagia, Jingyun Xiao, and Eva Dyer. Sceing the forest and the tree: Building representations of both individual and collective dynamics with transformers. Advances neural information processing systems, 35:2377-2391. 2022.
[21] Yong Liu, Tengge Hu, Haoran Zhang, Haixu Wu, Shiyu Wang, Lintao Ma, and Mingsheng Long. itransformer: Inverted transformers are effective for time series forecasting. arXiv preprint arXiv:2310.06625, 2023.
[22] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, arXiv preprint axX1v1171.05101, 2017.

---

[23] Simon Musall, Matthew I Kaufman, Ashley L Juavinntt,, Steven Gluf, and Anne K Churchland. Single-trial neural dynamics are dominated by richly varied movements. Nature neuroscience, 22(10):1677-1666 2019.
[24] Chethan Pandarinath, Daniel J O' Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew T Kaufman, Stephen Ryu, Leigh R Hochberg, et al. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature methods, (5(10))805-81) 81555-815,20 2018.
[25] Felix Pei, Joel Ye, David Zoltowski, Anqi Wu, Raeed H Chowdhury, Hansem Sohn, Joseph E O'Doherty, Krishna Shenoy, Matthew T Kaufman, Mark Churchland, et al. Neural latents benchmark 21: evaluating latent variable models of neural population activity. arXiv preprint Ntt::200900442.. 2021.
[26] Matthew G Perich, Charlotte Arlt, Sofia Soares, Megan E Young, Clayton P Mosher, Juri Minxha, Eugene Carter, Ueli Rutishauser, Peter H Rudebeck, Christopher Harvey, et al. Inferring brain- wide interactions using data-constrained recurrent neural network models. BioRxiv, pages 2020-12, 2020.
[27] Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language understanding by generative pre-training 2018.
[28] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188):9, 2019.
[29] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning,
pages 8748-8763. PMLR, 2021.
[30] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent, arXiv preprint arXiv:2205.06177. 2022.
[31] João D Semedo, Amin Zandvakili, Christian K Machens, M Yu Byron, and Adam Kohn. Cortical areas interact through a communication subspace. Neuron, 102(1):249-259, 2019.
[32] Joshua H Siegle, Xiaoxuan Jia, Séverine Durand, Sam Gale, Corbett Bennett, Nile Graddis, Greggory Heller, Tamina K Ramirez, Hannah Choi, Jennifer A Luviano, et al. Survey of spiking in the mouse visual system reveals functional hierarchy. Nature, 592(7852):86-92. 2021.
[33] Nicholas A Steinmetz, Cagatay Aydin, Anna Lebedeva, Michael Okun, Marius Pachitariu, Marius Bauza, Maxime Beau, Jai Bhagat, Claudia Böhm, Martijn Broux, et al. Neuropixels 2.0: miniaturized high-density probe for stable, long-term brain recordings. Science, 372(6539): eabf4588, 2021.
[34] Carsen Stringer, Marius Pachitariu, Nicholas Steinmetz, Charu Bai Reddy, Matteo Carandini, and Kenneth D Harris. Spontaneous behaviors drive multidimensional, brainwide activity. Science, 364(6437) 2019.
[35] Yi Tay, Mostafa Dehghani, Vinh Q Tran, Xavier Garcia, Jason Wei, Xuezhi Wang, Hyung Won Chung, Siamak Shakeri, Dara Bahri, Tal Schuster, et al. U12: Unifying language learning paradigms. arXiv preprint ar arXiv:2205.05131, 2022.
[36] Armin Thomas, Christopher Ré, and Russell Poldrack. Self-supervised learning of brain dynamics from broad neuroimaging data. Advances in neural information processing systems, 35:21255-21269 2022.
[37] Eric M Trautmann, Janis K Hesse, Gabriel M Stine, Ruobing Xia, Shude Zhu, Daniel J O' O'Shea, Bill Karsh, Jennifer Colonell, Frank F Lanfranchi, Saurabh Vyas, et al. Large-scale high-density brain-wide neural recording nonhuman primates. bioRxiv, pages 2023-02, 2023.
12

---

[38] Anne E Urai, Brent Doiron, Andrew M Leifer, and Anne K Churchland. Large-scale neural recordings call for new insights to link brain and behavior. Nature neuroscience, 25(1):11-19, 2022.
[39] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[40] Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le. Finetuned language models are zero-shot learners. arXiv preprint wrXiv:2109.01652 2021.
[41] Joel Ye and Chethan Pandarinath. Representation learning for neural population activity with neural data transformers. arXiv preprint arXiv:2108.012/11 2021.
[42] Joel Ye, Jennifer Collinger, Leila Wehbe, and Robert Gaunt. Neural data transformer 2: multi- context pretraining for neural spiking activity. bioRxiv, pages 2023-09, 2023.
[43] Zhiwen Ye, Andrew M Shelton, Jordan R Shaker, Julien Boussard, Jennifer Colonell, Sahar Manavi, Susu Chen, Charlie Windolf, Cole Hurwitz, Tomoyuki Namima, al. Ultra-high density electrodes improve detection, yield, and cell type specificity of brain recordings. bioRxiv 2023.
13

---

Appendix
A Additional results
Architecture agnosticity. We have added an experiment where we trained the LFADS (a sequen- tial VAE) architecture [24] on 5 test sessions using the temporal masking baseline and our MiM learning objective. The results of the LFADS analysis are shown in Table We found that the LFADS architecture trained with MtM outperforms temporal masking on all activity prediction tasks and is competitive with NDTI trained with MtM (although slightly worse). Behavior prediction comparable between MtM and temporal masking. This experiment suggests that MtM is an architecture-agnone learning objective that can be used with both transformers and RNNs.
Table 3: The performance of single-session NDT/ vs. LFADS trained with various masking schemes on activity reconstruction and behavior decoding lasks. The metric for activity reconstruction is in units of bits per spike (bps), averaged across all neurons in session. Behavior decoding assessed using accuracy for choice and Rº whisker motion energy. For all metrics, a higher value indicates better performance.
Activity Reconstruction
Behavior Decoding
Co-Smooth Forward Pred Intra Region Inter- Inter-Region Choice
Whisker Motion Energy
NDTI Baseline
0.55
0.17
-0.49
0.18
0.87
0.65
NDTI McM
1.00
0.28
0.70
0.83
0.85
0.67
LFADS Baseline
0.86
0.09
-0.34
0.50
0.86
0.67
LFADS MtM
0.87
0.26
0.65
0.76
0.82
0.67
Estimation of "functional connectivity". We utilized our pretrained MtM model to ask how well individual brain areas predict each other. We fine-tuned our pretrained MtM model on each test session with new inter-region task where we randomly dropout some of the regions when predicting one region. Then, we used MtM predict region A from region B for all regions in the session. The prediction matrix is shown in Figure 6 Even without explicit training to predict one region from another region, there interesting structure the matrix as some regions predict each other well, for example in visual areas. We also fine-tuned our pretrained MtM model on each test session to predict region from region B for selected regions. We compared performance NDT1 regression models trained from scratch for the same region-to-region predictions. In Figure 7 the prediction matrices show that the fine-tuned MtM outperforms the NDTI regression models region-to-regie prediction. Taken together, these results suggest that MtM can be used to explore the "functional connectivity" of different areas.
APN
0C-mo (25)
06 mo (47)
1.0
nn-ug (9)
Region
LP
(115)
0.
MRN (85)
0.0
vise23 113
VSS22d (26)
Pro 129)
Visad
ViSate
1.5
ViSata
ber (19)
0C-13
IP
vi5u00
05
1922
-2.
DG
Region 8
Region
Region
Figure 6: Predicting region A from region matrix using MrM. Each matrix test session and each entry
the co-bps when predicting region A from region B using MtM. We fine-tune MtM on each session using an inter-region masking scheme where randomly dropout regions to encourage the model to utilize every region. We filter out all regions with less than neurons.
Cross-species transfer. While the IBL repeated site dataset used for training MtM high diversity of brain regions, it lacks task (species) diversity. To demonstrate that MtM can be useful pretraining strategy for new tasks and species, we fine-tuned our 34-session pretrained NDT1-MtM model on the Neural Latents Benchmark (NLB) (25) MC. RTT dataset. This dataset consists of spiking activity

---

A
20
20
mean bps B-aA, >0.1Hz. (MLM finetune), eid Sdcee
mean bps B->A, >0.1Hz. INDT regression), eid Sdcee
VI5a
0.238
0.292
0.301
0.195
(-4)
0.313
3.323
0.293
0.172
1.0
10
EX1 1230
(23)
020
0.184
1.184
0.735
0.212
0.147
0.0
DG (28)
0.885
0.413
ron
0.188
0.091
-0.5
-0.5
.P
0.217
n2n
157)
nan
0.348
0.238
0.286
0.253
153)
0.888
0.209
3.236
0.25
2 -1.5
-15
3
5
VC2
8
9
region
-2.0
regton
Figure 7: A comparison of the pretrained MtM fine-tuned for region-to-region prediction the NDTI regression model trained from scratch for same task. The matrix panel displays region-to-region predictions from the fine-tuned MtM, while panel B presents the predictions from the NDTI regression model. Both matrices are from the same session, with each entry representing the co-bps when predicting region A from region B.
Co So-SSoothhing
Ground Truch
MEM
Baselne
Fring Rate
Nou from CAL
15
10
0.5
0
0.5
1.0
1.5
0.5
0
0.5
1.0
1.5
0.5
0
0.5
1.0
1.5
Forward Predittion
bas: 0.17
bast 0.11
Baseline
Fring
40
30
Neurons
20
Figure 8: Rastermap visualizations of MrM and temporal masking for a single session on co-smoothing and forward prediction tasks. Each row corresponds to different predictive task. The far left column is schematic of the predictive task, the second column is the ground-truth spiking activity in region, and the third and fourth columns predicted rates from MtM and temporal masking (baseline), respectively.
recorded from the primary motor cortex of monkey performing divided this dataset into two "brain regions" based on the NLB held-in and held-out neuron split. We then fine-tuned our model using the MtM learning objective. The MtM performance (0.19 bps) significantly improved upon NDTI's previous best performance (0.16 bps).
reaching task. We
Additional rastermap visualizations. We have added visualizations of the rastermaps for NDT1 trained with MtM and temporal masking for each predictive task: neuron, causal, inter-region, and intra-region prediction. As can be seen in both Figure and Figure 8 MtM leads to better predicted rastermaps for inter-, intra-, and neuron masking. For forward prediction, the improvement is more subtle, largely because temporal masking performs well at temporal predictive tasks. Overall, these visualizations give confidence that MtM significant improvement over the baseline.
B Single-session details
Masking scheme ablation. understand the importance of each masking scheme, we evaluate NDTI and NDT2 trained with various masking schemes, including temporal, neuron, causal, intra-
15

---

Table 4: performance of single-session NDTI trained with various masking schemes on activity recon- struction and behavior decoding tasks. The metric for activity reconstruction is in units of bits per spike (bps), averaged across all neurons in one session. Behavior decoding is assessed using accuracy for choice and R for whisker motion energy, For all metries, higher value indicates better performance. Masking
Activity Reconstruction
Behavior Decoding
Co-Smooth
Forward Pred Intra-Region
Inter- Aierr-Reiion
Choice
Whisker Motion Energy
Temporal (Baseline)
0.84
0.42
-0.20
0.57
0.66
0.56
Neuron
-0,21
-0.22
0.78
0.60
Causal
0..88
-0.36
0.23
0.75
0.59
Intra-Region
-9.86
-2.97
0.32
-9.06
0.55
0.38
Inter- Maer-Region
0.92
0.01
-0.58
0.90
0.52
0.59
MIM (Not Prompted)
0.99
0.54
0.42
0.83
0.69
0.61
MtM (Prompted)
0.98
0.57
0.43
0.84
0.63
0.61
Table 5: The performance single-session NDT2 trained with various masking schemes activity recon-
struction and behavior decoding tasks. metric activity reconstruction units of bits per spike (bps), averaged across all neurons in one session. Behavior decoding is assessed using accuracy for choice and for whisker motion energy. For all metries, higher value indicates hetter performance.
Masking
Activity Reconstruction
Behavior Decoding
Co-Smooth
Forward Pred Intra- Inter-Region Choice
Whisker Motion Energy
Random Token (Baseline)
-6.94
0,50
-0.43
-6.95
0.74
0.58
Neuron
0.91
0.18
-0.26
0.62
0.65
0.62
Causal
0.02
0.52
-0.42
-0.20
0.69
0.59
Regi-R
-10.10
-1.30
0.21
-8.17
0.65
0.43
Region
0.63
-0.63
0.66
0.75
MtM (Not Prompted)
0.90
0.56
0.80
0.68
0.62
M(M (Prompted)
0.69
region, inter-region, as well as the proposed MtM method with and without the prompt token on a selected single session. For both NDTI 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 its corresponding task. Figure Sacureeccom 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.
Whisker Motion Energy
Co Smooth
Forward Prediction on
Intra Region
Region
Choice
0
-2.5 0.0 Baseline
2.5
-2.5
0.0 Baseline
-2.5 0.0 Baseline
2.5
0.5
1.0
Baseline
Baseline
Baseline
Figure 9: Comparison the random token masking baseline and the proposed MtM method for NDT2 activity reconstruction and behavior decoding across 39 sessions. For activity reconstruction, each point shows the average hps across neurons in one session. For choice (whisker motion energy) decoding, each point represents the average accuracy (R²) across all test trials in one session.
Binary mask token ablation for NDT2. To optimize NDTI and NDT2 for predicting held-out neurons, sub-token channels are masked during so-omoothing 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, this approach did not show benefits, leading to the use of zero masking instead. We also tested concatenating binary mask with NDT2, which showed a slight improvement in Table though overall performance metries for random token masking remained negative.
16

---

Co-Smooth Bps
Forward Prediction Eps
Intra Region Bps
Inter Region Bps
0.0
2.5
average
Trial average
average
average
Log Rates (Hz2)
0.4
MMM
0.4
0.4
0.2
0.2
0.2
0.0
0.0
-2
0.0
0.5
0.0
0.5
0.00
0.25
0.50
0.00
0.25
0.50
PSTH R2
PSTH R2
PSTH R2
PSTH R2
0.75
0.75-
0.75
0.75
0.50
0.50
0.50
MtM
0.25
0.25
0.25
0.00
0.00
0.00
0.00
-0.25
0.25
0.25
0.25
0.0
0.5
0.0
Temporal
Temporal
Temporal
Temporal
Figure 10: Single neuron activity reconstruction analysis for NDTI one session. To evaluate the reconstruction quality for each neuron, multiple metrics are computed: Bits per spike (Bps). H² between the ground truth and predicted peristimulus time histogram (PSTH R²), and the single-trial R averaged across all trials (Trial average R²). Each point represents one neuron, with the color indicating the neuron' 's log firing rates in Hertz (Hz).
Table 6: Ablation study on the impact of using binary mask token with NDT2 for downstrean activity reconstruction tasks. Results based single selected test session.
Co-Smooth Forward Pred Intra Intra-Region Inter-Region
NDT2 Random Token
-6.84
-0.84
-6.39
NDT2 Random Token Rinary Mask Token NDT2 MtM
-2.27
0.41
-0.55
-2.10
0.85
0.51
0.32
0.73
Single neuron evaluation. To better understand neural activity prediction performance at single- neuron level, we conducted an evaluation of single-session NDT1 on each neuron using bits per spike (bps), R² between the ground truth and predicted peristimulus time histogram (PSTH R²) ), and the single-trial R² averaged across all trials (trial average R²), on one session, in Figure 10 We find that MtM outperforms the temporal baseline across most neurons regardless of firing rate, We did find 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 be more concentrated, while scores for inactive neurons are relatively dispersed, exhibiting both extremely low and high values. For both It2 metries, the performance shows a positive correlation with the mean firing rates. In particular, those neurons with extremely low mean firing rates typically exhibited It2 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² metric might tend to yield values closer to zero. Across all three metries (Bps/ Trial average RP/PSTH R²), our model demonstrated substantial improvements for neurons with relatively higher mean firing rates. However, for neurons with lower mean firing rates, notable improvements only observed in the bps metric.
17

---

Whisker Motion Energy
Co-Smooth
Forward Prediction
Intra-Region
Inter-Region
Choice
1.5
0.50
1.0
1.5
0.75
1.0
1.0-
0.5
0.25
0.75
0.50
Baseline
Baseline
Baseline
Baseline
Baseline
Baseline
Figure 11: Fine-tuning performance comparison of NDT2 pretrained with MIM vs. random token masking for activity reconstruction and behavior decoding across held-out sessions, For activity reconstruction, each point shows the average bps across all neurons in held-out session. For choice (whisker motion energy) decoding. each point represents the average accuracy (H²) across all test trials one session.
Energy
Avorage Bits por Saha
0.95
1.05
0.90
0.70
1.00
0.3
0.65
0.85
0.85
0.65
0.2
KDTZ
0.60
0.90
10
Eeerage atz per Spite
3.5
-0.2
Accurcy
0.0
0.90
0.70
3.40
0.5
0.35
-0.8
1.0
3.30
0.80
0.60
-0.8
10
34 Num of Sessions
10
34 Num of Sessions
10 34
10
34
10
34
10
24
Nur Sessions
Num of Seasions
Num of Scssions
Num of Sessions
Figure 12: Comparison of scaling curves between NDT2 pretrained with MrM method vs. random token masking baseline. The reported metrics neuron-averaged hits per spike, choice decoding accuracy, and whisker motion energy decoding R² are averaged over all held-out sessions used for fine-tuning on both activity reconstruction and bchavior decoding tasks. "Num of Sessions" denotes the number of sessions used for pretraining.
C Multi-session details
For NDT2, report results for pretraining using MtM and the baseline random token masking on 34 sessions of data in Table 7 Figure [1] Table 7 we show -vverane 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-sessio results, the biggest improvements for MtM are for the unsupervised activity metrics, especially inter- and intra-region prediction. In Figure I1 we show 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 MiM vs. random token masking on activity reconstruction and behavior decoding. Activity reconstruction performance is reported in souron-averaged bps. For behavior decoding, trial-averaged accuracy (H²) for choice (whisker motion energy) decoding is shown. All metrics averaged
held-out sessions, and higher value indicates hetter performance.
Activity Reconstruction
Behavior Decoding
Training
Masking
Co-Smooth Forward Pred Intra Rogion
Inter Region
Choice Whisker Motion Energy
Single-Session Session
Random Token (Baseline)
-0.87
0.35
-0.99
-0.75
0.84
0.65
MIM (Prompted)
0.36
0.84
0.85
0.67
Random Token (Buseline)
0.26
0.38
-0.49
-0.39
0.89
0.66
Multi-Session
MIM (Prompted)
1.01
0.34
0.71
0,87
0.87
0.67
18

---

Scaling analysis. To examine NDT2's ability of scaling data, Figure 12 shows that NDT2 multi- session pre-training also benefits from scaling from 1 10 sessions, but we only observe marginal gains or no improvement going from 10 to 34 sessions. NDT2 (Figure 12 benefits less from multi- session IBL pre-training compared to NDTI-stitch (Figure 4 likely due to the inability of NDT2 with 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 NDTI 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 we ablate four prompt modes.
Table 8: Evaluation of NDTI's 's behavior decoding performance when pretrained and fine-tuned using the MIM approach, tested with different prompt tokens inference time. Behavior decoding is assessed using trial-averaged accuracy for choice and trial-averaged R2 for whisker motion energy, with the reported metrics averaged over 5 held-out session.. For both metrics, a higher value indicates better performance.
Behavior Decoding
Prompting
Choice
Whiker Motion Energy
Neuron
0.88
0,68
Causal
0.89
0,68
Intra- Intra-Region
0.88
0,68
Inter-Region
0.88
0.68
D Experiment details
Architecture details. We created schematic for the NDTI and NDT2 architectures in Figure 13 to clarify how these architectures tokenize and reconstruct neural data. NDTI uses the timestep for the positional encoding. NDT2 uses 2 types of position encoding: (I) timesteps and (2) neuron group identity.
NDT-1 MtM Architecttre
icconttctionlee
NDT-2 MtM Architecture
reconatuctionles los
reut
projetion
Transfsmer
Transformer
postoor
road
wesilo speatli projedor
Inver
slare scooe papu prooi
2
Figure 13: Architectures for NDTI and NDT2. Each architecture tokenizes using temporal tokens although NDT1 uses the full population and NDT2 uses patch of neurons. Following a lincar transformation of the token, a position embedding is used before passing into the transformer. The output of the transformer is transformed back into neural activity using linear readout weights.
Training details. The 34-session NDTI model has 25.55 million parameters, many of which come from cession-specifii stitchers. However, these stitchers use only one linear layer per session, making the model capacity equivalent to that of single-session NDTI model. We do not provide model
19

---

sizes for the single-session NDTI, 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
stitchers and consist of 1.09 million parameters. We trained our model using AdamW W optimizer [22] for 1000 epochs with learning rate of 1e-4 using a cosine scheduler. We put weight decay 0.01 to avoid overfitting. We utilized 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 trial-average R² of top-50 active neurons, which we selected top-50 active neurons and calculated each neuron's 's It2 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 and 3 days, respectively, while the 10-session and 34-session NDT2 models took and 5 days, respectively. Single-session NDTI and NDT2 models, well as finetuning, were trained on a single Tesla V100 GPU, requiring to hours. We make sure our experiments are reproducible by seeding.
Hyperparameters details. The masking ratio is an important model hyperparameter for NDTI 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 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 NDTI performance in neural activity reconstruction. The reported metries quantify performance in terms of average bits per spike (bps) across all neurons from a selected session. higher bps value indicates better performance.
Masking
Mask 0,)
Musk
Co Smooth
Freward Pred Intrs Resion
Inter Resion
Ci Smooth
Freweed Proi
Intra Retion
liter Rezion
Co Siaoohhhh
Forwar Prol lnsta Region
laiz Repion
Tempural (Baucinne)
0.67
0.75
-0.1V
0.46
1.21
0.98
0.51
0.52
0.92
0.68
0.42
1156
Nearoz
1.41
-0.17
1120
0.55
1.34
-1104
0.22
0.79
1:25
0.10
01.6
1.10
Cazed
0.46
0.92
0.46
0.92
0.14
Iamm- Regivn
0:49
1.34
1.17
0.49
-000
1.14
1:13
4000
111
MM (Pocepted)
1.31
0.79
0.92
1.17
1.24
0.74
0.77
1.8
0.53
0.53
0.71
0.58
Ablation study on positional embeddings. We perform an additional experiment to evaluate the impact of the positional embedding MIM's performance for different downstream tasks. In Table we compare the effects of RoPE and learnable positional embeddings on the NDTI architecture pretrained using the MtM method. The results demonstrate that RoPE consistently outperforms learnable positional embeddings across all downstream tasks.
Table 10: Ablation study examining the impact of RoPE vs. learnable positional embeddings in NDTI pretrained with the MrM method across various downstream tasks. The displayed results are averaged over three sessions.
Activity Reconstruction
Behavior Decoding
Co-Smuoth Forward Pred Intra-Region Inter-Region Choice Whisker Motion Energy
Learnable PE
0.74
0.27 0.36
0.54
0.54
0.87
0.69
RoPE
0.78
0.60
0.60
0.90
0.70
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 neuron is defined as follows:
bits/spike
(C(r; Xn,)) - C(rm,si X...)),
log
where Xn,, represents the neural activity of the target neuron n at time f, r is the inferred firing rate from the model, Po.: is the mean firing rate for neuron n, and nop is its total number of spikes. We then average the bits per spike across all neurons in test session to get the final performance.
20

---

The bits spike metric is closely related to the deviance, standard metric generally used in Statistics, e.g.. generalized linear models (GLMs) deviance goodness-of-lit measure for statistical model. The deviance compares the goodness-of-fit of the model interest, e.g.. MtM, to baseline null model, where goodness-of of-fft fit measured by the model log likelihood. The bps is 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 that the metric be comparable across neurons regardless of whether the neurons active 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 []] For whisker motion energy, we report the test R that will be reached when consistently using the train set's trial average. Despite imbalanced datasets (e.g., c7bf2d49), Figure shows that decoding accuracy and R from selected regions are generally larger than these chance levels.
Table 11:
The chance for decoding choice and whisker motion energy in each held-out session.
Chance Level
5dcee0eb c7bf2d49 3d9a098 d57df551
824cf03d
Choice
0.58
0.91
0.52
0.49
0.62
Whisker Motion Energy
0.44
0.30
0.59
0.43
0.39
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reffect the paper's contributions and scope?
Answer: [Yes]
Justification: We clearly declare and introduce our main contributions as the "universal translator" for neural dynamics.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Please refer to the Discussion section for the limitations.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results these assumptions (e.g.. independence assumptions, noiseless settings, model sell-specification asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
violations of
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting.
speech-to-texx system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms how they scale with dataset size.
If applpicaale, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't ack acknowledged in the paper. The authors should that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
their best judgment and recognize
Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: This paper does include thcoretical results.
Guidelincs:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
22

---

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors encouraged to provide 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 supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
Experimental Result Reproducibility
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 Dand include the seed, model architecture, and hyperparameter choices in the attached code.
Guidelines:
The answer NA means that the paper does not include experiments.
the paper includes experiments, 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 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 novel architecture, describing the architecture fully might suffice, if the contribution is specific model and empirical evaluation, it may be necessary to either make 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 hosted model (e.g., in the case of a large language model), releasing of model checkpoint, or other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g.. a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g.. with open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, which case authors are welcome to describe the particular way they provide for reprodueibility. In the case of closed-source source models, it may be that access to the model is limited in some way (c.g.. registered users), but it should be possible for other rescarchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code. with sufficient instructions to faithfully reproduce the main experimental results, described in supplemental material?
Answer:
[Yes]
Justification: We use public datasets for our experiments. Our code is attached for reproducibility of our results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines httpeee////pps.ccpppppleee
Gaadee//CoeSuban
for more details.
While we encourage release of code and data, we understand that this might not possible,
so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g.. for a new open-souree benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: : //:/://nppp cc/public/ qul.dess/CooeSuubmiiissooooooo.. for more details.
23

---

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reprodueible, they should state which ones are omitted from script and why.
At submission time, to preserve anonymity, authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
Experimental Setting/Details
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 understand the results?
Answer: [Yes]
Justification:
We specify all the training and hyperparameters details in Appendidx 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 level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined other appropriate informa- tion 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.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results 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 variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some paramcter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to 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 the standard error of the mean.
is OK report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures
symmetric error hars that would yield results that are out of range (c.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain 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 the computer resources (type of compute workers, memory, time exccution) needed 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.
24

---

The paper should provide the amount of compute required for each of the individual experimental well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, every respect, with NeurIPS Code of Ethics hattp:://eerrti hwtpc//pabcccccccccauac.a.aaaaaa
Answer:
[Yes]
Justification: We carefully read and follow the NeurIPS Code of Ethics. All of our cade and main text are anonymous.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure preserve anonymity (e.g., if there is special consideration due laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts the work performed?
Answer: [NA]
Justification: Foundational research, not specifically linked to any social applications, except for its relevance to fundamental neuroscience research.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
the authors answer NA or No, they should explain why their work has no socictal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it not needed point out that generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitor- ing misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g.. pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does not include any high risk data and all results and models are free to share.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines restrictions access the model implementing safety filters.
25

---

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 encourage authors take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We clearly indicate cited existing assets by citing the work.
Guidelines:
The answer NA means that the paper does not 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 (c.g., CC-BY 4.0) should be included for cach asset. scraped from particular source (e.g.. website), the copyright and terms of service of that source should provided.
assets are released, the license, copyright information, and terms of use the package should be provided. For popular datasets, paperswi thcode. www.coom//aaac /attttttstts has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-per re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
this information not available online, the authors are encouraged to reach out to the asset's creators.
13. 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 datasct/code/model as part of their sub- missions via structured templates. This includes details about training, license, limitations,
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 amonymized URL or include anonymized zip file.
14. 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.
Guidelincs:
The answer NA means that the paper does not involve crowdsourcing 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 included the main paper.
According to the NeurIPS Code of Fthics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals Equivalent for Research with Human Subjects
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?
Question:
26

---

Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer means that paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, approval (or equivalent) may be required for any human subjects research. If you ohtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applica- ble), such the institution conducting the review.
27


Reviews:
[{"review_id": "ecregalMlD", "reviewer": "Reviewer_BBJQ", "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.", "questions": "1.\tIt 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.\n2.\tIt 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?\n3.\tIt 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). \n4.\tEquation (1): Does the transformer use any kind of positional encoding? If so, what type of positional encoding? \n5.\tFor 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? \n6.\tLine 174: possibly a typo, should be 1% of the trial (20ms).\n7.\tHow 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?\n8.\tSimilarly, was the motion energy decoded at each timestep of $\\hat{X}$ or at some particular point during the trial?\n9.\tFigure 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": "The authors have addressed the limitations in the manuscript.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1.\tThe paper tackles an important problem of building “foundation models” for neuroscience at the level of spiking activity that can help with the understanding of the brain structure and facilitate the decoding of animal behaviors. \n2.\tThe proposed MtM method is novel and demonstrated effectiveness over the existing temporal masking method in learning different structures of neural activity.\n3.\tThorough evaluation on multi-region, multi-subject IBL dataset provides valuable insights for the neuroscience community.  \n4.\tThe manuscript is well-written. Text and figures are logical and easy to follow.", "weaknesses": "1.\tThe 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. \n2.\tThe 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.\n3.\tSome experiment details are missing or not presented clearly (elaborated in Questions section below).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses:**\n1. We thank the reviewer for the suggestion to run more baseline architectures with the MtM objective. Our goal was to compare masked modeling approaches for neural population modeling which just includes temporal masking as a baseline. Based on this feedback, however, we ran a new experiment where we trained the LFADS architecture (a sequential VAE) with both temporal masking and MtM. The results can be seen in Table 1 of the Global Response. MtM training improves all activity prediction metrics for LFADS especially intra-region masking and inter-region masking. The best performing model is still NDT1-MtM, but the LFADS-MtM model is quite competitive. This experiment provides evidence that MtM is a model-agnostic objective that can teach region-level population structure to both transformer and RNN architectures.\n2. We agree with the reviewer that there is a lack of task and species diversity in our original submission. To partially address these concerns, we fine-tuned our 34-session pretrained NDT1-MtM on the Neural Latents Benchmark (NLB) MC_RTT dataset. For details on this experiment, please see the Global Response. With a single private submission to the leader board, the MtM performance (.19 bps) significantly improves upon NDT's previous best performance (.16 bps). These results show promising transfer of our pretrained MtM models to a new task and species. We also want to point out that while the IBL datasets lack species and task diversity, they have high diversity of recorded brain areas. Brain region diversity is an unexplored axis for large-scale pretraining as the POYO and NDT2 papers only train on data from motor areas. We wanted to tackle this challenge first before moving on to cross-species transfer. However, we believe that these approaches are complementary and that building large-scale models that generalize across different tasks, species, and brain areas will be an exciting future direction. \n\n**Questions:**\n1. We agree with the reviewer that our paper lacks a simpler NLB baseline and have implemented an LFADS baseline model trained with temporal and MtM masking (results in Table 1 of the Global Response). NDT1 is still the top performing model when trained with the MtM objective although LFADS-MtM is competitive. While we agree that the LFADS RNN architecture is simpler, transformers have been shown to scale extremely well to massive datasets and also have benefits for computational efficiency (please see Figure 6 of [1]). We will clarify these reasons for using transformers in the updated manuscript. We did not include an additional MLP baseline for behavior decoding because we are decoding from the inferred rates and, therefore, would not expect to do better than an MLP. This is the strategy used in the Neural Latents Benchmark [2] to evaluate behavior prediction.\n\n2. This is a great question and something we are interested in exploring as future research. To attempt to answer this question, we fine-tuned our 34-session pretrained NDT1-MtM on the Neural Latents Benchmark (NLB) MC_RTT dataset. All details for this experiment can be found in the Global Response.\n\n3. Yes, there is some overloading of terminology in our paper between behavioral \"tasks\" in neuroscience and the predictive \"tasks\" of MtM. In MtM, there are multiple neuron-level and region- level predictive tasks. We will make this distinction clear in the final version of the paper.\n\n4. Yes, both NDT1 and NDT2 use a positional embedding layer for each token. NDT1 uses the timestep for the positional encoding. NDT2 uses 2 types of position encoding: (1) timesteps and (2) neuron group identity. We did not include the positional embedding step in Equation (1), but we can correct this omission in the final version. We also created schematic diagrams for both NDT1 and NDT2 (see Figure 3 in the one page PDF) which we will add to the supplement to improve our description of the underlying architectures.\n\n5. In the held-out sessions, there are 397 ~ 579 trials in the train split, 57 ~ 83 trials in the validation split, and 114 ~ 166 trials in the test split. \n\n6. Yes, good catch. This should read \"10% of the trial (200ms)\". We will correct this in the final version.\n\n7. For each trial, we predict the binary choice using the entire $N \\times T$ matrix $X$, which represents the spiking activity of all neurons across all timesteps.\n\n8. In each trial, we use the entire $N \\times T$ matrix $X$ to predict the motion energy at every timestep $t = 1, 2, …, T$, where $T$ is the total number of timesteps.\n\n9. For binary choice, we report the test accuracy that will be reached when constantly predicting the train set's majority class. For whisker motion energy, we report the test $R^2$ that will be reached when consistently using the train set’s trial average. Despite imbalanced datasets (e.g., c7bf2d49), Figure 5 shows that decoding accuracy and R2 from selected regions are generally larger than these chance levels. \n\n    |      Chance Level     | 5dcee0eb | c7bf2d49 | 3d9a098 | d57df551 | 824cf03d |\n    |:---------------------:|:--------:|:--------:|:-------:|:--------:|:--------:|\n    |         Choice        |   0.58   |   0.91   |   0.52  |   0.49   |   0.62   |\n    | Whisker Motion Energy |   0.44   |   0.30   |   0.59  |   0.43   |   0.39   |\n    |                       |          |          |         |          |          |\n\n[1] Ye, Joel, and Chethan Pandarinath. \"Representation learning for neural population activity with Neural Data Transformers.\" arXiv, 2021.\n\n[2] Pei, Felix, et al. \"Neural latents benchmark'21: evaluating latent variable models of neural population activity.\" arXiv preprint arXiv:2109.04463 (2021).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tYTknytcGA", "reviewer": "Reviewer_JcgB", "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.", "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": "The authors adequately state the limitations of their modeling approach, training paradigm, and datasets.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The work is original and forward thinking. The writing is admirably clear and well referenced.", "weaknesses": "The paper would benefit from:\n- 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?\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses:**\n1. We thank the reviewer for this suggestion to compare to a wider range of neural population models. We want to clarify that MtM can be used as a learning objective for any architecture (not just transformers). To demonstrate this, we ran a new experiment using the LFADS architecture (a sequential VAE) trained with both temporal masking and MtM. The results can be seen in Table 1 of the Global Response. MtM training improves all activity prediction metrics for LFADS especially intra-region masking and inter-region masking. The best performing model is still NDT1-MtM, but the LFADS-MtM model is quite competitive. This experiment provides evidence that MtM is a model-agnostic objective that can teach region-level population structure to both transformer and RNN architectures. We also attempted to compare MtM to mDLAG on three regions (VISa2/3, VISa4, VISa5) of a single test session. For these regions, there were a total of 42 neurons, 828 trials, and 100 time bins (20 ms bin size). We trained mDLAG for 5.5 hours (5000 EM iterations) and the model did not converge. A visual inspection shows that the learned latent factors for each region were nearly identical and that activity reconstruction from the learned latents consistently led to low co-bps (negative to zero). This experiment provides evidence that scaling mDLAG to full IBL datasets with tens of regions and thousands of neurons is not feasible. We will provide all details for this experiment in the updated manuscript.\n\n\t\n2. We agree with the reviewer that this is an important detail missing from the initial submission. We will make sure to add a discussion of both temporal masking for NDT1 and random token masking for NDT2 in the final submission. We also plan to add schematics for the NDT1 and NDT2 architectures in the updated manuscript (Figure 3 of the one page PDF) to demonstrate how they tokenize and reconstruct neural data.\n\n**Questions:**\n1. This is a great question and something we are also thinking hard about! We believe that by incorporating data from many different sources, a neural foundation model can help us answer important questions about cross-individual (and even cross-species) neural variability. We also believe that neural foundation models, which are trained on many different combinations of brain regions, can eventually be used to predict missing brain regions in new recordings. We also believe that by training neural foundation models on healthy animals, we can test hypotheses about neural activity in diseased animals. This advancement could potentially lead to new treatments and interventions for neurological disorders. Finally, it will be possible to fine-tune pretrained foundation models on new tasks (brain computer interfaces, etc.) to achieve state-of-the-art performance [1]. We believe our work is a step towards realizing this vision and we are hopeful that the community can start working with and building similar models. We will improve our discussion of foundation models of neural data in the updated manuscript.\n\n[1] Azabou, Mehdi, et al. \"A unified, scalable framework for neural population decoding.\" Advances in Neural Information Processing Systems 36 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mcMRkArZrY", "reviewer": "Reviewer_xe9o", "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).\n\nThe 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.", "questions": "184: Please include a formal definition of \"bits per spike\"\n\nTable 1: How a \"bits per spike\" measure can be negative? What does it mean?\n\nFigure 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": "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?", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* In my opinion, this work is a step in a very important direction, indeed with the wealth of neural data being collected foundation models for various modalities can become very important tools for modelling brain dynamics\n\n* Clearly and simply written, easy to read and follow", "weaknesses": "* 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.\n\n* 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\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses:**\n\n1. We agree with the reviewer that this is a important detail missing from our paper. To address this, we added schematics for NDT1 and NDT2 that demonstrate how tokenization and neural reconstruction works (please see Figure 3 in the 1 page PDF). We will add this figure to the supplement in the final version of the paper.\n\n2. We agree with the reviewer that the lack of visualizations is a problem with our original submission. To partially address this, we provide visualizations of the rastermap reconstructions for all predictive tasks in Figure 1 of our one page pdf. The difference in the bps metric for temporal masking and MtM can clearly be seen in the region-level predictive tasks and can also be seen in neuron and forward prediction tasks. We will make sure to add rastermap visualizations for many different sessions in the supplement of our final paper. Hopefully, this will convince readers to try out our approach for NeurIPS 2024! :-)\n\n**Questions:**\n1. We define the bits per spike for a neuron with the same equation as [1] on the bottom of page 6. We then average the bits per spike for each neuron in a test session to get the final performance. We will include a version of this equation in the updated manuscript.\n\n2. The bits per spike (bps) can be negative when the predicted rates for a neuron are less informative than just predicting the average firing rate of the neuron for the trial. This can be seen in the above equation where we subtract the likelihood of the neural activity assuming an average firing rate from the likelihood of the neural activity with the predicted rates of the model. The bps also tends to be more negative when trying to predict low-firing rate neurons.\n\n3. \"Choice\" refers to the choice decoding accuracy for each method. \"WME\" refers to the whisker motion energy decoding $R^2$ for each method . We will improve the axis labels in the final version of the paper.  We will also try fixing the axis for each session as suggested by the reviewer.\n\n[1] Pei, Felix, et al. \"Neural latents benchmark'21: evaluating latent variable models of neural population activity.\" arXiv preprint arXiv:2109.04463 (2021).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UDHt61ecmY", "reviewer": "Reviewer_P8rK", "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.", "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?\n\n2. 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": "see Weaknesses.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The deep thinking about foundation model in neurosciece.\n\n2. The model significantly improves generalization by modeling across different brain areas, individuals, and sessions.\n\n3. The multi-task-masking (MtM) strategy effectively enhances the capabilities of the Transformer model.", "weaknesses": "1. Lack of Qualitative Analysis:\n\nThere 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?\n\n2. Insufficient Ablation Studies:\n\na) 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?\n\nb) What are the performance differences when using multi-session data from a single individual versus cross-individual multi-session data?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Weaknesses:**\n1. We agree with the reviewer that our submission was lacking in qualitative analysis. To partially address this, we provide visualizations of the rastermap reconstructions for all predictive tasks in Figure 1 of our one page pdf. For the final version of our paper, we will include rastermap visualizations for many different animals/sessions, so readers can see differences in neural activity across different individuals. We believe there might be some confusion about the neural dynamics learned in our models. Since NDT1 is a transformer-based model, neural dynamics are learned implicitly. We are not explicitly learning neural dynamics like a state-space model would. So it is not possible to visualize the dynamics directly in a transformer. Building transformer architectures with explicit latent dynamics is an interesting future direction.\n\n2. These are very interesting questions! (a) When training on single region data, we would recommend just using the neuron and causal masking, as region-level predictions are not informative in this setting. When training on data from multiple brain areas, we uniformly sample brain areas for inter/intra region masking and then we average the loss per-neuron and per-timestep; this makes sure that no brain region dominates the loss. (b) We are using the IBL repeated sites dataset [1] which only have repeated sessions for one individual so we cannot easily address the difference in performance when using multi-session data from a single individual versus cross-individual multi-session data. Using the full public IBL brain-wide map dataset [2] would allow us to test this phenomena and is an exciting future direction.\n\n[1] International Brain Laboratory, et al. \"Reproducibility of in-vivo electrophysiological measurements in mice.\" bioRxiv (2022): 2022-05.\n\n[2] International Brain Laboratory, et al. \"A brain-wide map of neural activity during complex behaviour.\" Biorxiv (2023): 2023-07.\n\n\t\n**Questions:**\n1. This is an interesting question and something we have been thinking about as well. As a first attempt to answer this question, we ran a new experiment where we try to predict region A from region B for all regions in a test session using a model trained with MtM. As this is different from our current inter-region predictive task (which predicts one region from all other regions), we fine-tuned our pretrained MtM model with a new inter-region task where we randomly dropout some of the regions when predicting one region. The prediction matrix (region A to region B for all regions) is shown in Figure 2 of the one page PDF. It can be seen that there is interesting structure in the matrix where some regions can predict each other well, for example in visual areas. We believe that exploring the \"functional connectivity\" of different areas using MtM is an interesting future work.\n\n2. Good question! Models trained with MtM will be immediately useful for understanding interactions between different brain areas. As discussed above, we were able to construct a \"functional connectivity\" matrix using MtM that can show how predictive different regions are from each other. We believe it may also be possible to examine the attention weights when predicting region A from region B to interpret temporal delays between the two areas. \nUsing MtM to explore common patterns among different individuals is more challenging. An interesting result of our paper is that when pretraining with more individuals, the single individual performance improves after fine-tuning. This indicates that there are common patterns among individuals that are being captured by MtM. If we were to utilize the MtM objective with a latent variable model like LFADS, it may be possible to look at how dynamics in different areas differ among different individuals. This is an exciting figure direction we hope to explore.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for the thoughtful and detailed feedback on our manuscript. We are excited to hear that the reviewers thought that our work "is original and forward thinking" (**JcgB**), represents "deep thinking about foundation models in neuroscience" (**P8rK**), and marks a "step in a very important direction" (**xe9o**). They remarked that "The proposed MtM method is novel" (**BBJQ**), and we provided a "thorough evaluation on the multi-region, multi-subject IBL dataset that provides valuable insights" (**BBJQ**). Multiple reviewers also comment on the clarity of the writing and that "The writing is admirably clear and well referenced." (**JcgB**)

We have done our best in the rebuttal to address any concerns. To this end, we ran four new experiments and added additional figures (in the one page pdf) detailed below:

* An additional population model baseline: LFADS
  * We have added an experiment where we trained the LFADS architecture (a sequential VAE) on 5 test sessions using the temporal masking baseline and our MtM learning objective. The results of the LFADS analysis are shown in Table 1 of our global response. We found that the LFADS architecture trained with MtM outperforms temporal masking on all activity prediction tasks and is competitive with NDT1 trained with MtM (although slightly worse). Similar to our original manuscript, behavior prediction is comparable between MtM and temporal masking (see Section 6.1 and Figure 5 of our manuscript for a more detailed discussion about behavior decoding). This experiment suggests that MtM is an architecture-agnostic learning objective that can be used with both transformers and RNNs. We do not plan to include any other population model baselines as we are focused on models that can do multi-session training such as NDT1, NDT2, and LFADS. We will add details about this LFADS experiment in the updated manuscript. 

    | Table 1 | Co-Smooth | Forward Prediction | Inter-Region | Intra-Region | Choice | Whisker Motion Energy |
    | -------- | ------- | ------- | ------- | ------- | ------- | ------- | 
    | NDT1 Baseline  |  0.55 | 0.17 | 0.18 | -0.49 | 0.87 | 0.65 |
    | NDT1 MtM| 1.00 | 0.28 | 0.83 | 0.70 | 0.85 | 0.67 | 
    | LFADS Baseline | 0.86 | 0.09 | 0.50 | -0.34 | 0.86 | 0.67 |
    | LFADS MtM    | 0.87 | 0.26 | 0.76 | 0.65 | 0.82 | 0.67 |

* Visualizations of the predictions for MtM and temporal masking 
  * We have added visualizations of the rastermaps for NDT1 trained with MtM and temporal masking for each predictive task: neuron, causal, inter-region, and intra-region prediction. As can be seen in Figure 1 of our one page PDF, MtM leads to better predicted rastermaps for inter/intra/neuron masking. For forward prediction, the improvement is more subtle, largely because temporal masking performs well at temporal predictive tasks. Overall, these visualizations give us confidence that MtM is a significant improvement over the baseline. We will include rastermaps for additional test sessions in the updated manuscript.

* Fine-tuning our pretrained MtM model on a monkey recording
  * While the IBL repeated site dataset used for training MtM has high diversity of brain regions, Reviewer BBJQ suggested that it lacks task/species diversity. To demonstrate that MtM can be a useful pretraining strategy for new tasks and species, we fine-tuned our 34-session pretrained NDT1-MtM model on the Neural Latents Benchmark (NLB) MC_RTT dataset. This dataset consists of spiking activity recorded from the primary motor cortex of a monkey performing a self-paced reaching task. We divided this dataset into two "brain regions" based on the NLB heldin and heldout neuron split. We then fine-tuned our model with using the MtM learning objective. With a single private submission to the leader board, the MtM performance (.19 bps) significantly improved upon NDT's previous best performance (.16 bps). We will include all details in the updated manuscript.

* Attempting to run mDLAG on a subset of IBL data
  * To address Reviewer JcgB's suggestion that we test MtM again mDLAG, a multi-region latent variable model, we attempted to run mDLAG on three regions (VISa2/3, VISa4, VISa5) of a single test session. For these regions, there were a total of 42 neurons, 828 trials, and 100 time bins (20 ms bin size). We trained mDLAG for 5.5 hours (5000 EM iterations) and the model did not converge. A visual inspection shows that the learned latent factors for each region were nearly identical and that activity reconstruction from the learned latents consistently led to low co-bps (negative to zero). This experiment provides evidence that scaling mDLAG to full IBL datasets with tens of regions and thousands of neurons is not feasible. We will provide all details for this experiment in the updated manuscript.

* We utilized MtM to estimate a "functional connectivity" matrix for three test sessions.
  * Inspired by Question 1 of Reviewer P8rK, we utilized our pretrained MtM model to ask how well individual brain areas predict each other. We fine-tuned our pretrained MtM model on each test session with a new inter-region task where we randomly dropout some of the regions when predicting one region. Then, we used MtM to predict region A from region B for all regions in the session. The prediction matrix is shown in Figure 2 of the one page PDF. Even without explicit training to predict one region from another region, there is interesting structure in the matrix  as some regions predict each other well, for example in visual areas. This experiment suggests that MtM can be used to explore the "functional connectivity" of different areas. This is an exciting future direction that we would like to explore.
* We created a schematic for the NDT1 and NDT2 architectures to clarify how these architectures tokenize and reconstruct neural data. This can be seen in Figure 3 of our one page PDF. We plan to add a version of this figure to the updated manuscript.

Meta-review: 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.
Decision: Accept (poster)