Towards a "universal translator" for neural dynamics at single-cell, single-spike resolution

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 multitask 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://ibl-mtm.github.io/

Paper

References (42)

Scroll for more · 30 remaining

Similar papers

Peer review

Reviewer P8rK6/10 · confidence 4/52024-06-17

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.

Strengths

1. The deep thinking about foundation model in neurosciece. 2. The model significantly improves generalization by modeling across different brain areas, individuals, and sessions. 3. The multi-task-masking (MtM) strategy effectively enhances the capabilities of the Transformer model.

Weaknesses

1. Lack of Qualitative Analysis: There is a lack of qualitative analysis regarding the neural dynamics obtained from population neural activity across different individuals, such as visualizing features output by NDT. Are there similar patterns? Are there similarities between single-neuron and population-level activities? 2. Insufficient Ablation Studies: a) What are the differences between training on multi-region data versus single-region data? When training with data from three brain areas, is there a phenomenon where the activity of one brain region dominates? b) What are the performance differences when using multi-session data from a single individual versus cross-individual multi-session data?

Questions

1. Typically, cross-individual modeling aims to explore common patterns among different individuals, while cross-area modeling aims to understand information interactions between different brain areas. How much do you think the current foundation model can contribute to these explorations? 2. Generally, in neuroscience, there is a greater interest in understanding the interactions with other internal brain regions, which would be helpful for functional modeling in network level. What insights can the current foundation model provide for research in these areas?

Rating

6

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

see Weaknesses.

Reviewer xe9o6/10 · confidence 4/52024-07-08

Summary

This work proposes a transformer architecture (based on previously existing ones) and, more interestingly, a training procedure that, when applied to spiking neural data, should result in a foundation model for spiking neural data. The clever bit about the procedure is that the learning model is asked to reconstruct differently-masked portions of the data: activity of one held-out neuron, future activity of a neuron, activity of a neuron based on neurons in the same region, activity of a neuron based on activity of neurons in other regions. During training these tasks are presented to the model interchangeably, so it has to learn to do all of them (although with a help of an embedding vector that can decode for the type of the task). The results show that the proposed model achieves better numbers in a variety of tests, however I found it hard to evaluate the significance of the result using the reported metric of bits per spike and how it translates to real-world performance. Please see my comments below and I hope to learn more about this from authors' rebuttal.

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

Questions

184: Please include a formal definition of "bits per spike" Table 1: How a "bits per spike" measure can be negative? What does it mean? Figure 5: Please explain "Choice" and "WME" and axis units in the caption and/or on the figure. Perhaps it would be better to fix the range on the axis, so that one could easily estimate also the variability between the sessions and MtM improvement in the sessions.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

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?

Reviewer xe9o2024-08-08

Figures 1 and 3 in the rebuttal PDF are very nice and informative. Thank you for providing those. I think Figure 1 should even be part of the manuscript itself (with more visualizations in the supplementary materials) Q1: It seems to me that MtM smoothes adjacent time bins quite a lot. Ground truth often has almost empty time bins, followed by a sharp increase in the firing rate right in the next bin. But in MtM's reconstruction the change of firing rate is always gradual from bin to bin. And that worries me because sharp increases and sudden spikes or drops are very important for actual neuroscientific insight. Why dies this smoothing occur? Is adjacency in the time series plays some sort of a special role in a transformer? Or is it due to the loss function? Can you define a loss function that would strive to preserve sharp changes? Q2: I am still somewhat unhappy with the "bits per spike" measure :), however seeing the actual definition was very helpful and for sure needs to be part of the manuscript. The issue I take with it - what is the benefit of adding this additional layer of abstraction and measure via Poisson model? Why not go super simple and just count the number of spikes the prediction got wrong? (and then average in a suitable manner to make comparable across). Spikes are binary, and this makes it easy to count them. Maybe it's just me, but when I see a metric like "bps: 0.51" I have no idea how good or bad is it. Maybe if you provide info on what the bps would be for a perfect reconstruction? I guess what I am looking for is a more intuitive number here, like "73% of spikes are where they should be, 27% are missing / extra" or "the number of spikes MtM predicts within each bin is off by 2.7 spikes per bin, or 14%". Basically I am still not able to understand when I am looking at "Baseline bps: 0.03" vs. "MtM bps: 0.51" should I be impressed? How impressed should I be? Reconstruction-wise both Baseline and MtM look pretty far away from the ground truth. Q3 (exploratory, just for fun): Do I understand correctly that currently tokens are based on firing rates? Is yes, then my question is would it be possible to use analgues of "spiking words" or "spiking syllables" of fixed length instead? Something like ``` ||||__|||| -> token 1 |||||||___ -> token 2 ___||||___ -> token 3 ... etc ``` Where "|" are spikes and "_" are non-spikes. To basically build a vocabulary of all possible 10ms spiking patterns and use those as tokens? Or, actually, why use bins at all? What would happen if you use spiking sequences as-is, so that the time step is small enough to just be either | or _ ?

