Title: Towards a "Universal Translator" for Neural Dynamics at Single-Cell, Single-Spike Resolution
Venue: NeurIPS 2024
Accepted: yes
Keywords: neural dynamics

transformer

self-supervised learning

pretraining

multi-task learning

masked modeling

brain-computer interfaces
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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?\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nsee Weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-06-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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\n\nWEAKNESSES:\n* 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\"\n\nQUESTIONS:\n184: 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.\n\nLIMITATIONS:\nThe 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?", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-08"}

comments: 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.

STRENGTHS:
The work is original and forward thinking. The writing is admirably clear and well referenced.

WEAKNESSES:
The paper would benefit from:
- Comparing to a wider range of models of neural population dynamics (e.g., LFADS, GPFA, mDLAG, etc). Many previous approaches have reported co-smoothing performance (e.g., Neural Latents Benchmark). How does MTM-NDT1 compare against those previous approaches on co-smoothing? How does the proposed approach compare to multi-area models (e.g., mDLAG) for on the inter-region co-smoothing task?
- More fully describing the details of the “temporal masking” baseline. The details are referenced out to a citation rather than described within the text. Since this baseline appears throughout all of the results, it seems important to describe it within this paper.

QUESTIONS:
- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1.\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.\n\nWEAKNESSES:\n1.\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).\n\nQUESTIONS:\n1.\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.\n\nLIMITATIONS:\nThe authors have addressed the limitations in the manuscript.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}