SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mehdi Azabou $^{1,*}$ , Vinam Arora $^{1}$ , Venkataramana Ganesh $^{1}$ , Xineng Mao $^{2,3}$ , Santosh Nachimuthu $^{1}$ , Michael J. Mendelson $^{1}$ , Blake Richards $^{2,4}$ , Matthew G. Perich $^{2,3}$ , Guillaume Lajoie $^{2,3}$ , Eva L. Dyer $^{1,*}$ $^{1}$  Georgia Tech,  $^{2}$  Mila,  $^{3}$  Université de Montréal,  $^{4}$  McGill University
Our ability to use deep learning approaches to decipher neural activity would likely benefit from greater scale, in terms of both model size and datasets. However, the integration of many neural recordings into one unified model is challenging, as each recording contains the activity of different neurons from different individual animals. In this paper, we introduce a training framework and architecture designed to model the population dynamics of neural activity across diverse, large-scale neural recordings. Our method first tokenizes individual spikes within the dataset to build an efficient representation of neural events that captures the fine temporal structure of neural activity. We then employ cross-attention and a PerceiverIO backbone to further construct a latent tokenization of neural population activities. Utilizing this architecture and training framework, we construct a large-scale multisession model trained on large datasets from seven nonhuman primates, spanning over 158 different sessions of recording from over 27,373 neural units and over 100 hours of recordings. In a number of different tasks, we demonstrate that our pretrained model can be rapidly adapted to new, unseen sessions with unspecified neuron correspondence, enabling few-shot performance with minimal labels. This work presents a powerful new approach for building deep learning tools to analyze neural data and stakes out a clear path to training at scale.
Recent advances in machine learning, particularly in the context of large-scale pretrained models like GPT [1, 2, 3, 4], have showcased the immense potential of scaling up, both the terms of the size of datasets and models [5, 6]. Similarly, in neuroscience, there is a growing need for a foundational model that can bridge the gaps between diverse datasets, experiments, and individuals, allowing for a more holistic understanding of brain function and information processing [7]. The development of such a model would allow researchers to uncover underlying patterns and interactions within neural populations [8], and potentially allow for more robust decoding of brain states.
However, creating a large-scale neural decoding model that can effectively combine spiking datasets from various sources is a complex challenge [9]. One of the central challenges is the lack of a shared "vocabulary" in neural recordings. Unlike the case for text—wherein every document written in a given language shares a basic lexicon for tokenization—there is no one-to-one correspondence between neurons in different individuals. As such, every recording from a different individual involves a unique set of neurons that cannot be easily aligned with another set. An additional core challenge lies in the inherent variability of neuron sets observed across different days [10]. Even when monitoring the same individual, variability in electrode/tissue interface can lead to distinct
*Contact: {mazabou,evadyer} @gatech.edu. Project page and code: https://poyo-brain.github.io
neuron sets which, despite advanced sorting methods, can lead to inconsistencies in input channels across sessions [11, 12, 9]. Overall, this lack of correspondence across recordings complicates the integration of information from different experiments and individuals, ultimately hampering efforts to construct a unified perspective on population-level interactions and dynamics in the brain.
In response to these challenges, we propose a new framework for large-scale training on neural spiking data called POYO (Pre-training On manY neurOns). This framework is designed to enable scalable and efficient training across multiple sessions of neural recordings, even when spanning different sets of neurons with no known correspondence. Our approach centers around a novel tokenization scheme that transforms individual neural action potentials, or "spikes", into discrete tokens, thereby preserving the neural code's finest temporal structure while simultaneously enhancing computational efficiency. The resulting tokenization not only allows for a more effective representation of neural activity but also paves the way for training on larger volumes of data. We combine this input tokenization method with an architecture that builds on the PerceiverIO [13] to compress the input spikes into a latent space and learns interactions across spikes in time and across neurons.
We evaluate the performance of our proposed approach on data from over 158 sessions from open electrophysiology datasets from seven non-human primates (NHPs), spanning over 27,373 units and 100 hours of recordings. We demonstrate that through pretraining on large amounts of data, we can transfer with very few samples (few-shot learning) and thus improve overall brain decoding performance. Our work not only presents an innovative framework for training large models on neuroscience datasets, but also offers insights into the scaling laws that govern decoding from neural populations. By enabling the development of large pretrained models for neural decoding, our approach advances the field of brain-machine interfaces and other decoding applications.
The main contributions of this work include:
- A framework for large-scale training on neural recordings: We present a novel framework for training transformer models end-to-end on multi-session and across-individual electrophysiology datasets derived from neural populations, enabling efficient and effective decoding from a diverse range of neural recordings.
- Innovative spike-based tokenization strategies: We introduce a fundamentally different way to tokenize neural population activity. Our approach tokenizes individual spikes (events) across neural populations, preserving fine temporal structure and enhancing computational efficiency by adopting a sparse representation of the data.
- Pre-trained models for neural decoding: We build two large pretrained models (POYO-1, POYO-mp) that can be fine-tuned on new sessions and across recordings from different animals and new behavioral tasks. We will make the weights and code available, and provide both pretrained models as a resource to the community.
2 Approach
The transformer architecture [14], originally introduced in the context of natural language processing (NLP), has shown remarkable flexibility and effectiveness in various domains, especially in the presence of large and diverse datasets [15, 3]. In this work, we explore how to leverage this versatility in the neural data domain.
2.1 Tokenizing neural population activity
Neurons communicate asynchronously using electrical impulses called spikes. The timing and frequency of spikes encode signals that convey information about the external world and coordinate the internal dialogue within the brain. In many neuroscience experiments, neural activity is recorded from the brain through multiple electrodes, and then processed [16] to extract the spiking events for a set of neural "units" [17].<sup>3</sup> The resulting data are multi-variate event-based sequences that are typically very sparse relative to the number of recorded time points.
$^{2}$ Poyo is the exclamation used by Kirby, who has an insatiable appetite. Similarly, POYO consumes spikes from many neural datasets and combines them into one unified model.
3We use the term units instead of neurons because recordings may include both single-units (isolated neurons) and multi-units (contain spikes from multiple neurons).
Figure 1: Overview of our approach. Input spike tokens are compressed into a smaller set of latent tokens which are processed through multiple self-attention blocks, finally time-varying outputs are predicted by querying the latent space. All tokens in this model are assigned a timestamp, which are used towards rotary position encoding.
Neurons and their spikes are of course not independent of one another. Rather, they are part of a complex, interwoven tapestry of neural activity, with each spike contributing to a collective dialogue across billions of neurons. The true challenge and opportunity lie in interpreting these spikes not in isolation, but in the context of this broader "conversation."
When trying to scale up and train on many sessions of data, we are faced with the challenge of having recordings of neural conversations with completely different set of "speakers" for which we don't know the identity or their functional tuning (or what they respond to). This is because each time we record from the brain, we are tuning into a conversation between a new set of speakers. However, as with language, there is some reason to think that neurons are ultimately communicating on similar topics, e.g. sensations from the external world, internal states of the body, muscle commands, etc. In other words, the lexicon may be different, but everyone is talking about similar subjects. Despite the variability of our data, our aim is to decipher these "conversations" in a way that generalizes to new neural datasets without fixed or known correspondence across their inputs.
The neural tokenizer. Based upon these motivations, we propose a novel approach for tokenizing neural population activity where each spike is represented as a token (see Figure 1). In this case, each token can be defined in terms of which unit it came from (via a learnable embedding) and the time that the spike event was detected. By representing our data in this way, we never define or fix the expected number of units, and the model can ingest populations of arbitrary size, and thus can be trained across many datasets. At the same time, this approach also avoids having to specify a specific time resolution for the neural code, as is the case with binning, and gets rid of the accumulation of a lot of sparse tokens containing no events.
More concretely, we assign each unit a unique identifier and a corresponding  $D$ -dimensional learnable embedding. Let  $\mathrm{UnitEmbed}(\cdot)$  denote a lookup table that associates each unit to its unit embedding. Akin to word embeddings that encode semantic meaning and relationship between words, the unit embedding space encodes something about the "speaker's" identity and role in neural computation. A unit fires a sequence of spikes in a context window of time  $[0, T]$ . Each spike will be represented by a token characterized by  $(\mathbf{x}_i, t_i)$ , where
$$
\mathbf {x} _ {i} = \operatorname {U n i t E m b e d} (\text {s p i k e i ' s u n i t})
$$
is the learned embedding associated with the unit that emitted the spike and  $t_i$  is the event timestamp.
Collectively, and for an arbitrary set of units, we combine all the spikes into a sequence of length  $M$ . The neural population activity is represented by  $[\mathbf{x}_1,\dots ,\mathbf{x}_M]$  and their corresponding times  $[t_1,\ldots ,t_M]$ . While the size of our context window  $T$  stays fixed,  $M$  will vary depending on the number of units in the population and their firing rates (see Figure 1). Note that all spikes from a specific unit have the same unit embedding, and only differ in their timing.
2.2 Building the latent space
Compressing the input sequence. Rather than processing the input sequences with self-attention, which is quadratic in sequence length and can be expensive for long sequences, we use a perceiver encoder module [18] that summarizes the input sequence into a shortened "latent" sequence. Let's consider a sequence of  $M$  input tokens  $\mathbf{X} = [\mathbf{x}_1,\dots ,\mathbf{x}_M]$  and a sequence of learned latent tokens  $\mathbf{Z}_0 = [\mathbf{z}_{0,1},\ldots ,\mathbf{z}_{0,N}]$ , where  $\mathbf{z}_{0,i}\in \mathbb{R}^{D}$  and  $N\ll M$ . We use cross-attention to pull information from the input sequence into a compressed latent sequence. The queries  $\mathbf{Q} = \mathbf{W}_q\mathbf{Z}_0$  are a projection of the learned latent tokens  $\mathbf{Z}_0$  and the keys and values are a projection of the input tokens:  $\mathbf{K} = \mathbf{W}_k\mathbf{X}$  and  $\mathbf{V} = \mathbf{W}_v\mathbf{X}$ , respectively. The cross-attention operation can be expressed as follows:
$$
\mathbf {Z} _ {1} \leftarrow \operatorname {C r o s s - A t t n} (\underbrace {\mathbf {Q}} _ {\downarrow}, \underbrace {\mathbf {K}} _ {\downarrow}, \underbrace {\mathbf {V}} _ {\downarrow}) = \operatorname {s o f t m a x} \left(\frac {\mathbf {Q} \mathbf {K} ^ {T}}{\sqrt {d _ {k}}}\right) \mathbf {V}. \tag {1}
$$
We use the standard transformer block with pre-normalization layers and feed-forward nets.
Attention in the latent space. Following the compression of the input sequence through cross-attention, we apply multiple self-attention blocks on the latent token sequence, which now costs  $O(N^2)$  instead of  $O(M^2)$  with  $N \ll M$ . Let  $\mathbf{Z}_l$  denote the embedding of the latent tokens at the  $l$ th layer. The self-attention operation can be expressed as follows:
$$
\mathbf {Z} _ {l + 1} \leftarrow \text {S e l f - A t t n} (\overbrace {\mathbf {Q} , \mathbf {K}} ^ {\mathbf {Z} _ {l}}, \mathbf {V}) = \operatorname {s o f t m a x} \left(\frac {\mathbf {Q} \mathbf {K} ^ {T}}{\sqrt {d _ {k}}}\right) \mathbf {V} \tag {2}
$$
We denote the final latent sequence obtained after  $L$  layers as  $\mathbf{Z}_L = [\mathbf{z}_{L,1},\dots ,\mathbf{z}_{L,N}]$
2.3 Encoding relative timing information
To incorporate timing information, we leverage rotary position encoding (RoPE) [19] across all attention layers in the architecture. Unlike traditional position encoding methods that inject absolute position information into the token's embedding, RoPE applies a rotation operation in the query, key embedding space, allowing each token in a sequence to attend to others based on its relative positional information.
Recall that each input token  $i$  has an associated timestamp  $t_i$ . We will also assign a timestamp to each of the latent tokens: we divide the latent tokens into groups of equal size and spread each group uniformly over the context window  $[0, T]$ . This method allows us to capture temporal relationships between the latent tokens and the input tokens, enabling a temporal understanding of the encoded sequence. By distributing the latent tokens evenly within the context window  $[0, T]$ , we create a structured temporal representation, which preserves and propagates temporal information all the way through the model. A detailed formulation of the relative encoding can be found in Appendix A.1.
2.4 Querying the latent space
Having built a latent space which can encode and model any population of units, we now want a flexible way to readout behavioral variables. The output of the latent encoder is the latent sequence of size  $N$ , while the desired output can be any arbitrary sequence of length  $P$ . Let us consider the task of hand velocity decoding for example, in the context window  $[0, T]$ , the length of the output sequence will depend on the sampling frequency. Since we aspire to train on datasets sourced from various labs, the sampling rate can differ significantly. We thus need a flexible mechanism for predicting outputs of varying lengths and querying from the neural activity at specific points in time.
We define a sequence of output tokens  $\mathbf{Y}_0 = [\mathbf{y}_{0,1},\dots ,\mathbf{y}_{0,P}]$ , where  $P$  is the number of output time points, which can change across sequences. Each output token is defined by  $(\mathbf{y}_{0,i},t_i^{out})$ . Initially all the output tokens within a session are set to the same learned embedding  $\mathbf{y}_{0,i} = \mathbf{y}_0\in \mathbb{R}^D\forall i$ . The output  $\mathbf{Y}$  is obtained by querying the latent space through cross-attention:
$$
\mathbf {Y} \leftarrow \operatorname {C r o s s - A t t n} (\underbrace {\mathbf {Y} _ {0}} _ {\downarrow}, \underbrace {\mathbf {Z} _ {L}} _ {\downarrow}, \underbrace {\mathbf {V}} _ {\downarrow}) = \operatorname {s o f t m a x} \left(\frac {\mathbf {Q} \mathbf {K} ^ {T}}{\sqrt {d _ {k}}}\right) \mathbf {V} \tag {3}
$$
Since we use rotatory embeddings, and both the latent and output tokens have assigned timestamps, the querying mechanism can leverage the relative position of latent and output tokens to extract the temporally relevant context that enable prediction of the behavioral variable of interest.
Session embeddings. To account for variability of the experimental setups in real world settings, we propose to define a learnable session embedding which captures the hidden experimental variables. This information is injected into the output query  $\mathbf{y}_0$ , where again we use a lookup table to register each new session we encounter. To illustrate how we probe the model to generate an output, we can think about it as asking a question of the form:
"Predict [BEHAVIOR] at time [TIMESTAMP] under the experimental conditions of [SESSION ID]".
This produces a set of output tokens of the dimension of the number of queries which we then pass through an MLP to generate the behavioral variables of the desired dimension (e.g., 2D for hand velocities in a planar movement task).
2.5 Unit identification in new sessions
Our design of the unit embedding space allows our model to learn latent information about the units it encounters, as well as capture the relationship between units in the population. Given a new recording with unidentified units, we can transfer our model by mapping these new units into the unit embedding space. To do this, we introduce an approach that we call "unit identification", which leverages gradient descent to learn the embeddings of new units. In this approach, we freeze all existing weights of the model and simply add new rows to the UnitEmbed(\cdot) lookup table for each of the new units, as well as a new session embedding. Notably, the bulk of our model which maps the neural population activity to behavior is unchanged and is simply transferred to the new dataset. In our experiments, we find that this approach is surprisingly effective and allows us to rapidly integrate new datasets into the same underlying model. Further details and analysis on the robustness of this approach are presented in Appendix C.2.
In this section, we demonstrate the promise of our approach for large-scale training and examine the benefit of scaling in neural population decoding.
3.1 Datasets and experiment setup
One of the key advantages of our approach is its ability to scale to handle large amounts of neural data, including sessions from different numbers of neurons, across different tasks and recording setups, and from different animals. Thus we set out to build a diverse dataset large enough to test our approach. We curated a multi-lab dataset with electrophysiological recordings from motor cortical regions, where neural population activity has been extensively studied [10], and deep learning tools and benchmarks have recently been established [27]. In total, we aggregated 178 sessions worth of data, spanning 29,453 units from the primary motor (M1), premotor (PMd), and primary somatosensory (S1) regions in the cortex of 9 nonhuman primates (see Table 1). We place this in the context of standard analyses within a single lab or paper which typically involve 10's of sessions and a few hundred neurons.
All of these neural recordings were collected while the animals performed various motor tasks that vary in their inherent complexity (see Figure 2A-B). The center-out (CO) task is relatively stereotyped, with the animal making a reach to one of eight targets after receiving a go cue, and
Study	Regions	# Indiv	# Sess	# Units	# In	# Out	Tasks
Perich et al. [20, 21, 22, 23]	M1, PMd	4	117	11,557	143M	20M	CO, RT
Churchland et al. [24]	M1	2	9	1,728	706M	87M	CO
Makin et al. [25]	M1, S1	2	47	14,899	123M	15M	RT
Flint et al. [26]	M1	1	5	957	7.9M	0.3M	CO
NLB-Maze [27]	M1	1	1	182	3.6M	6.8M	Maze
NLB-RTT [27]	M1	1	1	130	1.5M	2.8M	RT
Table 1: Datasets used in this work. CO: Center-Out, RT: Random Target.
Figure 2: Building a multi-session model spanning multiple animals and tasks. (A) Center-out reach and (B) Random target task [20], along with examples of true and predicted behavior (x,y velocities). In (C-D), we show the decoding performance for single-session models (gray) and the P0Y0-mp multi-session model (green).
then returning to the center. In contrast, the random target (RT) task is significantly more complex. The animals make continuous and self-paced movements with new targets appearing in succession at random locations in the workspace. In addition to the greater exploration of the workspace and heterogeneity in movements, this task allows individual to plan their next movement while finishing the execution of the current movement leading to greater complexity in neural dynamics. Other sources of variability that we find across labs include the: choice of pre-processing algorithms (spike sorting or threshold crossing analysis), type of controller (manipulandum, touch screen), and sampling rate when recording behavior (100Hz to 1kHz). We do not re-process the data or attempt to standardize it across labs or tasks. Further details on the datasets are provided in Appendix B.
Experiment setup. Throughout all of our experiments, we use a context window of  $1s$  and do not segment data into trials during training. We only use the trial structure when reporting the decoding performance, in particular, center-out sessions are evaluated during the reaching movement. We train the model with  $N = 512$  latent tokens and a dimension  $D = 128$ . We use the LAMB optimizer [28], and employ a cosine decay of the learning rate at the end of training. For every session, we holdout  $20\%$  of the trials for testing, and  $10\%$  for validation. We use 1-GPU and 8-GPU setups for single-session and multi-session models, respectively. All details can be found in Appendix A.
3.2 Testing the model on single sessions
To first investigate the performance of our architecture when trained on a single session, we trained single-session models on 100 different recording sessions acquired from three nonhuman primates (Monkey C, Monkey M, Monkey Ja), each containing anywhere from 10 to 106 minutes of data [29]. In all of these recordings, the same behavioral task setup, behavioral recording apparatus (10 ms temporal resolution), and spike sorting procedure was used.
Across all 100 single-session models, we obtained an average R2 of 0.9347 on CO and 0.8402 for RT sessions. When we compared these single-session results with existing models for neural decoding, including a Wiener Filter, GRU and MLP [30], we found that our approach consistently outperforms these baselines, with even greater improvements observed on the RT task. Additionally, we found that our single-session models are stable for a wide range of hyperparameters (Appendix C.3).
3.3 POYO-mp: Building a large pretrained across-animal, multi-session model
To investigate the question of how training with more sessions of data can improve brain decoding, we trained a large model (POYO-mp, 24 layers) on all 100 sessions that we studied in our single-session analysis. In total, we used 9,789 units and 4,367 neuron-hours (number of neurons  $\times$  amount of time recorded) to train our POYO-mp model.
Method	Same animal, New day Monkey C - CO (2)	New animal
Monkey T - CO (6)	Monkey T - RT (6)
From scratch	Wiener Filter	0.8860 ± 0.0149	0.6387 ± 0.0283	0.5922 ± 0.0901
GRU	0.9308 ± 0.0257	0.8041 ± 0.0232	0.6967 ± 0.1011
MLP	0.9498 ± 0.0119	0.8577 ± 0.0242	0.7467 ± 0.0771
POYO-[Single-session]	0.9682 ± 0.0111	0.9194 ± 0.0185	0.7800 ± 0.0702
Pre-trained	POYO-mp + Unit ID	0.9675 ± 0.0079	0.9012 ± 0.0271	0.7759 ± 0.0471
POYO-mp + Finetune	0.9708 ± 0.0116	0.9379 ± 0.0193	0.8105 ± 0.0561
POYO-1 + Unit ID	0.9677 ± 0.0096	0.9028 ± 0.0248	0.7788 ± 0.0548
POYO-1 + Finetune	0.9683 ± 0.0118	0.9364 ± 0.0132	0.8145 ± 0.0496