Authorsrebuttal2024-08-08

Response to additional questions from Reviewer xe9o - Part 1

We thank the reviewer for getting back to us so promptly! We agree that Figure 1 should be a part of the manuscript and we will include a version of this figure in the main text of the updated manuscript. We believe there might be a misunderstanding about the goals of statistical models of neural data that we hope to clarify in the following response. Our response will be in two parts due to character constraints (we will answer Q2 and Q1 in the first comment). If we answer all questions satisfactorily, we hope the reviewer would consider raising their score. **Q2** We are glad the reviewer found our definition of bps helpful and we will make sure to add this in the updated manuscript. We suspect there might be a misunderstanding here about prior literature on neural population models. As described in [1], there is a rich literature of population models that utilize either Gaussian or, more recently, Poisson observation models (please see Table 2 of [1]) to predict neural spiking activity. The goal of neural population models is to predict the underlying firing rates of each neuron because the exact spike times are believed to be observations from a stochastic process. Similar to this seminal paper [2], we choose to model the observed spikes as samples from an inhomogeneous Poisson process whose rate corresponds to the inferred firing rate for the given neuron. To measure the quality of fit of our Poisson process, it is most natural to use a version of the log-likelihood such as the bps metric utilized in the Neural Latents Benchmark [3]. Using this commonly agreed upon metric makes sure that our work is consistent and comparable with existing research on neural population models. In addition, the "bits per spike" metric is closely related to the "deviance", a standard metric generally used in Statistics, e.g., generalized linear models (GLMs) use deviance as a goodness-of-fit measure for a statistical model. The deviance compares the goodness-of-fit of the model of interest, e.g., MtM, to a baseline null model, where the goodness-of-fit is measured by the model log likelihood. The bps is simply a normalized version of the deviance metric, which compares the model predictions to the average firing rate of the neuron for the trial. The bps further normalizes the deviance by the spike count so that the metric can be comparable across neurons regardless of whether the neurons are active or inactive. Along with the bps, we would be happy to report the R-squared in the revision to provide further intuition, but note that the R-squared depends strongly on bin size (as bin size gets small so will R-squared, since the precise timing of spikes is likely not predictable in this setting). **Q1** In Figure 1 of the one page PDF, the ground-truth rastermaps are the spike counts of the trial while the MtM/NDT rastermaps are the inferred firing rates. These quantities are not directly comparable because we expect the underlying firing rate of a neuron to vary smoothly while its observed spikes (which are assumed to be stochastic) may have sharp increases. A great example of this phenomena is shown in [4] where they demonstrate that when the inferred firing rates exactly fit each spike time, this is actually overfitting to the neural data. They introduce a novel dropout strategy, similar to MtM's neuron masking, that reduces this overfitting. In fact, it has been shown by [2] (in Figure 2e) that the smoothed firing rates from a Poisson model actually allow for decoding animal behavior more accurately than the raw spike counts. Whether to smooth or not involves a bias-variance tradeoff. More smoothing means more bias and less variance. When we optimize for prediction accuracy, we are implicitly optimizing this tradeoff, and empirically we (and others) find that some smoothing (bias) leads to better predictions. [1] Hurwitz, Cole, et al. "Building population models for large-scale neural recordings: Opportunities and pitfalls." Current opinion in neurobiology, 2021. [2] Pandarinath, Chethan, et al. "Inferring single-trial neural population dynamics using sequential auto-encoders." Nature methods, 2018. [3] Pei, Felix, et al. "Neural latents benchmark'21: evaluating latent variable models of neural population activity." arXiv 2021. [4] Keshtkaran, Mohammad Reza, et al. "A large-scale neural network training framework for generalized estimation of single-trial population dynamics." Nature Methods, 2022.

