SlowGuess/NeurIPS_2023_OCR huggingface dataset:
As discovered in the Perceiver framework [13], we observed an improvement in performance upon rotating only half of the head dimensions. That is, half of the  $\mathbf{R}_{2\times 2}$  matrices in  $\mathbf{R}$  were replaced by  $2\times 2$  identity matrices.
A.2 Time-based context window
We consider an arbitrary window of time  $[t, t + T]$ , meaning that we do not use the trial structure when training our model. In our experiments, we set  $T$  equal to 1 second. At inference time, we use a sliding window of step  $500\mathrm{ms}$  and average the predicted output in overlapping segments.
In language models, and more generally transformer-based models, the context window is defined with respect to the number of tokens. The neural data modality is a bit more complex, as the number of spikes scales with the number of units: Across multiple datasets, we record a variable number of units, ranging from tens of units to hundreds of units. Defining the context window with respect to number of tokens will result in a variable time duration, long for small populations and short for larger ones. Our context window is defined with respect to time (fixed to 1s), each sequence will have a variable number of tokens  $(M)$ . When creating batches, we pad to the largest sequence in the batch, and appropriately use an attention mask in the first cross-attention layer. GPU hardware is optimized to process fixed-size tensors, to make computation more efficient, we introduce a distributed training load balancing mechanism, which deals with variable-length input sequences (Appendix A.4.4).
A.3 Further details on tokenization
Delimiters. When extracting the spike data in a window of time  $[t, t + T]$ , we include delimiters for each unit, indicating to the model that we are tuning into the activity of that unit between times  $t$  and  $t + T$ . This choice is important because: 1) we are using relative position encoding, so the beginning or the end of the window is not clear, and 2) our tokenization only captures activity and not inactivity. If we consider a unit that is inhibited and does not fire in that window of time, the model would be unaware of the presence of that unit, unless we add [START] and [END] tokens. We learn two embedding vectors for these delimiters  $\mathbf{x}_{start}$  and  $\mathbf{x}_{end}$ . For each unit in the population, we add these two tokens with embedding equal to the sum of the delimiter embedding and the unit embedding, and assign each token timestamps  $t$  and  $t + T$  respectively.
Latent tokens. Recall that we divide the latent tokens into groups of  $n$  and spread each group uniformly over the context window. Specifically, we use a total of 256 latent tokens and divide them into groups of 8, this means that there will be 32 tokens that share the same timestamp. We experiment with weight-sharing, where we set the learned latent tokens in the same group to have the same weights  $(\mathbf{e}_0,\dots ,\mathbf{e}_{31}$  in Fig 1). This means that tokens in the same group (size 8) will have the same embedding  $\mathbf{z}_{0,i} = \mathbf{e}_{(i\mod 32)}$ , but different timestamps  $(t_0^l,\ldots ,t_7^l$  in Fig 1). We believe this is a reasonable inductive-bias to architecturally enforce since, intuitively, we want all sets of latent tokens that share a timestamp to query the spike tokens in the exact same manner, just with a different timestamps. These tokens will evolve independently once the information is pulled from the input space (after the first cross-attention).
A.4 Training details
The model is trained using the LAMB optimizer [28] with weight decay. The learning rate is held constant, then decayed towards the end of training (last  $25\%$  of epochs), using a cosine decay schedule. Single-session models are trained on a single GPU with a batch size of 128 while large models are trained with 8 GPUs with a total batch size of 1400. Note that we didn't see any benefits in increasing the batch size when training single-session models.
A.4.1 Compute
The large models are trained on a machine with an AMD EPYC 7452 32-Core Processor and 8 Nvidia A40 GPUs (48Gb memory), POY0-mp was trained for 2 days, and POY0-1 was trained for 5 days (both for a total of 400 epochs). Single-session models are trained with a single Nvidia GeForce RTX 3090 GPU, and take less than an hour to train. Finetuning models is also done with a single GPU. Unit identification converges very quickly and takes less than a minute on a single GPU or a few minutes on the CPU.
A.4.2 Data augmentation
Previous work [48, 49] explores the use of augmentations when training neural population models. In our experiments, we use the unit dropout augmentation which randomly samples a subset of the population. We set a minimum population size of 30 for this augmentation to ensure that it is not too destructive (note that our model is still trained with any number of units, this limit is only imposed for the augmentation). We did not explore the use of other augmentations, but believe it could be a promising direction to further improve the capabilities of the model.
A.4.3 Loss
We train our model using a mean-squared error loss over the hand velocity sequence. We normalize the velocity data using z-scoring at the lab level, this means that all velocity data in datasets from the same lab is rescaled using the same scalar.
Dealing with variable sampling rate. When training on mixtures of datasets from multiple labs, any given batch will contain behavior sequences that have different sampling rates. In a window of 1s, some sequence will have more timepoints over which we will evaluate the MSE compared to other sequences. To avoid the over-representation of samples that simply have a high sampling rate, we evaluate the error on a randomly sampled 100 timepoints (these points are randomly sampled at each pass).
Weight in training mix. While behavior during random target tasks can be more complex and noisy, behavior during center-out reaching tasks is more structured, the monkey usually gets a preparation phase, where the movement can be planned [20]. The latter task has been studied extensively for that reason. Knowing that neural activity is very salient during center-out reaching, we increase the weight of the prediction loss during the reaching segments by a factor of 5. The idea of over-weighing samples of "good quality" is not novel and is commonly used when training language models [15].
A.4.4 Load Balancing
We distribute our training over multiple GPUs. Note that, in our experiments, the input sequence length  $(M)$  can vary anywhere between 1k and 20k tokens. Because of the variability in the length of the input sequences, we distribute sequences that have close length to same GPU node.
We now provide further details on our load balancing strategy:
- Each GPU node has a local batch size. A GPU that is assigned large sequences will have a smaller local batch size than one that is assigned shorter sequences.
- We create a number of data buckets equal to the number of GPUs. Each bucket  $i$  is assigned sequences that are of length  $M \in [M_{min}^{i}, M_{max}^{i}]$ , and a local batch size  $B_{i}$ .
- We select these parameters in such a way that we effectively utilize all available compute capacity, and minimize the amount of padding used.
A.5 Evaluation Details
During training, we train on any 1s segment of the recording, but at evaluation we report the decoding performance using the same evaluation strategy used in previous work [27, 10]: for center-out reaching datasets, we report the decoding score during the reaching movement (only when the trial is completed successfully), for random target datasets that are trialized (hold period followed by 3/4 random reaches), we report the decoding score during the reaching as well, and finally for all other continuous random target datasets, we report the decoding score for all segments. Note that none of the segments used for testing are seen during training.
Midline
Churchland, et al.
Monkey Je (Left M1)
Makin et al. Monkey I (Right M1 & S1)
Flint, et al. Monkey Mi (Right M1)
Figure 6: Diversity of neural recordings used to train P0Y0. We visualized the array locations for all of the datasets used in our current study. P0Y0-mp is trained on recordings from PMD and M1 (CL, CR, M, Ja) and P0Y0-1 is trained on recordings from (CL, CR, M, Ja, I, L, Mi) spanning PMD, M1, and S1 across both hemispheres and in seven animals.
A.6 Fine-tuning Details
When finetuning the model, we perform gradual unfreezing of the weights. For the first few steps, we optimize the unit embeddings and the session embedding only (unit identification) and then unfreeze the rest of the weights. We use the LAMB optimizer with batch size 128, learning rate  $10^{-4}$  and weight decay  $10^{-4}$ . We use the exact same hyper parameters and finetuning process for all experiments we report. This strongly suggests that our method is general, and easy to adapt to new data, without the need for expert neural network training knowledge, or expensive hyperparameter search.
To train and evaluate our framework, we curated a large number of open and publicly available datasets that represent a range of neural and behavioral recording platforms.
B.1 Datasets used for training
Datasets used to train POYO-mp. To build our first large-scale model (POYO-mp), we acquired 100 sessions of unique recordings from three non-human primates performing two different movement tasks that were previously studied in four publications [21, 22, 20, 23]. We make these datasets publicly available on the DANDI open repository [29]. In all of these datasets, a nonhuman primate is seated in a primate chair and executes movements using a custom 2-D planar manipulandum to control a cursor on a computer screen. They performed one of two tasks within a session.
- Center-out Task (CO): They performed a standard center-out reaching task involving eight outer targets arranged around a circle with a radius of  $8\mathrm{cm}$ . The monkeys were required to move to the center target, wait for a variable hold period, and then move to one of the outer targets upon receiving an auditory go cue. The task aimed to study neural activity during movement planning, preparation, and execution.
- Random Target Task (RT): The setup is similar, the targets are not arranged around a circle but are randomly placed. The monkeys were required to move between 4 targets upon receiving an auditory go cue.
After extensive training, the monkeys were implanted with chronic multi-electrode arrays in the primary motor cortex (M1) and the dorsal premotor cortex (PMd). Neural activity was recorded using
a Blackrock Cerebus system, and the data was manually processed offline to identify single neurons. To ensure only independent channels were included, steps were taken to identify and disable potential candidates with high crosstalk and to exclude cells with a high percentage of coincident spikes.
Datasets used to train the multi-lab, multi-session model (POYO-1). When training POYO-1, we acquire more data as shown in Table 1, adding all of the datasets outlined from Churchland et al., [24], Makin et al., [25], and Flint et al., [26]. We still hold out the 12 sessions from Monkey T, 2 sessions from Monkey C, and the two Neural Latents Benchmark (NLB) [27] datasets from training.
As depicted in Figure 5, instead of using a manipulandum like in the previous experiments, these new datasets use a touch screen and have different sampling rates for their behavioral outputs. In addition, in the Churchland, et al. datasets, they use threshold crossing based processing rather than applying spike sorting to identify single units. We refer the reader to the listed publications for more details about how these datasets were collected.
We would like to reiterate that there are major differences in the way the datasets were collected. Yet, in spite of this variability, it is important that our method works without having to standardize the datasets across sites further and apply specialized techniques to one dataset but not other. Thus we did not attempt to homogenize or standardize the data. The only processing that we perform is on the behavior to clip outliers in the behavior based upon times of high acceleration. We use a simple threshold heuristic to avoid training on periods with extremely high acceleration.
1. We do not filter units based on a presence ratio, or reject multi-units.
2. We do not re-apply the same spike sorting algorithm on all datasets, and use the data as is. In [24], the spiking events obtained through threshold crossings (i.e. all units are multi-units), and in [25], we have both threshold crossing units and spike sorted isolated single units, and we use both. Each algorithm used in each lab was tuned differently, but we consider this a good example of diversity that a general model needs to deal with.
3. We do not resample or process the velocity timeseries.
Minimizing the amount of processing needed to integrate new datasets into our model is key for more easily scaling to more datasets, and also providing an accessible model that the community at large can use out of the box, without having to adapt to a specific standard.
B.2 Datasets for evaluation and fine-tuning
In addition to holding out  $20\%$  of the data from each session, we hold out all sessions from Monkey T (6 CO, 6 RT). We also use two datasets from the Neural Latents Benchmark [27], MC-Maze and RTT. The Maze dataset consists of recordings from primary motor and dorsal premotor cortices of a monkey performing reaches with instructed delays to visually presented targets while avoiding the boundaries of a virtual maze. This dataset offers a variety of behavioral configurations, with each configuration using a different combination of target position, number of virtual barriers, and barrier positions, resulting in a variety of straight and curved reach trajectories. With thousands of trials, the Maze dataset allows for a rich investigation into the structure of population activity, while its instructed delay paradigm enables a clear separation of neural processes related to preparation and execution. On the other hand, the RTT dataset introduces different modeling challenges, as it contains continuous, point-to-point reaches that start and end in various locations, have highly variable lengths, and few repetitions. We report the performance of our models on the same segments of data used to evaluate models in the NLB benchmark.
C.1 Comparison with single-session baselines
In our comparisons in Table 2, we compared POYO against existing baseline models that are commonly used for neural decoding [30]. In particular, we applied a Wiener Filter (WF), Feed Forward Network (MLP), and a Gated Recurrent Unit (GRU). These models predict the behavior at each timestamp  $t$  given a history window of neural activity  $[t - T, t]$ . For the NLB datasets, we also compare with
Auto-LFADS [31], an unsupervised approach that can be used to obtain single trial smoothed rate estimates for a population of neurons before applying a linear layer on top of the inferred rates to obtain an estimate of the behavior. While this model isn't designed for decoding per se, it is frequently used for denoising for BMI decoding applications.
For the NLB, we use a 5 ms binning of the neural activity for training models and for evaluation of the behavioral output as this is the finest resolution used in the benchmark. For the rest of the MP datasets, we use a 10 ms binning rate for the neural data and behavior as this corresponds to the behavior sampling rate for these datasets.
The hyperparameters we used for training our single-session POYO model can be found in Table 4.
Hyperparameter	Value
Embedding Dimension	128
Head Dimension	64
Number of Latents	128
Depth	6
Number of Heads	8
FFN Dropout	0.3
Linear Dropout	0.3
Attention Dropout	0.3
Weight Decay	1e-4
Learning Rate	4e-3
Batch Size	128
C.2 Robustness of Unit-identification
Unit-identification, introduced in Section 2.5, can be seen as an approach for identifying units from a new dataset. To evaluate the robustness of this process, we ran unit-identification on a session that was already seen during pre-training. We test whether finetuning on this "new" set of units, will map them close to their true embedding (found during pre-training). We use our largest model, POYO-1, start with randomly initialized unit embeddings and run unit-identification.
We perform this experiment on two different animals. In Figure 7, we show the evolution of unit-embeddings during training. To compare the newly calibrated set of units with its pre-trained version, we normalize the unit-embeddings and report the cosine similarity averaged over the set of units. Note that normalization is done because POYO's first cross-attention layer applies layer norm to these embeddings. We also report the "unit identification accuracy" which we define as the ratio of tuned unit embeddings that have their true unit embedding as their nearest neighbor. Results for these experiments are present in Table 5.
Table 4: Hyperparameters used for training all POYO single-session models
Dataset	Num of Units	Unit-Id Accuracy	Unit-Id Cosine-Similarity
Monkey C, CO, 2013/10/03	73	1.000	0.845
Monkey M, CO, 2014/02/03	116	0.966	0.802
Table 5: Unit re-identification results. Cosine-similarity measure is averaged over all units in a dataset. Accuracy is measured for 1-nearest-neighbor classifier.
Overall, this analysis suggests that the unit-identification approach is robust and reliable for identifying units. We believe that understanding the functional similarities of units that are mapped close to each other in the unit embedding space will be an interesting avenue for future work.
C.3 Hyperparameter sensitivity
When initially training the single-session models, we selected a subset of datasets (out of the 100 sessions) and performed a random hyperparameter search and selected the set of hyperparameters with the best performance on the validation set. These experiments quickly revealed that the model was very stable to a wide range of hyperparameters. In Figure 8, we report some examples of this for
Monkey M, CO 2014/02/03
Figure 7: Evolution of the unit embeddings with unit-identification on sessions from two different animals. From left to right, the finetuned embeddings (orange) are visualized at 10, 100, and 400 epochs, relative to their true embedding in the pre-trained model (blue).
Monkey C, CO 2013/10/03
two sessions. We perform 300 runs for each session, then report the average performance for each hyperparameter.
Given these findings, we selected the best set of hyperparameters for the subset and used the same parameters when training all of the 100 single-session models that we study in Figure 2.
Figure 8: Single-session model performance for a wide range of values for different hyperparameters. The  $R^2$  score is reported on the validation sets for both datasets.