SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 2: Behavioral decoding results across neural recordings from two nonhuman primates performing two different tasks. All the baselines and the single-session model are trained from scratch, while P0Y0-mp and P0Y0-1 are pretrained. The standard deviation is reported over the sessions. The number of sessions in each dataset is contained in  $(\cdot)$  and the top performing models in each category are indicated in boldface (first) and underlined (second).
This model achieves an average test  $\mathrm{R}^2$  of 0.9512 on the center-out tasks and 0.8738 on the random target tasks, which is a marked improvement over the single-session average. When we compared the performance of our multi-session model to single-session models head-to-head (see Figure 2C-D), we found that across the board, multi-session training improves over single-sessions, and we observed even greater improvements for datasets with fewer trials (indicated by the size of the circles in the plot). On RT sessions, we observe an even bigger improvement over the single session models. These results suggest that there are significant benefits in joint-training across multiple sessions, especially when decoding more complex behaviors like the random target task.
Scaling analysis. By enabling multi-session training, we can start to ask questions about the extent to which having more data or more parameters can improve decoding (Figure 3). Thus, we studied the improvements obtained for three different depths of models, with  $\mathrm{L} = 6$  (3.8M),  $\mathrm{L} = 14$  (7.4M), and  $\mathrm{L} = 26$  (13M) layers (parameters). For both CO and RT tasks, we find that even at the same depth  $(\mathrm{L} = 6)$ , our 1 multi-session model shows an improvement over the average performance of 100 single-session models. We see further improvements with increasing the model depth, with the RT task benefiting even fur
Figure 3: Scaling curves.  $\mathbb{R}^2$  is evaluated and averaged over the test splits on both CO and RT tasks. The single-session performance is the mean over 100 different models.
ther from scaling. When we fix the model depth to  $\mathrm{L} = 6$ , and study how training jointly with more data contributes to improvement in performance, we find an improvement in performance when comparing single-session training (average 1.2M tokens) to training jointly on all RT sessions (26M tokens). When we also train jointly with other CO sessions (134M tokens in total), and increase model depth to help in accommodating the more growing and diverse training set, we continue seeing improvement in the decoding performance on RT sessions, accumulating to  $4\%$  over our single-session baseline.
3.4 Transferring to new sessions
After pretraining, we can then test on new sessions with unknown neurons using either the (i) unit identification approach we described in Section 2.5, or (ii) full finetuning of the model weights. Recall that when we use a unit identification approach, we freeze the weights of the model and only learn new unit and session embeddings.
Results on held out sessions from the same animals. We first tested the unit identification approach on held-out sessions from Monkey C that are unseen during training (Table 2, Left). In this case, we don't have correspondence between units in the training and the testing conditions. Surprisingly, we find that we can achieve comparable performance with unit identification on the pretrained model (0.9675) with that of the single-session models trained fully from scratch (0.9682). With further finetuning of the rest of the model's weights, we can improve the accuracy further (0.9708) to go beyond the accuracy of the single-session models. This highlights the robustness
(a) Monkey C, heldout CO session
(b) Monkey T (new subject), RT session
Figure 4: Sample and compute efficiency for training, unit identification, and fine-tuning approaches. On the left, we show the sample and compute efficiency for a heldout CO session. On the right, we plot the sample and compute efficiency for a new animal not seen during training.
of our model and its flexibility to accommodate fresh data with even a simple input mapping (unit identification).
We compare with the single session model and a GRU baseline (Figure 4A) with our multi-session model trained on the same number of trials. Both our single-session and finetuning approach achieve good performance scaling as we increase the number of samples, with the finetuning maintaining a gap over the single session models over the entire sampling space.
Results on a new animal performing the same tasks. Next we tested our model on 12 sessions from a completely new animal (6 CO sessions, 6 RT sessions from Monkey T) that was not included in the training of the model (Table 2, Right). When we applied our unit identification approach on the multi-session model (P0Y0-mp + Unit ID), we see a bit of a dip in overall accuracy from the single-session model (P0Y0-[Single-session]); however, when we fine-tune the weights of the model further, we achieve an accuracy of 0.9379 on CO, which is a significant improvement over all of the other baselines. For the RT task, the single session model is 0.7569 and with unit identification we achieve 0.7669 and with full fine-tuning we get up to 0.7916 (Figure 4B). These results are promising and show that we can use our pretrained model on new animals with only a few minutes of labeled data.
Performance on the Neural Latents Benchmark (NLB). To understand how well our pre-trained model performs on data collected from new animals performing novel tasks with different equipment (example: touch screen vs. manipulandum), we applied our pretrained model to the MC-Maze (Monkey L) and MC-RTT (Monkey I) datasets from the NLB (Table 3)[27]. The NLB serves as a benchmark for neural representation learning and decoding and thus we can include other single-session baselines to our comparisons, including self-supervised models AutoLFADS [31] and NDT [32] which produce denoised firing rate estimates over which we fit a linear layer (+ Linear), a supervised variant of NDT
[33] (NDT-Sup), and EIT [33]. Both datasets contain single sessions from new animals performing movement tasks that we haven't seen during training.
Method	NLB-Maze	NLB-RTT
From scratch	Wiener Filter	0.7485	0.5438
GRU	0.8887	0.5951
MLP	0.8794	0.6953
AutoLFADS + Linear [31]	0.9062	0.5931
NDT + Linear [32]	0.8929	0.5895
NDT-Sup [33]	0.8708	0.4621
EIT [33]	0.8791	0.4691
POYO-[Single-session]	0.9470	0.6850
Pretrained	POYO-mp + Unit ID	0.8962	0.7107
POYO-mp + Finetune	0.9466	0.7318
POYO-1 + Unit ID	0.9329	0.7294
POYO-1 + Finetune	0.9482	0.7378
Table 3: Behavioral decoding results for datasets from the Neural Latents Benchmark [27]. Best performing model is in bold and second best model is underlined.
On the NLB-Maze dataset, we obtain a R2 of 0.8952 after unit identification, which is competitive with the baselines. These results are surprising since we do not modify the model's weights, and yet our pretrained model yields competitive results on a dataset collected under very different conditions. When we finetune the model, we boost the performance even further establishing a  $4.4\%$  gap over the best baseline. Similar trends can be observed for the RTT task (Monkey I), with even larger  $(2\%)$  improvement after finetuning.
3.5 POYO-1: A multi-lab, multi-task model for neural decoding
Given the impressive transfer results of POYO-mp to datasets from different labs, we ask whether we can use our approach to build a model that spans even more diverse recording setups that we expect
A.
Random target
Controller: Touch screen
C.i
B
Maze
Controller: Touch screen
Figure 5: Scaling up to more tasks and diverse neural and behavioral recording conditions. (A) Random target task from Makin et al. included in training of P0Y0-1 and (B) Maze task from NLB heldout for transfer testing. In (C) decoding accuracy for the single-session (gray) and the P0Y0-1 model (blue). (D) PCA projection of the learned session embeddings.
C.ii
D Learned Session Embeddings in POYO-1
to encounter when trying to unify data from many sources. In total, we used datasets from seven nonhuman primates spanning three different labs, with a total of 27,373 units and 16,473 neuron-hours for training our model. We call this pretrained multi-lab, multi-task model POYO-1.
Even in light of the high amounts of variability across these different datasets, POYO-1 provides consistent improvements over the single-session models (Figure 5C). When tested on a number of different transfer tasks (Tables 2, 3), we again find that the unit identification and finetuning approaches provide effective strategies for adapting our pretrained model to new datasets. We obtain notable performance on the NLB-Maze, where we find that we obtain a  $R^2$  of 0.9329 with only unit identification remapping, an almost  $4\%$  improvement over the unit identification result for our POYO-mp pretrained model, suggesting that we cover more of the task space with POYO-1.
When comparing the POYO-1 model with POYO-mp model, it is clear that both methods have their strengths. POYO-mp excels on the datasets sourced from the same lab, with 0.8788 on RT sessions compared to POYO-1's 0.8664. On the other hand, both models exhibit great transfer capabilities, with POYO-1 having the edge especially when using unit identification, indicating its ability to generalize better across diverse experimental conditions. This flexibility and scalability make these methods promising tools for future research in analyzing neural data from diverse sources.
Visualizing the session embedding space. We visualized the learned task embeddings to see if the model learns relationships between sessions seen during training (Figure 5D), even though it was not explicitly trained with this information. This analysis revealed clusters corresponding to different data sources and tasks, suggesting that our model has not only learned to identify patterns within each session but also recognized overarching structures across sessions. In particular, we find that the datasets collected in the Miller Lab (C,M,J) used to train the POYO-mp model are mapped into similar regions of the session latent space, and (I and L) sessions are mapped to their own distinct clusters.
Transformer architectures for time-series data. Transformers have emerged as a powerful model for processing time-series data due to their ability to capture long-range dependencies [34, 35]. The key to their success lies in the self-attention mechanism [14], which allows the model to weigh the importance of each time step in the input sequence when producing an output. A typical approach in applying transformers to time-series data involves discretizing the continuous time domain into discrete bins [35], and treating the binned data as a sequence of tokens. Each bin is then linearly projected into a high-dimensional embedding, which is input to the transformer model. The Perceiver framework [18, 13] moves away from the idea of patches, by directly processing the input bytes,
using cross-attention layers to reduce the complexity of attention. In addition to regular timeseries, there has been interest in applying transformers to model irregular timeseries including event stream and point process data [36, 37, 38]. However, many of these models work on univariate event streams and when they extend to multivariate cases, they assume fixed and known input channels.
Transformers applied to neural data. With the recent advances in transformers for sequence modeling in many different time-series, transformers have also recently found successful applications in building representations of neural population activity [32, 33, 39]. In these models, the spiking activity is first binned and then the estimated bin counts are tokenized. In the neural data transformer (NDT) model [32], the firing rates of all neurons in the population are embedded jointly in one token (time step or bin). In the embedded interaction transformer (EIT) [40], neurons are considered independently in one stage of processing and at a population level in a second stage, and thus the whole dataset is tokenized over both neurons and time. In the spatiotemporal (STNDT) model [39], two different tokenizations are also considered, one in space and one in time, and two representations are learned jointly for both tokenizations. In all cases, binned data are used and the models are trained on a single session and fixed set of neurons.
Multi-session training and alignment. The idea of decoding across multiple sessions has been explored in previous work [41, 42, 43, 44]. In many of these works, an initial baseline representation is formed on one day and alignment-based approaches are used to transfer a model trained on one session across recording days [23, 45, 12, 46, 47]. A subset of these methods [41, 43] can be trained on many sessions jointly, but rely on the assumption of shared dynamics or structure of a single task to achieve alignment. To the best of our knowledge, our work is the first to demonstrate multi-session transfer across subjects performing different tasks, and the first to demonstrate scaling across different data and model sizes.
In this paper, we introduce a novel framework for training transformers on large multi-session, multi-task neural activity datasets. To tackle the challenges of training on such large heterogeneous sources, we introduce a novel spike-level tokenization scheme and architecture that enables the model to learn from populations with varying numbers of neurons. We show that training a single unified model on multiple recordings is possible, and find that it leads to improved decoding performance. Finally, we build two large pretrained models (POYO-1, POYO-mp) that can be efficiently fine-tuned on new datasets, and make them available as a resource to the community.
In contrast to models trained on a single dataset, the pretrained models that we have developed provide a potential way to compare and contrast datasets, and also understand common motifs of activity and dynamics that may be shared across different sessions, tasks, and individuals. Thus, it will be critical to develop tools to probe the patterns and motifs learned by such models and characterize the neural mechanisms underlying different tasks and computations. In particular, we look to understand how spike tokens are grouped across different latent tokens and how the dynamics of the population are modeled in this latent space. Additionally, our proposed unit embedding space allows us to map units into a high-dimensional space; thus understanding how unit projections are organized might help reveal the similarities between different neurons and the nature of their interactions. Similarly, we can analyze the session embeddings to glean insights into inter-session and across-animal differences.
Our work shows how pretraining on diverse data, including datasets from animals performing different tasks and across different laboratories, can all help to improve our ability to decode from novel and unseen neural datasets. Already, our results demonstrate the positive effect of scale for neural data analysis. However, to scale this approach further and integrate even more diverse brain regions and tasks, it will be critical to move toward a self-supervised objective. Thus, our current architecture and multi-session framework could be also extended to self-supervised tasks like generative next-event prediction or masked modeling to allow for even larger datasets to be ingested.
This framework has the potential to advance neuroscience research in several ways. By enabling the development of large pretrained models that can be fine-tuned for various downstream tasks, our work can accelerate progress in brain-machine interfaces and disease modeling applications. The ability to quickly and effectively decode neural activity from diverse datasets and modalities can have significant implications for improving our understanding of the brain and developing new therapeutic interventions.
We would like to thank Patrick Mineault for providing a lot of great feedback on the work, and Jingyun Xiao for helping to run baselines for NLB. This project was supported by NIH award 1R01EB029852-01, NSF award IIS-2146072 as well as generous gifts from the Alfred Sloan Foundation (ED), the McKnight Foundation (ED), and the CIFAR Azrieli Global Scholars Program (ED). BAR acknowledges support from NSERC (Discovery Grant: RGPIN-2020-05105; Discovery Accelerator Supplement: RGPAS-2020-00031; Arthur B. McDonald Fellowship: 566355-2022) and CIFAR (Canada AI Chair; Learning in Machine and Brains Fellowship). MGP acknowledges support the Fonds de recherche du Quebec Santé (Grant: Chereurs-boursiers en intelligence artificielle). GL acknowledges support from the Canada CIFAR AI Chair program, the Canada Research Chair for Neural Computation and Interfacing, and NSERC Discovery grant RGPIN-2018-04821.
[1] A. Radford, K. Narasimhan, T. Salimans, I. Sutskever, et al., “Improving language understanding by generative pre-training,” 2018. 1
[2] A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. Chung, C. Sutton, S. Gehrmann, et al., “Palm: Scaling language modeling with pathways,” arXiv preprint arXiv:2204.02311, 2022. 1
[3] M. Dehghani, J. Djolonga, B. Mustafa, P. Padlewski, J. Heek, J. Gilmer, A. Steiner, M. Caron, R. Geirhos, I. Alabdulmohsin, et al., "Scaling vision transformers to 22 billion parameters," arXiv preprint arXiv:2302.05442, 2023. 1, 2
[4] A. Radford, J. W. Kim, T. Xu, G. Brockman, C. McLeavey, and I. Sutskever, "Robust speech recognition via large-scale weak supervision," arXiv preprint arXiv:2212.04356, 2022. 1
[5] R. Bommasani, D. A. Hudson, E. Adeli, R. Altman, S. Arora, S. von Arx, M. S. Bernstein, J. Bohg, A. Bosselut, E. Brunskill, et al., “On the opportunities and risks of foundation models,” arXiv preprint arXiv:2108.07258, 2021. 1
[6] J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei, "Scaling laws for neural language models," arXiv preprint arXiv:2001.08361, 2020. 1
[7] A. E. Urai, B. Doiron, A. M. Leifer, and A. K. Churchland, "Large-scale neural recordings call for new insights to link brain and behavior," Nature neuroscience, vol. 25, no. 1, pp. 11-19, 2022. 1
[8] C. Hurwitz, N. Kudryashova, A. Onken, and M. H. Hennig, "Building population models for large-scale neural recordings: Opportunities and pitfalls," Current Opinion in Neurobiology, vol. 70, pp. 64-73, 2021. Computational Neuroscience. 1
[9] M. Dabagia, K. P. Kording, and E. L. Dyer, "Aligning latent representations of neural activity," Nature Biomedical Engineering, pp. 1-7, 2022. 1, 2
[10] C. Pandarinath, K. C. Ames, A. A. Russo, A. Farshchian, L. E. Miller, E. L. Dyer, and J. C. Kao, “Latent factors and dynamics in motor cortex and their application to brain-machine interfaces,” Journal of Neuroscience, vol. 38, no. 44, pp. 9390–9401, 2018. 1, 5, 16
[11] W. M. Grill, S. E. Norman, and R. V. Bellamkonda, "Implanted neural interfaces: biochallenges and engineered solutions," Annual review of biomedical engineering, vol. 11, pp. 1-24, 2009. 2
[12] J. Jude, M. Perich, L. Miller, and M. Hennig, "Robust alignment of cross-session recordings of neural population activity by behaviour via unsupervised domain adaptation," in Proceedings of the 39th International Conference on Machine Learning (K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari, G. Niu, and S. Sabato, eds.), vol. 162 of Proceedings of Machine Learning Research, pp. 10462-10475, PMLR, 17-23 Jul 2022. 2, 10
[13] A. Jaegle, S. Borgeaud, J.-B. Alayrac, C. Doersch, C. Ionescu, D. Ding, S. Koppula, D. Zoran, A. Brock, E. Shelhamer, et al., "Perceiver io: A general architecture for structured inputs & outputs," arXiv preprint arXiv:2107.14795, 2021. 2, 9, 15
[14] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin, "Attention is all you need," Advances in neural information processing systems, vol. 30, 2017. 2, 9
A Additional Model Details
A.1 Rotary position encoding
Rotary Position Embeddings (RoPE) are used to incorporate relative positional information into transformer models in a memory and compute efficient manner[19]. Unlike absolute position embeddings which are added or appended to the input embeddings, these are directly incorporated in the attention mechanism.
First, we define a  $2 \times 2$  rotation matrix  $\mathbf{R}_{2 \times 2}(t,T)$  for a given timestamp  $t$  and time-period  $T$ :
$$
\mathbf {R} _ {2 \times 2} (t, T) = \left[ \begin{array}{c c} \cos (2 \pi t / T) & - \sin (2 \pi t / T) \\ \sin (2 \pi t / T) & \cos (2 \pi t / T) \end{array} \right]
$$
$\mathbf{R}_{2\times 2}(t,T)$  provides a time/positional encoding with a time period  $T$ . Intuitively, the sinusoidal values in this rotation matrix can help the model learn and distinguish variations in timestamps at a timescale  $T$ .
Given the head dimension of the attention block  $D$ , we create a  $D \times D$  rotary embedding matrix  $\mathbf{R}(t)$  for a timestamp  $t$  by constructing a block-diagonal matrix of rotation matrices  $\mathbf{R}_{2 \times 2}$ , where the time-periods are varied logarithmically between  $T_{min}$  and  $T_{max}$
$$
\mathbf {R} (t) = \left[ \begin{array}{c c c c} \mathbf {R} _ {2 \times 2} (t, T _ {0}) & 0 & \dots & 0 \\ 0 & \mathbf {R} _ {2 \times 2} (t, T _ {1}) & \dots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \dots & \mathbf {R} _ {2 \times 2} (t, T _ {\frac {D}{2} - 1}) \end{array} \right]
$$
Given  $D$ ,  $T_{min}$ , and  $T_{max}$ , we generate  $T_{i}$  by the following equation:
$$
T _ {i} = T _ {m i n} \left(\frac {T _ {m a x}}{T _ {m i n}}\right) ^ {2 i / D}
$$
In our models,  $T_{\text{min}}$  and  $T_{\text{max}}$  were set to 1ms and 4s respectively.
This rotation matrix is incorporated into the attention mechanism by replacing the attention score calculation with the following equation:
$$
\mathbf {a} _ {i j} = \mathrm {s o f t m a x} \big (\left(\mathbf {R} (t _ {i}) \mathbf {q} _ {i}\right) ^ {T} \left(\mathbf {R} (t _ {j}) \mathbf {k} _ {j}\right) \big)
$$
Since  $\mathbf{R}(t_i)^T\mathbf{R}(t_j) = \mathbf{R}(t_j - t_i)$ , the attention score  $\mathbf{a}_{ij}$  depends only on the relative timing between the query and key tokens. This is what makes RoPE relative in nature.
What is explained above has been introduced in [19], and use successfully in many models [18]. In our experiments, we found an interesting extension of RoPE. In addition to incorporating RoPE in calculating the attention scores, we also incorporate it in the weighted-summation of values inside the attention mechanism. We do this to encode the relative timing between values and queries in the output of the attention block. A direct, but computationally inefficient, way of doing so would be:
$$
\mathbf {z} _ {j} \gets \mathbf {z} _ {j} + \sum_ {i} \mathbf {a} _ {i j} \mathbf {R} (t _ {i} - t _ {j}) \mathbf {v} _ {i}
$$
This is inefficient since it requires us to explicitly compute  $\mathbf{R}(t_i - t_j)$  for all  $N^2$  values of  $(t_i - t_j)$ . We use the fact that  $\mathbf{R}(t_i - t_j)$  can be factorized as  $\mathbf{R}(-t_j)\mathbf{R}(t_i)$  to reach the following, much more efficient, implementation:
$$
\mathbf {z} _ {j} \gets \mathbf {z} _ {j} + \mathbf {R} (- t _ {j}) \sum_ {i} \mathbf {a} _ {i j} \mathbf {R} (t _ {i}) \mathbf {v} _ {i}
$$
That is, we first apply rotation matrices  $\mathbf{R}(t_i)$  to all values  $\mathbf{v}_i$  before passing them to as inputs to the attention block. Then we apply the rotation matrices  $\mathbf{R}(-t_j)$  to its outputs.
This value-rotation mechanism captures relative-timing information more explicitly in the output of the attention block, allowing the feedforward network to take advantage of timing information. In the absence of this mechanism, the only way timing information is added to the latent tokens is through the attention scores. We observed in early experiments that our model performed better with value-rotation enabled for the input cross-attention and all self-attention layers.