Authorsrebuttal2024-08-08

Response to additional questions from Reviewer xe9o - Part 2

**Q3** Before answer the question, we want to clarify that MtM is agnostic to the underlying transformer architecture and that the tokenization scheme will depend on the architecture that is used. To answer the question, a token for NDT is the observed spike counts for a single 20ms timestep of the population (please see Figure 3 in [5]). This means that NDT is using the exact spike times (up to 20 ms resolution) to fit the model. The reviewer's suggested tokenization scheme would be comparable to NDT with more fine binning (1 ms bins). While this is feasible, given the sparsity of neural data, this would give rise to many more empty bins and would also be computationally prohibitive (i.e., self-attention scales quadratically with the number of tokens). As the reviewer mentioned, "why use bins at all?". Treating each spike as a token, as done in the POYO paper [6], is also a potential solution. As this is computationally prohibitive, they adopted a perceiver IO architecture [7] so the attention is only applied to a set of latent tokens, which is much lower dimensionality. While this architecture is a promising step in the right direction, it has only ever been used as a supervised model to decode behavior from spikes. Adapting POYO and more interesting transformer architectures to be self-supervised such that they can be used with MtM is an exciting future direction. [5] Ye, Joel, and Chethan Pandarinath. "Representation learning for neural population activity with Neural Data Transformers." arXiv, 2021. [6] Azabou, Mehdi, et al. "A unified, scalable framework for neural population decoding." NeurIPS 2024. [7] Jaegle, Andrew, et al. "Perceiver io: A general architecture for structured inputs & outputs." arXiv, (2021).

Reviewer JcgB6/10 · confidence 4/52024-07-13

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?

Rating

6

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

The authors adequately state the limitations of their modeling approach, training paradigm, and datasets.

Reviewer BBJQ6/10 · confidence 4/52024-07-14

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.

Strengths

1. The 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. 2. The proposed MtM method is novel and demonstrated effectiveness over the existing temporal masking method in learning different structures of neural activity. 3. Thorough evaluation on multi-region, multi-subject IBL dataset provides valuable insights for the neuroscience community. 4. The manuscript is well-written. Text and figures are logical and easy to follow.

Weaknesses

1. The number of baselines is quite limited. Only two transformer architectures (NDT1 and NDT2) were benchmarked, and the proposed MtM method is mainly compared with one baseline of temporal masking. It appears that MtM could be used as the training objective for similar existing models that also assume a Poisson emission model, e.g. RNN-based LFADS (Pandarinath et al. 2018) and other models in the NLB benchmark (Pei et al. 2021). Evaluating MtM on more architectures could be helpful to gauge if MtM effectiveness is truly model-agnostic as the authors noted or is biased more toward certain type of architecture. 2. The dataset only consists of one behavior task (decision making), while existing pretrained models POYO and NDT2 referenced in the paper were trained/evaluated on multiple behavior (motor) tasks, even across species (monkey vs. human). It is unclear if and how the community could benefit from finetuning the proposed pretrained model on other datasets with a different behavior paradigm. 3. Some experiment details are missing or not presented clearly (elaborated in Questions section below).

Questions

1. It would be helpful to compare the proposed model with more NLB baselines to justify the architecture choice. Does NDT1 with MtM significantly outperform simpler baselines like RNN with MtM to justify the computation cost? Even knowing how well baselines like MLP at directly decoding the behavior could help readers better evaluate the strength of the proposed model. 2. It is not clear if the pretrained model would provide some leverage on other datasets with different behavior structures. Would it outperform a model trained from scratch on the new dataset? 3. It might be worth making it clear that “task” in “multi-task” used in the paper refers to the machine predictive task (decoding choice and motion energy), rather than animal behavior task (decision making – the only behavior investigated). 4. Equation (1): Does the transformer use any kind of positional encoding? If so, what type of positional encoding? 5. For finetuning on a held-out session, I assume the reported metrics were calculated on the test trials within the session. How many trials are there in each train/val/test split of each session? 6. Line 174: possibly a typo, should be 1% of the trial (20ms). 7. How was choice decoded from $\hat{X}$? I assume from Figure 1 that $\hat{X}$ is a N x T matrix, where N is number of neurons and T is number of timestep. If so, at which timestep of $\hat{X}$ that the linear layer was applied to decode the choice? 8. Similarly, was the motion energy decoded at each timestep of $\hat{X}$ or at some particular point during the trial? 9. Figure 5: what is the chance level for each session? For imbalance datasets, the chance level might not be 0.5, especially if the data was preprocessed for quality control in some ways. If so, it’s worth plotting the chance level in the figure.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors have addressed the limitations in the manuscript.

Reviewer P8rK2024-08-08

The rebuttal PDF meets my expectations. Although the learning of neural dynamics is implicit, some patterns appear to be discernible. Please consider further joint analysis with interpretable methods in the future work, otherwise this kind of work is useless for neuroscience. I would like to maintain my score. Good Luck!

Authorsrebuttal2024-08-08

Applying MtM to interpretable models of neural activity (in Global Response)

We thank the reviewer for responding back to us so quickly. We believe that MtM can easily be combined with more interpretable models of neural populations. For example, as part of our Global Response, we combined our MtM learning objective with LFADS, an interpretable, explicit dynamics model. We were able to improve the performance of LFADS on all predictive tasks introduced in the paper. Although we have no more space for figures, interpreting the latent factors learned by LFADS when trained with MtM is an exciting future direction and something we can do in the updated manuscript. We hope the ability of MtM to be applied to interpretable latent factor models like LFADS can improve the reviewer's score of our manuscript. Although explicit dynamics models are very important for neuroscience, we still believe that purely predictive models like transformers are also useful. Neural decoding/encoding algorithms have been developed for many years to determine how much information about a stimulus or behavior is encoded in a neural population [1]. Much like these algorithms, we are interested in determining how much information each brain area has about other brain areas. Understanding which brain areas predict each other best can allow us to make hypotheses about distributed neural computation. [1]​​ Paninski, Liam, Jonathan Pillow, and Jeremy Lewi. "Statistical models for neural encoding, decoding, and optimal stimulus design." Progress in brain research 165 (2007): 493-507.

Reviewer JcgB2024-08-12

Thank you for the thoughtful rebuttal. I believe the final product will indeed be improved by incorporating the ideas from these reviews and rebuttals. I stand by my score.

Authorsrebuttal2024-08-12

Response to Reviewer JcgB

We thank the reviewer for their positive response to our rebuttal. If there is anything we can further clarify or do to improve the reviewer's overall assessment of our work, please let us know.

Authorsrebuttal2024-08-13

Response to Reviewer BBJQ

We hope the reviewer found our rebuttal helpful. If there is anything we can further clarify, we are happy to do so!

Reviewer BBJQ2024-08-14

Thank you for the explanation. The extra experiments alleviated my concerns. I've increased my rating accordingly.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC