Neural Data Transformer 2: Multi-context Pretraining for Neural Spiking Activity
Authors: Joel Ye
Jennifer L Collinger
Leila Wehbe
Robert Gaunt
Venue: NeurIPS 2023 poster (2023)
Keywords: Pretraining
Scaling Laws
Neuroscience
Brain-computer interfaces
TL;DR: Unsupervised pretraining works for transformers on motor BCI cortical spiking activity.

Abstract: The neural population spiking activity recorded by intracortical brain-computer interfaces (iBCIs) contain rich structure. Current models of such spiking activity are largely prepared for individual experimental contexts, restricting data volume to that collectable within a single session and limiting the effectiveness of deep neural networks (DNNs). The purported challenge in aggregating neural spiking data is the pervasiveness of context-dependent shifts in the neural data distributions. However, large scale unsupervised pretraining by nature spans heterogeneous data, and has proven to be a fundamental recipe for successful representation learning across deep learning. We thus develop Neural Data Transformer 2 (NDT2), a spatiotemporal Transformer for neural spiking activity, and demonstrate that pretraining can leverage motor BCI datasets that span sessions, subjects, and experimental tasks. NDT2 enables rapid adaptation to novel contexts in downstream decoding tasks and opens the path to deployment of pretrained DNNs for iBCI control. Code: https://github.com/joel99/context_general_bci

Neural Data Transformer 2: Multi-context Pretraining for Neural Spiking Activity
Joel Yel.2.3, Jennifer L. 1. Leila Wehbe Robert Gaunt 1,3,4,5,6+
Rehab Neural Engineering Labs, University of Pittsburgh,
2 Neurosccience Institute, Carnegie Mellon University,
CCnter for the Neural Basis of Cognition, Pittshurgh,
Department of Physical Medicine and Rehabilitation, University of Pittsburgh, DDeparteent of Bioengineering, University of Pittsburgh,
6 Department of Biomedical Engineering, Carnegie Mellon University, Machine Learning Department, Carnegie Mellon University
Abstract
The neural population spiking activity recorded by intracortical brain-computer interfaces (iBCIs) contain rich structure. Current models of such spiking activity are largely prepared for individual experimental contexts, restricting data volume that collectable within single session and limiting the effectiveness of deep neural networks (DNNs). The purported challenge in aggregating neural spiking data is the pervasiveness of context-dependent shifts in the neural data distributions. However, large scale unsupervised pretraining by nature spans heterogeneous a fundamental recipe for successful representation
data, and has proven learning across deep learning. We thus develop Neural Data Transformer 2 (NDT2), a spatiotemporal Transformer for neural spiking activity, and demonstrate that pretraining can leverage motor BCI datasets that span sessions, subjects, and experimental tasks. NDT2 enables rapid adaptation to novel contexts downstream decoding tasks and opens the path to deployment of pretrained DNNs for iBCI control. Code: http://ww..ceeooo////////eeeeeee
Introduction
Intracortical neural spiking activity contains rich statistical structure reflecting the processing subserves. For example, motor cortical activity during reaching is characterized with dineennnli dynamical models [ [2), and these models can predict behavior under external perturbation and provides interpretive lens for motor learning [3 However, models are currently prepared for each experimental context, meaning separate datasets are collected for each cortical phenomena in each subject, for each session. Meanwhile, spiking activity structure is at least somewhat stable across these contexts; for example, dominant principal components (PCs) of neural activity can remain stable across sessions, subjects, and behavioral tasks (6.9. This structure persists in spite of turnover in recorded neurons, physiological changes in the subject, or task changes required by the experiment [TO [.]. Conserved neural population structure suggests the opportunity for models that span beyond single experimental contexts, enabling more efficient, potent analysis and application.
In this work focus computer interfaces (iBCIs). With electrical recordings of just dozens to hundreds channels of neuronal population spiking activity, today's is iBCIs relate this observed neural activity to behavioral intent, achieving impressive milestones such as high-speed speech decoding [12 and
primary case: neuroprosthetics powered by intracortical brain
"Correspondence rag53epitt.co edu
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

high freedeo control of robotic arms [3]. Even so, these iBCIs currently require arduous supervised calibration in which neural activity on that day is mapped behavioral intent. At best. cutting-edge decoders have included training data from across several days, producing thousands of trials, which still modest. by deep learning standards [2]. Single-session models still dominate the Neural Latents Benchmarks (NLB), primary representation learning benchmark for spiking activity [4]. Thus, despite the scientifically observed conserved manifold structure, there has been little adoption of neural population models that can productively aggregate data from broader contexts.
One possible path forward deep learning's seemingly robust recipe for leveraging heterogeneous data across domains: a generic model backbone (e.g. Transformer [[51), unsupervised pretraining over broad data, and lightweight adaptation for target context (e.g. through fine-tuning). The iBCI community has set the stage for this effort, for example with iBCI dataset releases (Section A.1 and Neural Data Transformer (NDT) [66, which shows Transformers, when prepared with masked autoencoding. model single-session spiking activity well. We hereafter refer to NDT as NDT1. Building on this momentum, we report that Transformer pretraining can apply to motor cortical neural spiking activity from iBCIs, and allows productive aggregation of data across contexts.
Here, we contribute NDT2, a Transformer that pretrains over broad data sources of motor cortical spiking activity. NDT2 modifies NDTI to improve scaling across heterogeneous contexts in three ways (Section 3) spatiotemporal attention, learned context embeddings, and asymmetric encode-decods [17]. We find positive transfer with data from different data sessions, subjects, and tasks, and quantify their relative value. Once pretrained, NDT2 can be rapidly tuned in novel experimental sessions. We focus offline evaluation on motor applications, demonstrating NDT2's value in decoding unstructured monkey reaching and human iBCI cursor intent. We also show cooof-of-principle real-time cursor control using NDT2 in a human with an iBCI.
Contributions:
2 Related Work
Unsupervised neural data pretraining. Unsupervised pretraining's 's broad applicability is useful in neuroscience; little is common across neural datasets except the neural data itself. Remarkably, pretraining approaches across neural data modalities are similar; of the sampled Table [ use masked autoencoding as pretraining objective (EEG uses contrastive lcarning), and 3 use a Transformer backbone (ECoG uses a CNN). Still, each modality poses different challenges for pretraining; for iBCIs that record spiking activity, the primary challenge data instability [8] The high spatial resolution of iBCI microelectrode arrays enables recording of individual neurons and provides the signal needed for high-peeformance rehabilitation applications, but this fine resolution also causes high sensitivity to shifts in recording conditions [1]] 18]. iBCIs typically require recalibration within hours, relative to ECoG-BCIs that may not require recalibration for days [19]. At the macroscopic end, EEG and fMRI can mostly address inter-measurement misalignment through preprocessing (e.g. registration to an atlas).
Table 1. Neural data pretraining, NDT2, like contemporary neural data models, are intermediate pretraining efforts, comparable early-modern large language models like BERT 20 Neural models vary greatly in task quality and data encoding: implanted devices (Microelectrodees, ECoG, SEEG) severely restrict the subject count available (especially with public data). Volume is estimated as full dataset size model input size. ECoG
Electrocorticooraphy SEEG = Stereo electroencephaloggrapyy EEG = Electroencephaloggaahy. fMR1
Functional magnetic resonance imaging.
Modality
Task
Estimated Pretraining Volume
Subjects
Microelectrodes (ours)
Motor reaching
0.25M trials
12
ECoG: LFP 21
Naturalistic behavior
0.04M trials/ / 108 days 22
12
SEEG:
LFP
23
Movie Viewing
3.2M trials / 4.5K electrode-hrs
10
fMRI
24
Varied (34 datasets)
1.8M trials (12K scans)
1.7K
EEG
25
Clinical assessment
0.5M trials / (26K runs 26
11K
BERT
20
Natural Language
IM 'trials' (3.3B tokens)
Data aggregation for iBCL. Data aggregation for iBCI has largely been limited to multi-session aggregation, where a session refers to an experimental period lasting up to several hours. These efforts often combine data through method called stitching [27]. For context, the extracellular spiking signals recorded microelectrode arrays are sometimes "sorted" by the shape of their electrical

---

A.
Binned Single-trial spiking activity
Context Embed
More data volume
Unmasked
Session Subject
Task
Encoder
More data relevance
Masked
Neural Decoder
Behavior Decoder
Poisson NLL
Figure 1. A. NDT2 a spatiotemporal Transformer encoder-decoder the style of He al. 17 operating on binned spiking activity. During masked autoencoding pretraining. spike rate decoder reconstructs masked spikes from encoder outputs; downstream, additional decoders similarly use encoder outputs for behavior prediction (e.g. cursor movement). The encoder and decoders both receive context embeddings as inputs. These embeddings are learned vectors for each unique type of metadata, such as session subject ID. B. NDT2 aims enable pretraining over diverse forms of related neural data. Neural data from other sessions within a single subject are the most relevant but limited volume, and may be complemented by broader sources of data.
waveforms, with different shapes attributed to spikes from distinet putative neurons. Sorted data has inconsistent dimensions across sessions, but as mentioned, activity across sessions has been observed to share consistent subspace structure, as e.g. identified by PCA (7]. Stitching aims to extract this stable subspace (and also resolve neuron count differences) by learning readin and readout layers per session. Stitching has been applied BCI and broader neuroscience applications over half However, even linear layers can add thousands of parameters, which risks overfitting in clinical iBCI data that comprise only dozens of trials.
year 30
31
Alternatiell, many iBCI systems simply forgo spike sorting, which appears to have a minor perfor- mance impact [TO
Then, multi-session data have consistent dimensions and can feed directly (even if the units recorded in those dimensions shift (34)). Note that these referenced models also typically incorporate augmentation strategies centered around channel gain modulation, noising, dropout, emphasizing robustness as a design goal.
[8
into a single model 10
32
33
Relative to existing efforts, NDT2 aims aggregate larger volume of data by considering other subjects and motor tasks (Fig. (B). The fundamental tension we empirically study is whether the increased volume of pretraining is helpful despite the decreased relevance of heterogeneous data for any particular target context.
Nomain-adaptive vs. domain-robust decoding. Given little training data, special approaches help enable BCI decoder use new, shifted contexts. For example, decoders can be designed to be robust to hypothesized variability in recorded populations by promoting invariant representations through model or objective design 33 35| 36). Alternatively, decoders can be adapted to novel context with further data collection. which is reasonable especially if only unsupervised neural data required. Several approaches align data from novel contexts by learning an input mapping that minimizes distributional distance between novel encodings and pretraining encodings explicitly [TT 37 38)
NDT2 also allows unsupervised adaptation, simply by fine-tuning without additional objectives.
3 Approach
3.1 Designing Transformers for unsupervised scaling neural data
Transformers prepared with masked autoencoding (MAE) 20 are competitive model for represen- tation learning on spiking neural activity in single contexts, as measured by their performance on the NLB 14]. Cross-domain Transformer investment has also produced vital infrastructure for scaling to large datasets. Thus, we retain Transformer MAE 17 as the unsupervised pretraining recipe. This architecture reviewed in Fig. IA.

---

The primary design choice is then data tokenization, i.c. how to decompose the full iBCI spatiotem- poral spiking activity into units that NDT2 will compute and predict over. Traditionally, the few hundred neurons in motor populations have been analyzed directly in terms of their population-level temporal dynamics []. NDTI [66] follows this heritage and directly embeds the full population, with one token per timestep. Yet across contexts, the meaning of individual neurons may change, so operations to learn spatial representations may provide benefits. For example, Le and Shlizerman [39] and Liu et al. [36] add spatial attention to NDT's s temporal attention. Yet separate space-time attention can impair performance [40] and requires padding in both space and time when training over heterogeneous data. Still, the other extreme of full neuron- wiise spatial attention has prohibitive computational cost. We compromise and group K neurons to a token (padding if needed), akin to pixel patches in ViTs [4I]. Neural "patches" are embedded by concatenating projections of the spike counts in the patch. In pilot experiments, we find comparable efficiencies between this patching strategy and factorized attention. We opt for full attention design to easily adopt the asymmetric encoder-decoder proposed in [17]- This architecture first encodes unmasked tokens and uses these encodings to then decode masked tokens. This 2-step approach provides memory savings over the canonical uniform encoder architecture, as e.g. popularized in the language encoder, BERT [20].
We next consider token resolutions. In time, iBCI applications benefit from control rates of 50- 100Hz 42
we adopt 50Hz (20ms bins) and consider temporal context up to 2.5s (i.e. 125 tokens in time). Given total context budget of 2K tokens (limited by GPU memory). this leaves a few dozen tokens for spatial processing. is currently common to record from 100-200 electrodes with an interclectrode spacing of 400μm (Blackrock Utah Arrays) 29 43), our budget forces 16-32 channels per token. We note that future devices are likely to record from thousands of channels at a time with much higher spatial densities [4). which will warrant new spatial strategies.
NDT2 also receives learned context tokens. These tokens use known metadata to allow cheap spe- cialization, analogous to prompt tuning [45] from language models or environment embeddings [46 from robotics. Specifically, we provide tokens reflecting session, subject, and task IDs.
3.2 Datasets
We pretrain models over an aggregation of datasets; the relative volume of different datasets are compared in Fig. 2A, with details in Section A.I All datasets contains single-unit (sorted) or multi- unit (unsorted) spiking activity recorded from either monkey or human primary motor cortex (M1) during motor tasks. particular, we focus evaluation on a publicly available monkey dataset, where the subjects performed self-paced reaching to random targets generated on 2D screen (Random Target Task, RTT) [47. and unpublished human clinical BCI datasets. RIT data contains both sorted and unsorted activity from two monkeys over 47 sessions (~20K seconds per monkey). RTT is ideal for assessing data scaling as features long, hourly sessions of a challenging task where decoding accuracy increases with data (88 For comparison, in the Maze NLB task, which uses cued preparation and movement periods, decoding performance saturates by 500 trials [14. Since RTT is continuous, we split each session into Is trials in keeping with NLB 14
We also study MI activity in 3 human participants with spinal cord injury (P2-P4). These participants have limited preserved motor function but can modulate neural activity in MI using attempted movements of their own limbs. This activity can decoded to control high degree-of-freete freedom behavior I we restrict our study to 2D cursor control tasks to be most analogous to RTT, which also restricts targets to a 2D workspace. All experiments conducted with humans were performed under an approved Investigational Device Exemption from the FDA and were approved by the Institutional Review Board at the University of Pittsburgh. The clinical trial is registered at clinicaltrials. (ID: NCTO1 894802)) and informed consent was obtained before any experimental procedures were conducted. Details on the implants and clinical trial are described in 13, and similar task data described in 49
Results
We demonstrate the three requirements of pretrained spiking neural data model for a BCI: multi- context capable architecture (Section 4.T, beneficial scaled pretraining (Section 4.2) and practical
deployment (Section 4.3 Section 4.4 We discuss the impact of context tokens in Section A.4

---

Obstarle-gudee rench
Weight Initialization
bolaydd roacy
Example Pretraining: 20K trials
OPe
1%
26%
Session
paced
Random
Humar BCI
139
(Bec 3,6
50%
A.
B. Per-Eval Tune Probe: 300 trials
Pretra red
Unsupervised Pretraining
Unsupervised Tuning
Supervised Probe
Learning?
Encoder
Encoder
Encccer
No
Yes
C. Neural
Behavior
Neural
Behavior
Neura
Behavior
Figure Model Training: We model neural activity from huraa and monkey reach. monkey models, evaluation sessions are drawn from a self- self-paced reaching dataset [47]: multi multi-session and multi subtict models pretrain with other sessions in these data. The multi-task model pretrains with the other monkey reach data. Human models use similar volume of data. B. A multi-session model pretrains with data from evaluation subject, with held-out evaluation sessions. Then, for each evaluation session, first do unsupervised tuning off the pretrained model, and then train a supervised prohe off of this tuned model. C. We show which model components are learned (receive gradients) during pretraining and the two tuning stages. For example, supervised probes use an encoder that received both pretraining and tuning on a target session. All tuning is end to end.
Model preparation and evaluation. Most experiments used a 6-layer encoder (~3M parameters). NDT2 adds 2-layer decoder (0.7M parameters) over NDTI [16]; we controls to ensure this extra capacity does not confer benefits to comparison models. To ensure that our models were not bottlenecked by compute capacity in scaling (Section 4.2) models were trained to convergence with early stopping and progressively larger models were trained until no return was observed. We pretrain with causal attention where tokens cannot attend to future timesteps, would be the case in realtime iBCI use (though bidirectional attention improves modeling). We pretrain with 50% masking and dropout of 0.1. Further hyperparameters were not swept in general experiments; initial settings were manually tuned in pilot experiments and verified to be competitive against hyperparameter sweeps. Further training details are given in Section A.3] We briefly compare against prior reported results, but to our knowledge there is no other work that attempts similar pretraining, so we primarily compare performance within NDT-family design choices.
Model preparation and evaluation follows several steps, as summarized in Fig. 2 and C. We evaluate models on randomly drawn held-out test data from select "target" sessions (selection is later detailed per experiment). Separate models are prepared for each target session, either through fine-tuning or by training from scratch for single-contex, models. For unsupervised evaluation, we use the Poisson negative log- ligeliooddoodd (NLL) objective. which measures reconstruction performance of randomly masked bins of test trials (Fig. 2 middle). For supervised evaluation, report the R² of decoded kinematics, which for these experiments are 2D velocity of the reaching effector (Fig. right). These supervised models are separately tuned off the per-evaluation session unsupervised model; so that all supervised scores receive the benefit of unsupervised modeling of target data [50].
4.1 NDT2 enables multicontext pretraining
We evaluate on five temporally spaced evaluation sessions of monkey Indy in the RTT dataset, with both sorted and unsorted processing. Both versions are important; sorted datasets contain information about spike identity and are broadly used in neuroscientific analysis, while unsorted datasets are frequently more practical in BCI applications. Velocity decoding is done by tuning all models further with 2-layer Transformer probe (matching the reconstruction decoder). Here we provide the models with minutes of data (300 training trials) for each evaluation session. This quantity a good litmus test for transfer as it is sufficient to fit reasonable single single-session models but remains practical for calibration. A 10% test split is used each evaluation session (this small % is due to several sessions not containing much more than 300 trials). We pretrain models using approximately 20K trials of data, either with the remaining non-evaluation sessions of monkey Indy Sessii-S the sessions from the other monkey (Multi- (Multi-Subject), or from non-RTT datascts entirely (Multi Maski see Fig. 2)
5

---

A. Sorted
B. Unsorted
Data Source Single-session (Scratch)
NDT2 (Ours)
NDT2 (Ours)
Multi Scssion Multi-Subject
Velecity R²
Velocty R2
& Multi-Task
*
rEFH (Nakin
NDT2 NDT1
NDT1-Stitch
0,290
0,793
0.161
0.172
0.480
Test NLL ( (-)
Test NLL (-)
Figure 3. NDT2 enables pretraining over multi-session, multi-subjeet, and multi-task data. We show unsupervised and supervised performance (mean sessions, SEM intervals of seeds) on sorted (A) and unsorted (B) spiking activity. Higher is better for R², lower is better for negative log-likelihood (NLL). Pretraining data is size-matched at 20Ks, except scratch single-session data. NDT2 improves with pretraining with all data sources, whereas stitching is ineffective. NDT1 aggregation is helpful but does not apply beyond session transfer. A reference well-tuned decoding score from the rEFH model is estimated 48
Prior work in multi-session aggregation either use stitching layers or directly train on multi-day data with consistent unit count. Thus, we use NDTI with stitching baseline for sorted data, and with or without stitching for unsorted data. NDT2 pads observed neurons in any dataset to the nearest patch multiple. All models receive identical context tokens.
We show the performance of these pretrained models for sorted and unsorted data Fig. 3 For context, we show single-session performance achieved by NDTI and NDT2, and the decoding performance the nonlinear rEFH model released with the dataset [48] rEFH scores were estimated by linearly interpolating the 16ms and 32ms scores in Makin et al. [88 Single session performance for NDTI and NDT2 is below this bascline. However, consistent with previous findings on the advantage of spatial modeling [39]. we find single-session NDT2 provides some NLL gain over NDT1. Underperforming this established baseline is not too unexpected: NDT's performance can vary widely depending on the extent of tuning (Transformers span a wide performance range on the NLB, see also Section A.3) Part of the value of pretraining is that it greatly simplifies the tuning needed for model preparation [51].
However, sorted data shown in Fig. BA, all pretrained NDT2 models outperform rEFH and single- session baselines, both in NLL and kinematic decoding. Surprisingly. multi multi-subject data work as well as muulti-esssion data, and multi-task data provide an appreciable improvement as well. NDT-Stitch performs much worse in all cases, and in fact, cross-task pretraining brings NDT-Stitch below the single- senslonssssion baseline. We expect that stitching is less useful here than in other works because other works initialize their stitching layers by exploiting stereotyped task structure across settings (see PCR in (28) ) and we cannot do this because RTT does not have this structure.
The unsorted data has consistent physical meaning (electrode location) across datasets within a subject, which may particularly aid cross-session transfer. Indeed, unsorted cross-session pretraining achieves the best decoding (> 0.7/2²) in these experiments (Fig. 3 blue diamond). The consistent dimensionality should not affect cross-task and cross-subject decoding, yet they also improve vs their sorted analogs, indicating the unsorted format benefits decoding in this dataset. Given this, unsorted formats in subsequent analysis of RTT. Otherwise, relative trends are consistent with the sorted case. Both analyses indicate that different pretraining distributions all provide some benefit for modeling a new target context, but also that not all types of pretraining data are equivalently useful.
4.2 NDT2 scaling across contexts
Naturally, cross-session data are likely the best type of pretraining data for its close relevance, but less relevant data can also occur in much greater volumes. To inform future pretraining data
'The rEFH model data splits vary slightly from ours: its data splits are sequential and contiguous in time, whereas we use random draws in keeping with NLB. NDT2 sorted decoding drops from 0.54 to 0.52 when mirroring Makin al's data splits.

---

Pretrailing data
Targe data
Periarmanc
Perfrmanc
06-
Velocky n2
0.61
Performanng
64
0.60
Pret
30
105
10P
10'
Pretraining triels
Pretraining trials
Target context trials
Single-session (Scratch)
Cross-Session
Cross- Subject
Cross-Task
Figure Scaling transfer on RTT. We compare supervised H²² (A) and unsupervised NLL scaling (B) as we increase the pretraining dataset size. Each point is a model; non single-session models calibrate to evaluation sessions with 100 trials. All pretraining improves over the leftmost 100-trial single-session from-scratch model, though scaling benefits vary by data sources. C. We seek convergence point between pretraining and training from scratch, increase the number of trials
in our target context. Models converge by 3K trials.
efforts, we perform three analyses to coarsely estimate the data affinity [52] of the three different context classes (cross-scssion, cross-subject, and cross-task). Previously these relationships have been grounded in shared linear subspaces [ we now quantify this in the more general generative model encompassed by DNN performance scaling.
Scaling pretraining size. We consider transfer as we scale pretraining size, so as to extrapolate trends that might forecast the utility of progressively larger data. Specifically, we measure performance after tuning varied pretrained models with 100 trials of calibration a novel context. We do this for both supervised (Velocity I?? and unsupervised (NLL) metrics in Fig. 4A and B respectively. To contex- tualize performance of cross cross-context scaling, we measure in mn-distribution scaling of intra-session data (Scratch). For example, the largest cross-session model tuned with 100 trials achieves a NLL similar to IK trials of intra-so sesra-session data (Fig. 4B. Cross-session vs Scratch). This shows that cross- cross-session data can capture a practically long tail of single-session neural variance (experiments rarely exceed IK trials), with nonsaturated benefits in pretraining dataset size. Alternately, cross-session pretraining allows decoding performance similar the largest single-session model performance (Fig, [4 Cross-session vs Scratch), but this scaling is starting to saturate.
Shallower slopes for gross-subject and cross-task pretraining on both metrics indicate poorer transfer to evaluation data. In the unsupervised case (Fig. B), cross-subject and cross-task transfer never exceed the NLL achieved with 400 single-session trall. Note that our task scaling may be pessimistic as we mix human data (Table 2 with monkey data to prepare the largest model, but the trend before this point is still shallow. However, these limitations do not clearly translate to the supervised metric (note that discrepant scaling metrics are also seen in language models (33). For example, the decode R² achieved by the largest task pretrained model compares with in-session models at 800 trials; the same model's 's NLL is most comparable that of a 300 trial single-session model.
Convergence point with from-scratch models. It remains unclear how rapidly pretraining benefits decrease as we increase target context data. We thus study the returns from pretraining as we vary target context calibration sizes (54] (Fig. 4C). Both models yield returns up to 3K trials, which the size of the largest dataset in [47]. Session pretraining provides larger gains, but task pretraining e.g. also respectably halves the data needed to achieve 0.61 NLL. This indicates pretraining is complementary to scaling target session collection efforts. This need not have been the case: even Fig. 4B suggests that task transfer by itself ineffective at modeling the long tail of neural variance. Note that returns supervised evaluation are likely similar or better based Fig. 4W/B; explore related idea in Section 43
represents about 50m of data collection in the monkey datasets, and coincidentally
Overall, the returns using pretrained BCI models depends on the use case. If we are interested in best explaining neural variance, pretraining alone underperforms moderately large in-day data collection effort (single-session model achieves lowest NLL in Fig. 4B). However, we do not see
7

---

A
P4 Implant Day
Day
Cursor Control Pilots
0.8
0
61
63 + 65 S-Shot OLE
70
B
0.7
Pretrained (0-Shot)
O-Shot NDT2
Velacity R2
0.6
C
5-
&&&
OLE
Session
Broad
0.5
4-
0.4
Supervised tune
Reach Time
Scratch
0.3
Unsupervised tune
Scratch
0.2
(100 Supervised Trials)
102
103
larget session trials
0-Shot
Unsupervised Supervised
Figure 5. Tuning and evaluating pretrained decoders. A. In offline experiments, we compare multisession pretrained model calibrated to novel day against current approaches of pretrained model without adaptation (0 (O-Shot) and from-scratch training (yellow, orange). Both supervised and unsupervised outperforms these strategies. Standard error shown over seeds. B. In human control pilot experiments, we evaluated models on test days. "O-Shot" models use no data from test days. C. Average reach times in 40 center-out trials is shown for 3 decoders over 2-3 sessions. This average includes trials where target not acquired in 10s, though this occurs 1-2 times in 40 trials. Session NDT2 uses <250 trials of data, while Broad NDT2 also includes other human and monkey data. Pretrained models provide consistent O-Shot control while OLE ( can sometimes fail (shown by X). Control improves with either unsupervised or supervised tuning. However, OLE appears to overtake NDT2 with supervision.
interference [54] in our experiments, where pretraining then tuning underperforms a from-scratch model. Thus, so long as we can afford the compute, broad pretraining is advantageous; we show these trends are repeated for two other evaluation sessions in Section A.6 We reiterate that our pretraining effort is modestly scaled; the largest pretraining only has 2 orders more data than the largest intra-context models. These conclusions may further strengthen insofar if we are able to better scale curation of pretraining data over individual experimental sessions.
4.3
Using NDT2 for improved decoding on novel days
RTT Decoding. A continuously used BCI presents the optimistic setting of having both broad unsupervised data but also multiple sessions worth of supervision for our decoder. To evaluate NDT2 in this case, we follow 1st stage unsupervised pretraining with a 2nd stage of supervised pretraining of decoder, and finally measure the decoding performance in novel target session in Fig. 5A. We find that given either supervised or unsupervised tuning in our target session, beyond smoothly par
improving over the 0-Shot Pretrained model's 's performance, achieves decoding performance with the best from-scratch models at all data volumes. This is true both in the realistic case where the majority of target-session data unlabeled (100 Supervised Trials) and in the optimistic case when >K trials of supervised data are available (Scratch). As expected, though, gains against the Scratch models largest when target session data are limited. In sum, pretraining and fine-tuning enables practical calibration without explicit domain adaptation designs (as explored e.g. in 37 38))
Human BCI evaluation. We next evaluate NDT2 in offline decoding of human motor intent, i.e. open loop trials of 2D cursor control. This shift from actual reach in monkeys is challenging: human sessions often contain few trials (c.g. 40) and intent labels much noisier than movement recordings (intent labeling is described in Section ID. We now evaluate a temporally contiguous experimental block, but only tune one model over this block, rather than per session, due to the high session count. We also increase test split to 50% to decrease noise from evaluating only few trials.
We also compare broad pretraining (multi-session, subject, and task) performance (Table 2 row 1) with various data ablations. Ablating cross-subject data only has minor performance impact (row I vs 2), while ablating cross-task data, which here indicates removal of data collected during online control and leaving only open-loop observation trials, hurts performance more (row 3 vs 1). We note that NDT2 fails to train using only single-sessions given extremely low trial count; use
8

---

linear decoding performance as a single-session reference linear decoding performance. Overall, the plentiful cross-session data likely occludes further gains from cross- cross-subject and cross-task pretraining. but note that there is no negative transfer.
we assess whether monkey data in either pretraining or decoder preparation improves decoding (Table 2 rows 4-7), We find that monkey data, however incorporated, reduces offline decoding performance (row 4-7 < 1). This cross-species result is the first instance of possible harm from broader pretraining, but warrants further study given the potential of transferring able-bodied monkey decoders humans.
Next, given reports of monkey human transfer 155
Table 2. Human reach intent decoding. We analyze how pretraining data impacts offline decoding in two people, P2 and P3. Checks (V) indicate the data used in pretraining heyond task- and subject-specife data; human data totals 100K trials for P2 and 30K trials for P3. Intervals transfers across task and somewhat across subject, but there is no benefit from monkey data.
SEM over fine-tuning sceds. Pretraining
Neural data (Unsup. pretrain) Behavior (Sup. pretrain)
Velocity R² (t)
Cross- Subject
Cross-
+130K
+24K RTT (Monkey)
Task
Monkey
P2
P3
0.503=0.020
0.515+0.008
2)
0.487=0.007
0.50940.016
3)
0.444=0.007
0.493
4)
0.486=0.012
0.472=00019
/
5)
0.490=0.00
0.477
6)
0.474=0.00 000
0,491
7)
0.443
0.455
8)
Smoothed spike ridge regression (OLE)
0.077
0.208
4.4 NDT2 for human cursor control pilots.
Finally, to assess NDT2's potential for deployed BCIs, we run realtime, closed-loop cursor control with one person, P4. This person was implanted recently (~ months prior), and has high signal quality but limited historical data. We compare two NDT2 models, one of which pretrains with 10 minutes of participant-specfiic data, and one broadly pretrained on all human and monkey data, along with a baseline linear decoder (indirect OLE (56)). We evaluate the setting where decoders from recent days are available, and can either used directly (0-shot) or updated with calibration data on the test day (Fig. 5BB) All models can be supervised with test day data, but NDT2 models can also use unsupervised tuning with only neural data. In a standard center-out BCI-reaching task (methods in Section A..).. NDT2 allows consistent 0-shot use, whereas OLE sometimes fails (Fig. 5C, blue X). After either form of tuning, both NDT2 models (multisession, broad) improve. Importantly, NDT2 tunes without any additional distribution-alignmmen priors, as in [1] 35 33), showing that pretraining 20 may be a viable paradigm for closed loop motor BCIs.
and fine-tuning
Perhaps surprisingly, OLE provides the best control given supervised data. The performance gap and large distribution shift from offline analysis to online control is well known (18), though the specific challenge of DNN control only has basic characterization. For example, NDT2 decodes "pulsar" behavior as in 30 Costello et al. [57] provides a possible diagnosis: these pulses reflect the ballistic reaches of the open loop training data. whereas OLE, due its limited expressivity, will always provide continuous (but less stable) control that can be helpful in time-based metrics. Promising approaches to mitigate the pulsar failure mode include further closed-loop tuning (58) or open-loop augmentation 30) we leave continued study of NDT2 control for future work.
Discussion
NDT2 demonstrates that broad pretraining can improve models of neural spiking activity in the motor cortex. With simple changes to the broadly used masked autoencoder Transformer, NDT2 at once spans the different distribution shifts faced in spiking data. We find distinct scaling slopes for different context classes, as opposed to a constant offset in effectiveness (59]. For assistive applications of BCL, NDT2's simple recipe for multisession aggregation is promising even if the ideal scenario of

---

cross-species transfer remains unclear. More broadly, we conclude that pretraining, even at a modest 10-100K trials, useful in realistic deployment scenarios with varied levels of supervised data.
Supervising pretrained BCI decoders. Motor BCIs fundamentally bridge two modalities: brain data and behavior. NDT2 allows neural data pretraining, but leaves open the challenge of decoding of the full range of motor behavior. Without mapping that range, BCIs based on neural data pretraining alone will need continuous supervision. This is still practical path forward: beyond explicit calibration phases, strategies for supervising BCIs are diverse, and can derive from user feedback [60], neural error signals 661 62,, or task-based estimation of user intent 63
58 64
The ambition of pretrained
BCI models, with broad paired coverage of the experimental costs of data collection; the field of robotics suggests that scaled offline or simulated learning are important strategies given this expense. Since we lack convincing closed-loop neural data simulators (though see 65 63)) understanding how to leverage behavior from existing heterogeneous datasets is an important next step.
domain will be challenging given the
Negative NLB result. NDT2 performance did not exceed current NLB SoTA on motor datasets (RTT, Maze) [66). We first note the NLB evaluation emphasizes neural reconstruction of held- out neurons, which differs from our primary emphasis on decoding. Second, our unsupervised scaling analysis indicates that modest pretraining (100K trials) provides limited gains for neural reconstruction, especially when the target dataset has many trials (Fig. (C), as the NLB RTT dataset, so underperformance does not contradict our results. Revisiting the NLB after further scaling may be fruitful.
Limitations. NDT2's design space is under-explored. For example, we do not claim that full space- time attention is necessary over factorization. While NDT2 achieves positive transfer, further gains may come from precise mapping of context relationships [52]. Further, it is difficult to extrapolate the benefits of scaling beyond what was explored here, particularly with gains in unsupervised recon- struction appearing very limited. Our evaluation also has a limited scope, emphasizing reach-like behaviors. While these behaviors are more general than previous demonstrations of context trans- fer [1 38 33 35), evaluating more complex behavior decoding is practical priority. Further, NDT2 benefits are modest in human cursor control pilots, reiterating the broadly documented challenge in translating gains in offline analyses to online, human control [30, [8]. Finally, design parameters such as masking ratio may affect scaling trends, which we cannot assess due to compute limits.
Broader Impacts. Pretrained iBCI models may great improve iBCI usability. However, DNNs may require further safeguards to ensure that decoded behaviors, especially in real-time control scenarios, operate within reasonable safety parameters. Also, pretraining will require data from many different sources, but the landscape around human neural data privacy is still developing. While there are very few humans involved in these experiments, true deidentification remains difficult, requiring, at a minimum, consented data releases.
Acknowledgments and Disclosure of Funding
We thank Jeff Weiss, William Hockeimer, Brian Dekleva, and Nicolas Kunigk for technical assistance in trawling human cursor experiments. We thank Cecilia Gallego-Carraee and Patrick Marino for packaging data for this project (results do not include analysis from PM's data). We thank Joseph O'Doherty. Felix Pei, and Brianna Karpowiez for early technical advice. JY is supported by the DOE CSGF (i.e. U.S. Department of Energy. Office of Science, Office of Advanced Scientific Computing Research, Department of Energy Computational Science Graduate Fellowship under Award Number DE-SC0023112) RG consults for Blackrock Microsystems and is on the scientific advisory board of Neurowired. This research was supported in part by the University of Pittsburgh Center for Research Computing, RRID:SCR 022735, through the resources provided. Research reported in this publication was supported by the National Institute of Neurological Disorders and Stroke of the National Institutes of Health under Award Numbers RO ROINS121079 and UH3NS 1NNNS107144 14. The content is solely the responsibility of the authors and does not necessarily represent the official views of the National Institutes of Health.
This report was prepared as an account of work sponsored by an agency of the United States Government. Neither the United States Government nor any agency thereof, nor any of their employees, makes any warranty, express or implied, or assumes any legal liability or responsibility for the accuracy, completeness, usefulness of any information, apparatus, product, or process
10

---

disclosed, or represents that its use would not infringe privately owned rights. Reference herein to any specific commercial product, process, or service by trade name, trademark, manufacturer, or otherwise does not necessarily constitute or imply its endorsement, recommendation, or favoring by the United States Government or any agency thereof. The views and opinions of authors expressed herein do not necessarily state or reflect those of the United States Government or any agency thereof.
References
[1] Saurabh Vyas, Matthew D Golub, David Sussillo, and Krishna V Shenoy. Computation through neural population dynamics. Annual review of neuroscience, 43:249- 275, 2020.
[2] Ege Altan, Sara A. Solla, Lee E. Miller, and Eric J. Perreault. Estimating the dimensionality of the manifold underlying multi-electrodee neural recordings. PLOS Computational Biology, 17(11):1-23, 2021. doi: 10. 1/jouralll 1008000
1008591. URL http:::tttt/oiicttt /ttp::////oi..r/////// 11t.1/..111//ooaaaooooo pcbi. 1008591
[3] Daniel J O' Shea, Lea Duncker, Werapong Goo, Xulu Sun, Saurabh Vyas, Eric M Trautmann, Ilka Diester, Charu Ramakrishnan, Karl Deisseroth, Maneesh Sahani, et al. Direct neural perturbations reveal dynamical mechanism for robust computation. bioRxiv, pages 2022- 12, 2022.
[4] Patrick T Sadtler, Kristin M Quick, Matthew D Golub, Steven M Chase, Stephen 1 Ryu, Elizabeth C Tyler-Kabara, Byron M Yu, and Aaron P Batista. Neural constraints learning. Nature, 512(7515): 423.426, 2014.
[5] Saurabh Vyas, Nir Even-Chen, Sergey D Stavisky, Stephen I Ryu, Paul Nuyujukian, and Krishna V Shenoy. Neural population dynamics underlying motor learning transfer. Neuron, 1775) 2018.
[6] Juan A. Gallego, Matthew G. Perich, Stephanie N. Naufel, Christian Ethier, Sara A. Solla, and Lee E. Miller. Cortical population activity within preserved neural manifold underlies multiple motor behaviors. Nature Communications, 9 91))4233, Oct 2018. ISSN 2041-1723, doi: 10.1038/s41467-111 1--006600- URL
https://www.natt www.nnatturcnnoe com/ ar ticles/s41111- 0101767-888888-----111111
[7] Juan A. Gallego, Maithew G. Perich, Raeed H. Chowdhury. Sara A. Solla, and Lee E. Miller. Long-term stability of cortical population dynamics underlying consistent behavior. Nature Neuroscience, 23(2): 260-270, Feb 2020. ISSN 1546- 1726. doi: 10. 10.10088/115522--1--5555---
https://t:///ww http://wwwwaat...eeeee/ee com/
ar articles/s41593 019-3---11-555-----
[8] Mostafa Safaie, Joanna C. Chang, Junchol Park, Lee E. Miller, Joshua T. Dudman, Matthew G. Perich, and Juan A. Gallego. Preserved neural population dynamics across animals performing similar behaviour. Sep 2022. doi: 022110122222.0.22.5509222
26..0.9.......
URL/https::/wwwwbior wwwp.//wwww..//oeee//////t/// 1101/2022.09
26. 58948919
[9] Max Dabagia, Konrad P. Kording, and Eva I.. Dyer. Aligning latent representations of neural activity. Nature Biomedical Engineering. 7(4):337-343, Apr 2023. ISSN 2157-846X. doi: 10.
URL httt:hhtps:///mt http://wturre.coo/arrccceeeee//e
022-28
[10] David Sussillo, Sergey D. Stavisky, Jonathan C. Kao, Stephen I. Ryu, and Krishna Shenoy, Making train-maching interfaces robust to future neural variability. 7:13749, Dec 2016. ISSN 2041-1723. doi: 10. 00.1038/ccomms11337 URL httphhtps::/// //www nature. con.//articless/nooom133333.
[11] Brianna M. Karpowiez, Yahia H. Ali, Lahiru N. Wimalasena, Andrew R. Sedler, Mohammad Reza Keshtkaran, Kevin Bodkin, Xuan Ma, Lee E. Miller, and Chethan Pandarinath. Stabilizing brain-computer interfaces through alignment of latent dynamics. Nov 2022. doi: 10.1 111100222220000648887866 URL
https://www wworxvv.org/oontet/ooooo T181/2022 06.06. 487388v2
[12] Francis Willett, Erin Kunz, Chaofei Fan, Donald Avansino, Guy Wilson, Fun Young Choi, Foram Kamdar, Leigh Hochberg, Shaul Druckmann, Krishna V. Shenoy, and Jaimie M. Henderson. hegf-performance speech neuroprosthesis. Jan 2023. doi: 00.1101/002300..2..224444
URL https://www. bworxiv.org
Content/18. 10012/02.1102221111 01. 01. 524489v1 21.
[13] B Wodlinger, E Downey, E C Tyler-Kabara, A B Schwartz, M L Boninger, and J L Collinger. Ten- dimensional anthropomorphic arm control a human brain- machine interface: difficulties, solutions, and limitations. Journal of Neural Engineering. 1261)111611, Feb 2015. ISSN 1741 2560, 1741 -5522, doi: 10. 2550////41120000///10000000 URL http::://o/oosiieeeee iop wrg./rticceee...eer htt.://eceee.e.eee/11111111
016811

---

[14] Felix Pei, Joel Ye, David Zoltowski. Anqi Wu, Raeed H. Chowdhury, Hansem Sohn, Joseph E. O'Doherty. Krishna V. Shenoy, Matthew T. Kaufiman, Mark Churchland, Mehrdad Jazayeri, Lee E. Miller, Jonathan Pillow, Memming Park, Eva L. Dyer, and Chethan Pandarinath, Neural latents benchmark '21: Evaluating latent variable models of neural population activity, 2022.
[15] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkorcit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, Vishwanathan, and R. Garnett, editors, Advances Neural Information Processing Systems. volume 30. Curran Associates, Inc., 2017. URL URL
PaperraaaPaarpppe ppff
[16] Joel Ye and Chethan Pandarinath. Representation learning for neural population activity with neural data transformers, 2021, URL, htttps://dol.org//// 51620/00112.80.000 27358
[17] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. Dec 2021. doi: 10.48550karxi:22111..66 2111. 86377 arXiv:2111.06377 [es].
/ttp:///:///rxx wwwperwwww.oooveoooo.oe org//bs//
[18] Chethan Pandarinath and Sliman J Bensmaia. The science and engineering behind sensitized brain- controlled bionic hands. Physiological Reviews, 102(2):551-60 2022.
[19] Daniel B. Silversmith, Reza Abiri, Nicholas F. Hardy, Nikhilesh Natraj, Adelyn Tu- Tu-Chan, Edward F. Chang, and Karunesh Ganguly. Nature Biotechnology, 396)):266-335, Mar 2021. ISSN 1546-16990 1446-1696,..6 doi: 10. (1188/4115870020-6666
URL httpsttp////// // www. nature. http////rrtccccl//////////-----/ 0662-5
[20] Jacob Devlin, Ming-Wor Wei Chang, Kenton I Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirec- tional transformers for language understanding. 2019.
[21] Sabera J Talukder. Jennifer J. Sun, Matthew K Leonard, Bingni W Brunton, and Yisong Yue. Deep neural imputation: A framework for recovering incomplete brain recordings. In NeurlPS 2022 Workshop on Learning from Time Series for Health, 2022. URL httpp:///peereiiciee www.ppperreiiie..eeoerrrurrruree net:/oorunddd=s g8UrIcn
[22] Samuel M Peterson, Shiva H Singh, Benjamin Dichter, Kelvin Tan, Craig DiBartolomeo, Devapratim Theogarajan, Peter Fisher, and Josef Parvizi. Ajile12: Long-term naturalistic human intracranial neural recordings and pose. Scientific Data, 9(1):184, 2022. ISSN 2052-4463. doi: 10.1103//411977 599-02201-97-
URL httt::::/doii.rreeeeeee
/03/////1/9--///// 022- 022-5597772222----2-7- 01288-y
[23] Christopher Wang. Vighnesh Subramaniam, Adam Uri Yaari, Gabriel Kreiman, Boris Katz, Ignacio Cases, and Andrei Barbu. BrainBERT: superviee representation learning for intracranial recordings. In The Eleventh International Conference on Learning Representations, 2023. httttts://ppeeeee/////
foxum? reUn6
[24] Armin W. Thomas, Christopher Ré, and Russell A. Poldrack. Self-supervised learning of brain dynamics from broad neuroimaging data, 2023.
[25] Demetres Kostas, Stéphane Aroca- Ouellette, and Frank Rudzicz. BENDR:
Using transformers and contrastive self-supervised learning task learn from massive amounts of EEG data. Frontiers in Human Neuroscience, 15, 2021. ISSN 1662-5161. doi: 10.3389//hhum.2221.65333
URL hhtpp
7/ww. www..rrotierrinnoor//rrracceeeee
3389/ 3389/fnhun...m../ 2021.653659
[26] Tyad Obeid and Joseph Picone. The temple university hospital EEG data corpus. Frontiers in Neuroscience, 10:196, 2016. ISSN 1662-453X.
doi: (0.3389/ninn.20066.000 URL. http:s://ww..rootiirssii
articles/10. http:///1..33//finnn.eem 3389/fnins wewplins..e...... 2016.00195
[27] Srini Turaga, Lars Buesing, Adam M Packer, Henry Dalgleish. Noah Pettit, Michael Hausser, and Jakob H Macke. Inferring neural population dynamics from multiple partial recordings of the same neural circuit. In C.J. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K.Q. Weinberger, editors, Advances Neural Information Processing Systems, volume 26. Curran Associates, Inc., 2013. Utt/htts:://pp httpll::s///oreednnng/ /aalee//paperr/1////1eeeee wttp////////1////eeee/ http:///wereeeot..aa/aaaaaaaaaaa 13885666aee01ccaa4c7776666663333 Paperr
[28] Chethan Pandarinath, Daniel J. O' Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D. Stavisky, Jonathan C. Kao, Eric M. Trautmann, Matthew T. Kaufmman Stephen Ryu, Leigh R. Hochberg, Jaimie M. Henderson. Krishna V. Shenoy, L. F. Abbott, and David Sussillo. Inferring single-trie trial neural population dynamics using sequential auto- encoderss. Nature Methods, (50)))805 815, Oct 2018 ISSN 1548-7105. doi:
10. 10.1038/441992-018 011-1999--0111999999911111111999 URL http::tttt/////t nature, http://wwwnaauueee.o///rrrrrtees htt.uurrr.oo/rrrrrieee////122120
12

---

[29] Francis R. Willet, Donald T. Avansino, Leigh R. Hochberg, Jaimie M. Henderson, and Krishna V. Shenoy. High-performance brain-to-text communication via handwriting. Nature, 593(7888):2499255 May 2021. ISSN 1476-4687.
doi: 10.1088641336-021--3306---- URL https///www.ntuure..comrrrrrrtit
s41586- 123866-2118888 021-03506-2
[30] Darrel Deo, Francis R. Willett, Donald Avansino, Leigh R. Hochberg, Jaimic M. Henderson, and Krishna V. Shenoy. Translating deep learning to neuroprosthetic control. Apr 2023, doi: 10.1101/2023.04 21.537581. URL httLhhtts:///www wwt.:/oriii....rr/ooee////////// 11010//022//11111112223333 1101/2023. 84.21. 21.55.53781
[31] Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis. Learnable latent embeddings for joint behavioural and neural analysis. Nature, May 2023. ISSN 1476-4687. doi: 10. 10.180//41886 023-06031-
URL hittps://dof org/18. 1038/s41586 022441666602---60 2300601-666
[32] Nur Ahmadi, Timothy Constandinou, and Christos-Savee -Satvss-Svvvsssssse Bouganis. Robust and accurate decoding of hand kinematics from entire spiking activity using deep learning. Journal of Neural Engineering, 18(2): 026011, Feb 2021. ISSN 25522. doi: 10.1088/1 741-2552/abbd8a. URL https://iopsciences org/article/18 88 T088/ 1741 2552/ 2552/bbde8//mmtt
[33] Justin Jude, Matthew G. Perich, Lee E. Miller, and Matthias H. Hennig, Robust alignment of cross-session recordings of neural population activity by behaviour via unsupervised domain adaptation. Feb 2022. doi: 114855505n211.222220616555
httpp://axxw..or//dbbeee
06159 arXiv:2202.061.99 [es, q-bio]
[34] John E. Downey, Nathaniel Schwed, Steven M. Chase, Andrew B. Schwartz. and Jennifer L. Collinger. Intracortical recording stability in human brain-computer interface users. Journal of Neural Engineering. 154)00440,,,, Aug 2018. ISSN 1741-2552.
doi: 10.1088/7411-2552aaab
[35] Justin Jude, Matthew G. Perich, Lee E. Miller, and Matthias H. Hennig. Capturing cross-scssion neural population variability through self-supervised identification of consistent neuron ensembles. Proceedings the 1st NeurlPS Workshop on Symmetry and Geometry in Neural Representations, page 234-257. PMLR, Feb 2023. URL https::/pprcceeiinggg mlr. pttpps//v17/judd2ea.... html
[36] Ran Liu, Mehdi Azabou, Max Dabagia, Jingyun Xiao, and L. Dyer. Seeing the forest and the tree: Building representations of both individual and collective dynamics with transformers. Oct 2022. URL
https:://peerrevicw.. net
[37] Ali Farshchian, Juan A. Gallego, Joseph Cohen, Yoshua Bengio, Lee Miller, and Sara A. Solla. Adversarial domain adaptation for stable brain-machine interfaces. Jan 2019. URL httph:ht/////eeeeeeeie net net/forumiid-Hyyee @q
[38] Xuan Ma, Fabio Rizzoglio, Eric J. Perreault, Lee E. Miller, and Ann Kennedy. Using adversarial networks to extend brain computer interface decoding accuracy over time. Aug 2022. doi: 10.1101//022.00....5544 URI. https://www. www.xiii.owooooo....ooooooooooee wwt.://ee///co//eeeeeeee// 1181/2822 88. 504777v1
[39] Trung Le and Eli Shlizerman. Stndr: Modeling neural population activity with a spatiotemporal transformer, 2022.
[40] Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučié, and Cordelia Schmid. Vivit: A video vision transformer, 2021.
[41] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkorcit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference Learning Representations, 2021. htt.tsss///opeeevieeeeee//rrrrrr VicbFdNTTy
[42] Maryam M. Shanechi, Amy L. Orsborn, Helene G. Moorman, Suraj Gowda, Siddharth Dangi and Josc M. Carmena. Rapid control and foedback rates enhance neuroprosthetic control. Nature Communications, 8(1):13825, Jan 2017. ISSN 2041-1723. doi: 10. 1a..0033//ommm1.3333 13825. URL htttps://ww.. nwwura.tuecmmmm com/
articles/ncomms 13825
[43] Jennifer L Collinger, Brian Wodlinger. John Downey, Wei Wang. Elizabeth TTyler-Kabare Douglas Weber, Angus JC McMorland, Meel Velliste, Michael L Boninger, and Andrew B Schwartz. High- performance neuroprosthetic control by individual with tetraplegia. The Lancet, 38 (9886)) 2013.
13

---

[44] Eric M. Trautmann, Janis K. Hesse, Gabriel M. Stine, Ruobing Xia, Shude Zhu, Daniel J. O'Shea, Bill Karsh, Jennifer Colonell, Frank F. Lanfranchi, Saurabh Vyas, Andrew Zimnik, Natalie A. Steinmann. Daniel A. Wagenaar, Alexandru Andrei, Carolina Mora Lopez, John O' Callaghan, Jan Putzeys, Bogdan Raducanu, Marleen Welkenhuysen, Mark Churchland, Tirin Moore, Michael Shadlen, Krishna Shenoy, Doris Tsao, Barundeb Dutta, and Timothy Harris. Large-scale high-density brain-wide neural recording in nonhuman primates. bioRxiv, 2023, doi: 10.1 10111000222...2.00 content/earlylee content/earlly/222//8///// httpaaaly//e////////222202 02. 526664
URL http:://s://www wwtp:/wwwwwwiioooooooooooooo..oo org/
[45] Brian Lester, Rami AI-Rfou, and Noah Constant, The power scale for parameter-efficient prompt tuning, 2021.
[46] Xue Bin Peng, Erwin Coumans, Tingnan Zhang, Tsang- Wei Lee, Jie Tan, and Sergey Levine. Learning agile robotic locomotion skills by imitating animals. arXiv preprint arXiv:2004.00884, 2020.
[47] Joseph E. O'Doherty, Mariana M. B. Cardoso, Joseph G. Makin, and Philip N. Sabes. Nonhuman primate reaching with multichannel sensorimotor cortex electrophysiology, May 2017. URL https
//doi. org/10. 5281/zenodo. 788569
[48] Joseph G Makin, Joseph E O' O'Doherty, Mariana M B Cardoso, and Philip N Sabes. Superior arm-movement decoding from cortex with a new, ansupervised--caming algorithm. Journal of Neural Engineering. 15 doi: (108/1741-2552/aaaee //iopscience. Top orggarticle/10 1888/1741- 2aei112555222/55555
(2):026010, Apr 2018. ISSN 1741 -560,, 1741-2552,
URL https
[49] Brian M Dekleva, Jeffrey M Weiss, Michael L Boninger, and Jennifer L Collinger. Generalizable cursor click decoding using grasp-related neural transients. Journal of Neural Engineering, 18(4):0460e,, Aug 2021. ISSN 1541-5560, 1741-2552. doi: 10.108////44122555/ee12222 URL/httsss//opccieee..io...
article/18. 188//7741-25222cc6622
[50] Kundan Krishna, Saurabh Garg, Jeffrey P. Bigham, and Zachary C. Lipton. Downstream datasets make surprisingly good pretraining corpora, 2022.
[51] Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigeerver, Jessica Yung. Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning. 2020.
[52] Amir Zamir, Alexander Sax, William Shen, Leonidas Guibas, Jitendra Malik, and Silvio Savarese. Taskon- omy: Disentangling task transfer learning. 2018
[53] Ashish Teku Vaswani, Dani Yogatama, Don Metzler, Hyung Won Chung, Jinfeng Rao, Liam B. Fedus, Mostafa Dehghami, Samira Abnar, Sharan Narang, and Yi Tay. Scale efficiently: Insights from pre-training and fine-tuning transformers. 2022.
[54] Danny Hernandez, Jared Kaplan, Tom Henighan, and Sam McCandlish. Scaling laws for transfer, 2021.
[55] Fabio Rizzoglio, Ege Altan, Xuan Ma, Kevin L. Bodkin, Brian M. Dekleva, Sara A. Solla, Ann Kennedy, and Lee F. Miller. Monkey-to-human transfer brain-computer interface decoders. bioRxiv, 2022. doi: 1001101//022.1112255111
URL
12.51 52540
[56] Wei Wang. Sherwin S. Chan, Dustin A. Heldman, and Daniel W. Moran. Motor cortical representation of position and velocity during reaching. Journal of Neurophysiology, (0d))44228-27000 Jun 2007. ISSN 0022-3077. doi: 10. 1552/j..0 11011..0001000020000
[57] Joseph T Costello, Hisham Temmar, Luis H Cubillos, Matthew J Mender, Dylan M Wallace, Matthew S Willsey, Parag G Patil, and Cynthia A Chestek. Balancing memorization and generalization in mns for high performance brain-machine interfaces. bioRxiv, pages 2023-05, 2023.
[58] Beata Jarosiewicz, Anish A Sarma, Jad Saab, Brian Franco, Sydney S Cash, Emad Eskandar, and Leigh R Hochberg. Retrospectively supervised click decoder calibration for self-calibrating point-and-click interfaces. Journal of Physiology-Paris 1 10(4))382-391, 2016.
[59] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020.
[60] Jensen Gao, Siddharth Reddy, Glen Berseth, Nicholas Hardy, Nikhilesh Natraj, Karunesh Ganguly, Anca D Dragan, and Sergey Levine. X2t: Training an x-to-text typing interface with online learning from user feedback. arXiv preprint 2tl::2203.022.. 2022.
14

---

[61] Nir Even-Chen Sergey D Stavisky, Jonathan C Kao. Stephen Ryu, and Krishna V Shenoy, Augmenting intracortical brain-machine interface with neurally driven emror detectors. Journal of neural engineering. 14(6):066007, 2017.
[62] Iñaki Iturrate. Ricardo Chavarriaga, Luis Montesano, Javier Minguez. and José del R. Millán. Teaching brain-machine interfaces as an alternatiee paradigm to neuroprostheties control. Scientific Reports, 5(1): 13893. Sep 2015. ISSN 2045-2322. doi: 10.1038/srep13895. URL htttp::/www.nature.comm srep13893
[63] Guy H. Wilson, Francis R. Willett, Elias A. Stein, Foram Kamdar, Donald Avansino, Leigh R. Hochberg. Krishna V. Shenoy. Shaul Druckmann, and Jaimie M. Henderson. Long-term unsupervised recalibration of cursor bcis. Feb 2023. doi: 10.1101/2003302...5522222
htthtts::///ww biorxiv /
1101/2023, 03.03.22 527022v1
[64] Josh Merel, David Carlson, Liam Paninski, and John P Cunningham. Neuroprosthetic decoder training imitation learning. PLoS computational biology, 1204494, 2016.
[65] Francis R. Willett, Daniel R. Young, Brian A. Murphy, William D. Memberg, Christine H. Blabe, Chethan Pandarinath, Sergey Stavisky, Paymon Rezaii, Jad Saab, Benjamin L. Walter, Jennifer A. Sweet, Jonathan Miller, Jaimic Henderson, Krishna Shenoy, John D. Simeral, Beata Jarosicwicz, Leigh Hochberg, Robert F. Kirsch, and A. Bolu Ajiboyc. Principled bei decoder design and parameter selection using feedback control model. Scientific Reports, (1):8881, Jun 2019. ISSN 2045-2322. doi: 10.1038/
-41598-019-4416677 URL /ttp:///0000.///.///////11111110
[66] Unknown. EvalAT leaderboard. httpa//////////////////////////a 3184. 2022. Accessed May 16, 2023.
[67] Christoph Feichtenhofer, Haoqi Fan, Yanghao Li, and Kaiming He. Masked autoencoders spatiotemporal learners, 2022.
15

---

A Supplementary Material
A.1 Dataset Preparation
We aimed to perform minimal preprocessing datasets. For pretraining, do not explicitly filter for successful trial outcome as done in some neuroscientific analyses (though some datasets are released with erratic outcomes pre-filtered). Neither do we (beyond what is provided directly in datasets) filter for cross-correlated channels or low-firing neurons. We also do not z-score neuronal firing, both for simplicity and The one exception to this is that neurons with firing < 0.5Hz are removed in the sorted analysis the O' O'Doherty RTT dataset, to reduce the number of spatial channels below 288. As some datasets report single unit activity and some report multi-unit activity, the dynamic range of the input data varied by an order of magnitude, with baseline firing rates varying between 0. Hz upwards of 50Hz. We believe additional data curation is likely to improve model quality.
not remove any potential chanee/ssnnnn/sssioon information.
In total, the max number of pretraining trials pseudo trials on the order of 100K trials. Each trial lasted up to 2.5s (cropped or chunked trials were longer), and used all recorded MI activity, and PMd activity if available.
Decoding targets were either standard unit or in z-scores against the dataset mean and standard deviation (not session specific z-score), Standard units of meters/second were primarily used in most RTT analysis, except when preparing an RTT/human BCI decoder, which used respective 7-scores.
Quantifying volume Should neural data pretraining proliferate, it will be important to distinguish pretraining efforts by neural data volume. This is tricky as the most pertinent unit of neural data seems likely to depend convenient basis of time may lead to the illusion of standardization. For example, some datasets may contain much more behavior than others in the same time periods. This issue is still present, though less severe, if we presume to only differentiate data up to orders of magnitude: compare experimental behavior vs naturalistic behavior, which may have extended periods of rest. For now, we see it reasonable to quantify the value of data in pretraining with respect to semantic units of covariate data (experimental trials), and encourage documentation of pretraining composition.
the function of the neural data; we are concerned comparisons on the
Reaching datasets
Neural Latents Benchmark motor datasets (MC_Maze, MC_Mazeesmalll MC_Maze_med,
Sc__aze_larre MC_RTT): -3.7K trials.
Churchland et al. obstacle-guided (maze) reaching, 2 monkeys, 9 sessions / ~00 trials.
Nir-Even Chen et al., delayed reaching, 2 monkeys, 12 sessions / ~ 80K trials total.
O' D'Doherty et al., self-paced reaching, 2 monkeys, 47 sessions / ~40K seconds total.
Isometric manipulandum datasets
Gallego-Carracee et al. isometric center-out and hold, monkeys, 12 sessions, ~2.7K trials total.
Dyer al. 2 monkeys (same above), 4 sessions/~7700 trials total.
Human BCI datasets [Private]
Human participant data from ongoing clinical trials. Subsetted to 2D cursor control activity, either under observation/attemmpeed activity, partial, or full BCI control. During observation, participants observe a programmatically controlled cursors, which e.g. is performing center out at steady pace in trialized fashion. 130K trials. Note only observation trials are used for evaluating decoding.
A.2 Methods for online cursor control experiments
We evaluate realtime, online cursor control using a simple Center-Out reaching task. This is trialized task where the cursor initializes on the center of screen and the participant reaches to of 8 radially
16

---

spaced target that was cued at the start of the trial. The trial ends after the cursor is held at the target for 0.25s (preventing fly-bys from ending the trial). During open-loop calibration, the cursor autopiloted to the target. Velocity labels are produced by assuming that the participant's intention matches the autopilot presentation. Specifically: we take the velocity of the cursor on the screen and apply boxcar filter of 500ms to remove nonsmooth rendering, using the smoothed result as our velocity label. To span more reach conditions, we use an open-loop calibration that does not reset the cursor at the center after each trial, but instead immediately presents the next target. Our unsupervised tuning uses the neural data, but not the labels, collected in the same open loop calibration.
After each of 8 decoders (0-shot and supervised for OLE, O-shot; supervised and unsupervised for 2 NDT models) are prepared, we evaluate each decoder in a block of 40 trials. Evaluation blocks proceed without randomization (OLE always began evaluation as it is the fastest to prepare); our results may advantage earlier decoders due to participant fatigue which accumulates over the session.
A.3 Compute and Hyperparameter Tuning
training are available at interested users may compare specific hyperparameters logged there against those in the codebase during reproduction attempts.
The full.
uncurated
logs of all model
https//wwadbbaa//ocycc///tttttt gener
general bei pofPworrspace-user-jjee
Basic hyperparameters
1. In both pretraining and fine-tuning, we scale batch size (accumulating batches or using multi-GPU training when necessary) to be roughly proportional to full dataset so that each epoch requires 10-100 steps; we find performance is not too sensitive to batch size within an order of magnitude of this heuristic (especially in pretraining).
2. In pretraining we manually tuned LR to 5e - 4 in initial experiments and hold fairly constant in pretraining. We swept learning rate in hyperparameter comparisons below.
3. In pretraining, This is a high threshold that is typically not reached: training converges within 100-1K 100-1K epochs for manually tuned LR. In fine- tunintunnn, we experimented with similar ramping schedule but settled on fixed small LR (which are typically grid-searched).
learning rate warm-up for 100 epochs, and decay to by 2500 epochs.
4. For RTT, we swept and found that decoding lag of 120ms worked reasonably well. (This is similar to reports in [14]. For human BCI, we do not use decoding lag.
5. For human offline evaluation, we take the best of two evaluation hyperparameter settings: 10% or 50% masking during target-session tuning. We also report the R² only for times where the intent is non-zero; participants are not typically perfectly zero-intent during the majority of non-zero phases (i.e. data are noisy). We do not filter data by putative quality as measured by online performance in the experiment in which the data was collected; thus our calibration data includes several noisy, incomplete trials as well. Evaluation data are restricted to a contiguous set of sessions with non-trivial linear decoding.
6. Hyperparameters were not swept further during realtime pilots.
Compute costs We estimate computational costs with respect to data volume, as model size is held relatively static (6-12 layers, 128-384 hidden size). Most analysis was run on SLURRM clusters. Pilot realtime decoding used an NVIDIA 1060/2060, where loop time was <15ms (neural data was binned at 20ms). Model tuning was performed on remote server and took <5 minutes.
1. Fitting datasets on the order of 1K trials typically requires 20m- 1hr on 12G 1080/2080-series NVIDIA GPUs
2. 10K-20K trial datasets require 2-8 32G-V100 100 hours.
3. 100K+ datasets require 72 80G-A100 hours.
A.4 NDT2 Design Notes
NDT2 introduces two primary design elements: context tokens, and patch size. We first consider the effect of context tokens. NDT2 integrates context tokens directly by adding learned tokens
17

---

0.99
0.338
0.324
Test loss
11.39
0.312
0130
0.321
Test
Session Tokens
Subject Tokens
0.32
10?
10
Training Steps
Training Steps
Figure Context embedding ablations. A. Multi-session model training curves, with varying session context token count (3 seeds). Learning is improved, but additional tokens do not have notable effect. The converged score in only modestly affected. B. Subject transfer training curves with similarly varied token budget for subject embedding. Models receive 1 session token. There is clear additional improvement nor harm.
Sorted
0.6
Intra-session
Velocity R (t)
Cross-Session
Makin 18 (rEFH Est)
Cross-Subject
Cross-Task
Patch 32 (Pool)
Patch 32 (No pool)
Patch 8 (Pool)
Patch (No pool)
x
NDT1
0.289
0.291 Test NLL (
0.294
Figure 7. 32-neuron patches compared against 8-neuron patches. Smaller patches sometimes improve NLL but reliably decode poorly. (Pool)-ed ed models average-poo -ool all the neural data features for given timestep before being given to decoder, whereas regular decoders simply cross- gross-attend neural data features.
and adding them to the data token sequence (i.e. in-context aggregation). Our pilot experiments found difference using cross-attention integration. The training curves of sorted multisession models augmented with context tokens, shown in Fig. [6A, demonstrate a primary effect in speeding convergence, which can be valuable in large scale pretraining. The benefit converged NLL (some le 3) is modest but negliigi,,,,,, considering the NLL resolution in Fig. 3 This trend replicates smaller data scales (Fig. IT]. Providing one session token and additionally varying the available subject tokens (Fig. 6B) has much smaller effects. However, given no harm and negligible compute overhead, we hold as a default policy to provide one token for each of session, subject, and task.
Patch size is a less intuitive hyperparameter in NDT2 than similar architecture in the image domain, ViT. This is because units within patch not meaningfully related, especially across datasets; adjacent units often correspond to different physical locations on an electrode array, and this physical meaning typically underspecified in open data releases. Smaller patch sizes, ideally the single neuron level, thus appear to be a better design.
However, in carly experiments we explored patch size but quickly found tradcoff. Smaller patch sizes do appear to incrementally improve neural data models, but are both more expensive computationally (to train) and statistically (to learn decoders off of). We show this in Fig. 7 Note how the unsupervised NLL is similar or better with smaller patches, but decoding is dramatically worse, regardless of whether
pool across the population each timestep (Pool) or not. Smaller patches may be worth revisiting we have a high amount of supervised data to train the decoder.
Architectural details. We refer readers to the codebase for full details, but provide an overview here. For each patch, NDT2 linearly projects input spike counts and concatenates these projections to form the patch embedding, zero-padding if the patch isn't full. Input spikes pass through Transformer
18

---

Sample sweep for Indy 06/27/2016 (Sorted)
Sample HP sweep for Indy 06/27/2016 (Unsorted)
drpper
0.365
Poisson NLL (1)
0.1125
0.300
0.100
NLL
0.352
0.0050
0.350
0.0233
depest
0,000
0.345
E4
NDT2 (Patch 32)
NOT2 + Setcch
NDT1
NDY2 (Patch 32)
NOT1 + Stitch
NOT1
Sample HP sweep for Indy 06/27/2016 (Sorted)
Sample HP sweep for Indy 06/27/2016 (Unsorted)
0.385
5.15
0.5150
Wight cEC1)
D101
5.04
0012
0.360
0.022
Poisson NLL
[04)
0043
0.352
Poisson
0.0075
0,0050
0.350
0.0000
0:977
NDTZ (Patch 32)
NOT1 Stitch
NDT1
NDT2 (Patch 32)
NOT1 + Suichh
NOT1
Figure 8, Sweeps for regularization parameters. NDT2 requires lower dropout.
encoder and decoder layers, and the final layer's output then linearly projected to predict Poisson rates or effector velocity. NDT2 used pre-normalization layers but otherwise leave the Pytorch implementation of the Transformer layers untouched.
HP Sweeps.
We briefly show that NDT2 achieves higher performances than comparisons when sweeping across dropout ((0.1,0.4) 0.4), weight decay (lle 3, 5e 2]. and hidden size (128, 256). NDT2 does have higher variance, but the main sensitivity is to dropout. We run this sweep and test evaluation in one training stage, Our base NDT2 uses dropout 0.1, hidden size 256, weight decay le 2. In the code, this experiment is configured in exp/arch/tee sspparrhhtune_hpp
A.4.1 Mask Ratio.
We do not widely explore mask ratios due to compute constraints. In pilots throughout, we do not find that decoding is too sensitive to mask ratio (e.g. Fig. 9. but reconstruction quality is hard compare as the inference problem depends on the masking ratio itself. The reasonable effectiveness of high mask ratios is consistent with general observations of low dimensionality and high redundancy in the code, compared to say. language 67
A.5 Additional exploratory experiments
Stitching design Our stitching implementation randomly intializes linear readin and readout linear layer. For ease of implementation, we stitch at the output of the network encoder rather than the output of the decoder (the linear-exponential readout layer comes after per-context stitch layer). In unreported pilots, we find stitching into compressed dimensions (e.g. half of readin channels) to reduce the context-specific parameter count, only including stitching the readin or readout made no significant difference.
Kinematic decoder design There are three straightforward strategies for building a kinematic decoder with NDT2. In keeping with NLB, we could learn a linear probe on representations each timestep, or we could use thin Transformer decoder to allow information from multiple timesteps to aid the
19

---

Sorted
0.64
Velocity R (1)
0.62
Mask Ratio
50%
0.60
25%
75%
0.58
0.56
0.290
0.291 Test NLL (
0.292
Figure 9. Mask ratios over datasets. At test time, the given ratio is held out during unsupervised evaluation.
prediction. We experimented with both and chose the latter for minor gains. For simplicity, we run the Transformer decoder in one forward pass for all timesteps. i.e. there no autoregressive feedback of previous kinematic estimates. We find that cross-attention for decoding queries slightly edges out in-context attention on the higher ends of the pretraining data scales we explore (c.g. 100K), and report with that setting. The primary difference is that tross-attention restricts neural data tokens from attending the kinematic query tokens, while in-context strategies do distinguish the two. For decoding probes, where the decoder is prepared on only few hundred trials, we find it beneficial to mean-pool pool neural data tokens per timestep, and still use in-context attention (linear decoding directly works similarly).
A.6 Single-session breakdown of experimental results
Single-session variability We break open the aggregate results from our primary result figures. The primary takeaways are elaborated in each caption. Overall, we note that single datasets are insufficient make conclusions on design choices given variability in results.
20

---

Archictecture comparisons Velocity decode vs. NLL (Sorted)
dataset adoherty (b.indy-20180k00 02
dataret = cdoterty rtbindy-23166627 01
dataset odoharty rth rt-Indy-2016103 20161005 06
0.0
2.552
0.06
2.22
Volocty
0.62
Velocty
2300
3,473
0.0
2,52
2.00
24222
2/400
243
0,3145
0.341
0.254 0.25
Text
datasat awwhartyyyyaoo.o wttpa//mm./nnoo200102000 03
cataset - coohertyy wtt../ttdnnd2227133..00
Architecture
2.00
NDT2 (Patch 32)
5.8
0573
NDT1
NDT1-Stitch
Velocity 8
Velocity
Transfer Source
0.55
Session
Scratch (Single context)
282
0475
Subject
Task
0.07%
2244
2243
0.24
0.208
0224
Test
Test
Archictecture comparisons Velocity decode vs. NLL (Unsorted)
dataset adsherey ju--ndy-----ooo 20407707_00 ,02
dataset adohertty htloyuuut-nnyy222662777 .01
dataset adoherity_ wtt-iityy____yyy---- (triii/ttt--ddy2211155...33 20161305_0 .6
D41
R
Velaztt R2
R2
0.65
2.55
0.50
230
0.00
043
x
3528
0.14 0.re lest NLL
lest
delasel obblteeiyyuu-hyyy-22000 2016lll00-2..2226._... 03
delmell Buloheerly_U-hnuy--B 201701100
Architecture
NDT2 (Patch 32)
3755
0.63
NDT1
Valccity R²
2.70
+
Valcity
0.00
NDT1-Stitch
Transfer Source
Session
055
Scratch (Single context)
.PH
JEPH
630
Subject
6+
Task
6/3
acees 0.99
Test
Test
Figure 10. We breakout Fig. into individual datasets (points indicate 3 seeds). NDT2 shows consistent improvements over stitching, single session baselines, and rEFH most cases), but the ranking between data sources shifts, particularly for decoding scores.
21

---

0.6
5-1.e + 03 tra
Test
0.6
Test loss
0.5
5 5.5e+ 03 trais Context Tokers
= 2.2e + 04 trials
Test
0.1
10a
102
30
104
Training Steps
Figure 11. Session context tokens improve learning, as measured by unsupervised loss model training curves, though the majority benefit realized with token. We compare this for three data seales, annotated by S, where we scale the number of trials available per session. (The increments were 100%, 25%, and 6.25% of the data). We hypothesized that more data per session would make session tokens less relevant, but the primary effect of increasing convergence appears unchanged these scales.
22

---

Unsup. Transfer Scaling (100 Trial Calibration)
Initial
Middle
Final
5.6 10-1
32
10
Yest NLL
Test NLL
106
61x10-1
4.4 20
6x30
4030
Pretreining
Pretraining
Pretraining trials
Sup. Transfer Scaling (100 Trial Calibration)
045
06-
VolR"
0.85
Vel
Val
Soa
CA-
322
10²
10f
107
107
10"
12"
10
10"
30'
20"
107
Presraining prials
Pretraining trials
Pretraining br als
Pretrain-Posiin Positive Regimes Middle
Initial
Final
62x18
10
166
3
41xts
41
10
10
18xx0
Exts
Torget Figure 12. Individual dataset results for three scaling analyses. We presented the largest middle session in the primary text, here we also show results the first and final sessions in the dataset. The unsupervised and supervised transfer scaling reiterate the previous conclusions: cross-subject and task transfer provides low returns on scaling for unsupervised reconstruction, and decoding results are much more optimistic than unsupervised results. For convergence analysis (Pretrain-posiin regimes), all three trend lines suggest convergence beyond IK trials.
content
"Drget context
Tuning a Decoder (Unsorted) dataset owt.eetty.ttiinnee http:rttnnnm21111111... dataset adtheety wttpeety_ttinnd-20mmm0.. 002
dataset ottheet__ttimm/.2m6000700 on
Ontuned Predrained
Vel O2
Vel 92
Vee 9Z
Untuned Pretrained
Untaned Pretrained
session trais
Target session
Figure 13. The same breakout as above for pretrained decoder tuning. Again, we find that decoder tuning reliably outperforms non-adaptive decoders.
23


Reviews:
[{"review_id": "UTEtFBJp3P", "reviewer": "Reviewer_p9Yw", "summary": "The paper describes a transformer architecture NDT2 for spike train data. It splits a spike train in a set of 2D patches and employs a VIT with a causal auto-encoding approach to regenerate the spike-train patches. The architecture is tested on the RTT arm reaching dataset using multi-unit recordings from the primary motor cortex of a monkey and single-unit recordings from Humans with spinal cord injuries. The performance is tested in terms of negative log-likelihood and arm-velocity decoding R2. The authors report that the decoding improves when using multiple recording sessions.", "questions": "To improve this paper I recommend the authors to make the mathematical/technical description of the NDT2 more precise and unambiguous (showing equations, a pseudo code for the pre-training / evaluation algorithm, and neural network architecture with input/output tensor shapes would be a good way of doing so).\n\nOther than the main criticisms listed above. I also wonder:\n\nIn general, the unit order in a spike train is rather arbitrary, therefore with the patch-based architecture, the group of units forming each patch could be arbitrary. Is the patch always containing units that are neighbors in anatomical space? Is that choice purely arbitrary as given by the default dataset ordering? Would that change anything?", "limitations": "I see no specific limitation to be reported.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The topic of BCI and decoding performance on spike recording is moving quickly and is very important for industry and neuroscience research. The paper aims to find a generic architecture that scales well using transformers which is sounds and likely to be impactful given the growing quantity of neural recordings.\n\nThe idea of using VIT is sound given the success it has had with other modalities. The dataset of random arm reaching is a well-chosen benchmark, although I was more aware of the dataset from the Shenoy lab and not this particular one.\n\n", "weaknesses": "My main criticism is the lack of clarity in the description of the model being chosen. I find in general that the authors have a very good overview of the existing literature but they use generic imprecise terms like \"pretraining\", \"adaptation\", \"asymetric encoder-decoder\", \"single-sorted multi unit recording\", \"online iBCI decoding\" which are vague if the context is not clearly defined. I would have prefered unambiguous mathematical or technical precision in ways that make it reproducible, and very precise on the technical difference from other works.\n\nFor instance here are concrete things that I would have wished to understand in more precise terms:\n\n1) I did not understand what is meant by \"an asymmetric encoder-decoder\" architecture.\n\n2) I did not understand how the NLL is computed. I assume the patches are ordered in as in VIT, suggesting that the causal masking imposes the make a prediction from previous patches in the sequence order. But within a patch, how is a decoder implementing the generation of spikes on a time-step-to-time-step basis? What is the masking during the prediction of the NLL ? Future is always masked? previous patches are masked?\n\n3) The word pretraining is often used without a clear definition of what it means in the context. Are we talking about auto-encoder unsupervised pre-training before fine-tuning the velocity decoder?\n\n4) I did not understand what is meant by \"sorted\" or \"unsorted\". In the context of spike trains, it was not clear to me whether we are talking about spike-sorted spike trains versus raw electrode recordings or spike-sorted spike trains with sorted unit indices, versus spike-sorted spike trains with randomly ordered unit indices.\n\nOther sub-optimal aspects:\n\n5) Not only it would have been better to verbalize the technical details, but there are typical metrics using the arm reaching neural recording (see for instance https://neurallatents.github.io/datasets.html) which have been standardized for comparability across models. it would have been great to be sure that the metrics are consistent and used with the same dataset.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad the reviewer finds the approach sound and likely to be impactful, and agree that the initial submission can be made more technically precise and broadly accessible. We thank the reviewer for their specific feedback and have incorporated all points into the manuscript; we additionally detail them here.\n\n> Lack of clarity in model description and imprecise language.\n\nWe agree greater context is needed to communicate this interdisciplinary work. While descriptions of the breadth of techniques and settings in related work are necessarily imprecise, we have improved our descriptions where most relevant for our work, including the terms highlighted by the reviewer. Moreover, prompted by the reviewer’s suggestion, we add a figure overviewing the training and evaluation process, accompanied by the requested precise description. This description should complement the full code release, which provides another precise path for reproducibility.\n\nTo answer the concrete questions:\n> What is an asymmetric encoder-decoder?\n\nThe original NDT is a pure Transformer encoder, and processes all tokens, masked or not, in all its layers. He 21 introduces the term “asymmetric encoder-decoder.” This is a two stage architecture where only unmasked tokens are encoded in the bulk initial set of layers, and masked tokens are introduced and decoded in a final lightweight module with few layers [1]. The fact that the second stage has relatively few layers makes the model asymmetric. We have expanded this description in the text.\n\n> NLL computation and causality\n\nAs in NDT1 and other neural data models, NDT2 computes NLL with a Poisson likelihood model. Specifically, predicted firing rates for each patch are compared against the binned spike count observed in the patch. For example, with a patch size of 16 channels and 20ms, the model predicts for each patch a 16-D Poisson rate vector to compare with the corresponding 16-D spike count that occurred in this 20ms. The model always uses one time bin per patch, and cannot reason about causality below the binning resolution.\n\nWith regards to patch-wise causality: NDT2 attends in both space and time. Temporal attention is causal, but spatial attention is unrestricted. This contrasts with the Dosovitskiy Vision Transformer, which doesn’t restrict attention at all, and standard language models or PixelCNN, which is causal over a flat data stream. Our evaluation NLL is causal in that its computation is as done in training, where we randomly mask the input and report the average NLL on masked inputs. The data are large enough that there’s little NLL variability from masking.\n\nThese are subtle points so we thank the reviewer for highlighting this and clarify the text accordingly.\n\n> Defining Pretraining \n\nYes, our pretraining is the unsupervised masked autoencoding process, as in He 21 [1]. \n\n> Sorted vs unsorted\n\nWe briefly explained the sorting process on l68, but elaborate here. Microelectrode arrays record extracellular voltage signals and strong fluxes in this signal are recorded as spiking activity. Unsorted, multi-unit spiking activity is simply this activity grouped at the electrode level, so the data will have as many dimensions as there are electrodes. The sorting process further clusters this activity based on the waveform of the spikes, so sorted activity will have as many dimensions as identified spikes (which may vary across datasets).\n\nFor example, a dataset with 128 channels of recording, and 2 identified neurons per electrode, would produce 256-D spike sorted data but only 128-D multi-unit activity. We do not randomize the unit indices from how they are given in the data, but these indices do not have inherent meaning.\n\n> Controlled comparison against the Neural Latents Benchmark.\n\nWhile the NLB is a boon for comparing neural data models, it emphasizes the evaluation of model latents over expressive decoding. As such, the NLB restricts a private set of data on which a linear probe of behavioral decoding is separately evaluated, which precludes both the evaluation of our nonlinear decoders and a controlled analysis of the impact of decoder training data. The NLB also reports spike reconstruction on _held-out_ neurons, but these held-out evaluations are not the endpoint of our brain computer interface (BCI) evaluation. Our work aims to study how different types of BCI data can scale pretrained decoder performance. Therefore, we opted for the standard unsupervised NLL reported in scaling papers in machine learning.\n\nWe note that our discussion mentions our negative result on the NLB; we have clarified how the NLB setting is distinct from ours, but the negative result certainly warrants revisiting. NLB aside, we believe that a benchmark specific for BCI decoding would advance the field.\n\n> Question: Arbitrary ordering within each patch.\n\nIndeed, the spatial ordering is arbitrary in the model design, and the default of ordering them according to the datasets will place anatomically neighboring units in close patches. While we don’t believe this to affect modeling within a dataset, in early studies we did see that permuting spatial order across datasets harmed performance, suggesting that consistent patching strategies are important. However, in our supplement, A.3.2 shows that we find small patch sizes fail to decode well. We are excited for future improvements in model design that address this arbitrariness.\n\nWe thank the reviewer for the opportunity to improve the clarity and reproducibility of the paper and hope that they would reconsider their assessment if our responses resolve their concerns.\n\n[1] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. Dec 2021"}, {"review_id": "XsNegTAr85", "reviewer": "Reviewer_yEFS", "summary": "The authors present a method leveraging the original Neural Data Transformer model to span additional contexts and work across different trials and subjects. The essential claim is that pretraining using various existing datasets can improve the performance of individual iBCI decoding and control. \n\nUpdate: I have updated from score from a 6 to a 7 in light of the additional work the authors provided.", "questions": "i.\tLines 42-45 describe NDT1 and some of NDT2 but the explicit distinction between the two is not clear. \nii.\tLines 59-61 need citations. \niii.\tLines 69-71 need citations. \niv.\tLine 98: additional information on the NLB task, metrics, and top model would be useful here. \nv.\tLine 103-104 is not clear what is meant here. \nvi.\tLine 161. Is this data in the supplementary somewhere?\nvii.\tLine 220: Figure 10 is not in the main text", "limitations": "The authors discuss limitations clearly in their conclusion.\n\nCould the authors comment on line 202-203 about the lack of task structure impacting their model?", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality: While this method expands on the original NDT, it does so in an original way, using pretraining on additional datasets to include spatiotemporal attention mechanisms, learned context embeddings from diverse datasets, and asymmetric encoding and decoding methods.  \n\nQuality: The original NDT is well-studied and robust; this extension is well-motivated and uses techniques from the broader machine learning community to improve generalizability and performance in real neural datasets.\n\nClarity: The paper is written well but assumes details or familiarity with iBCI datasets that the NeurIPS audience may not have, or other acronyms (e.g. BERT is never defined). Much of the text is similarly dense with prerequisite knowledge. \n\nSignificance: This work has the potential to impact much of the BCI field, and future clinical applications. The dominant contribution appears to be using pretraining, which is generally known in the ML field but less so in BCI work, thus making this more significant for future endeavors. \n", "weaknesses": "The authors mention this method's use in offline decoding but it is unclear how useful this would be for practical iBCI usage. \n\nMuch of the specifics are unexplained; for example, the use of context tokens, (section starting on line 212) is described without mentioned what they are or why the authors chose to use them. \n\nWhile it is clearly stated that monkey pretraining data didn't improve performance of human data, there are open questions then as to what type of pretraining data is actually useful, and how to predetermine which data is useful or not. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are excited the reviewer values the central contribution of exploring the value of pretraining in BCI. We are grateful for the specific feedback points on clarity as noted in the questions, and revised our manuscript accordingly, and address them below.\n\n> Clarity: Assumed pre-requisite knowledge\n\nWe have revisited our main text and added exposition needed to introduce the previously assumed ideas from both ML and neuroscience.\n\n### Weaknesses\n> Applications to practical iBCI usage\n\nWe agree that it is critical to address how improvements in offline decoding translate to practical usage. We note that it is a substantial and active challenge to bridge this gap both technically and simply logistically, as online iBCI experiments require significant resources. See, for example, dedicated studies on deep network deployment in iBCI experiments [1, 2]. However, we are excited to add pilot online experiments (see global response) that show core benefits of pretraining and tuning: NDT2 allows decoders to work on new days, and for tuning to improve this performance.\n\n> Lack of detail on the specifics e.g. the use of context tokens.\n\nWe thank the reviewer for the specific highlight, and as mentioned have revisited the methods and added detail on the specific motivation and instantiation of method details. The context tokens are learned embeddings of metadata IDs, e.g. we learn vectors for subject 1, session 3.\n\n> Open questions on determining useful data\n\nIndeed, a major contemporary challenge of pretraining, even outside of neuroscience, is identifying the right datasets to train on. In general this is currently an empirical question that is difficult to forecast a priori. We dedicated our experiments to mapping data utility at a coarse level, identifying the relative value in the categories of session, subject, and task pretraining. We hope these results provide rough guidance on potential data utility, but also look forward to future work that provides more specific recommendations.\n\n\n### Questions\n\ni. The differences in NDT1 and NDT2 are specified in the Contributions section (l48). To reiterate: spatiotemporal attention, learned context embeddings, and asymmetric encode-decode.\n\niv. The NLB benchmark uses similar likelihood and decoding metrics as we do. The current top model on the O’Doherty dataset is a state space model (S5) but standard Transformers and RNNs are competitive.\n\nv. We have revised the full text (~l106-114) around this spacetime attention as follows: \n```\n…Yet factorization can impair performance [3] and requires padding in both space and time when training over heterogeneous data. However, full neuron-wise spatial attention is too costly. We thus adopt $K$-neuron patches akin to pixel patches in ViTs [4], padding data to the nearest multiple of $K$. The patch is embedded by concatenating projections of the spike counts in the patch. In pilot experiments, we find comparable efficiencies between this patching strategy and factorized attention. We opt for the former design as it enables easy adoption of the asymmetric encoder-decoder proposed in [5]....\n```\nvi. We have rewritten to clarify: model calibration is always through fine-tuning; only single-session models (which have no broader pretraining) “pretrain” through an unsupervised step on the data in that session.\nviii. Task structure impact:Thank you, we have added clarification that stitching typically is used after an initialization procedure called principal components regression (PCR). PCR projects neural data grouped according to experimental metadata (like reaching direction), intuitively so that multiple days initialize in similar subspaces. These methods are detailed in Pandarinath et al [6]. Private correspondence suggests PCR initialization is important for more robust stitching performance.\n\nWe thank the reviewer for specific feedback!\n\n[1] Deo et al. Translating deep learning to neuroprosthetic control. In bioRxiv 2023.\n\n[2] Costello et al. Balancing memorization and generalization in rnns for high performance brain-machine interfaces. In bioRxiv 2023.\n\n[3]  Arnab et al. Vivit: A video vision transformer. In CVPR 2021\n\n[4] Dosovitskiy et al. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR 2021.\n\n[5] He et al. Masked autoencoders are scalable vision learners. In CVPR 2021.\n\n[6] Pandarinath et al. Inferring single-trial neural population dynamics using sequential auto-encoders. In Nature Methods 2018. URL https://www.nature.com/articles/s41592-018-0109-9."}, {"review_id": "S8MbeBd1y0", "reviewer": "Reviewer_cSw7", "summary": "The paper introduces the NDT2 model for processing neural spiking activity data in intracortical brain-computer interfaces (iBCIs). The authors demonstrate the application and transferability of the NDT2 model across multiple experimental contexts through large-scale unsupervised pretraining. The model can aggregate neural activity data from diverse contexts and adapt rapidly to new environments. The model improves scalability in heterogeneous environments by incorporating spatiotemporal attention, learned context embeddings, and an asymmetric encode-decode structure. The paper also presents the datasets used and the evaluation results of the model.", "questions": "1. Did the authors optimize the computational efficiency of the model to address the computational cost issue caused by the full attention mechanism?\n2. Does the dataset used for evaluation have sufficient diversity to represent different real-world scenarios?\n", "limitations": "See weaknesses", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Large-scale unsupervised pretraining: The NDT2 model utilizes large-scale unsupervised pretraining, enabling the aggregation of neural spiking activity data from multiple experimental contexts, overcoming the limitations of traditional models that are designed for a single context.\n2. Rapid adaptation to new environments: The pretrained NDT2 model exhibits fast adaptation to new experimental environments, providing convenience for the deployment of neural brain-computer interfaces.\n3. High-performance decoding: The model demonstrates excellent performance in decoding unstructured monkey reaching behavior and human iBCI cursor intent.\n", "weaknesses": "1. Computational efficiency: The full attention mechanism may increase computational costs, potentially affecting the computational efficiency in practical applications. \n2. The method and background are not described in sufficient detail.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s generous summary. To address reviewer cSw7 (along with the other reviewers), we have updated the manuscript with increased detail in the method and more exposition in the background (see global response) and have provided new versions of the figures to improve the clarity of the methods and the results (see rebuttal document). To address the reviewer’s specific concerns:\n\n*Computational efficiency*:\nIndeed, full attention is more expensive than other variants such as attention that has separate space and time operations. We updated our text to clarify that we compared this computationally more efficient variant in early experiments and didn’t see severe costs when using full attention. We ultimately chose to use standard full attention as we wanted a stable proof of concept of pretraining using a standard architecture. We also note that full attention computation cost does not preclude real time use of the model (see global response attachment). Nonetheless, the reviewer’s concern is important and practical. We hope future works can explore effective, efficient alternatives to full attention.\n \n*Evaluation dataset diversity*:\nThe evaluation dataset is a self paced reaching task, which represents relatively naturalistic and more challenging behavior relative to the trialized, cued behavior in most other open source monkey datasets. Thus practically, we chose the most complex available open source evaluation dataset. Such data is regularly used to calibrate BCI decoders for translational control in lab environments. Identifying even more general data that applies in varied real world settings is an active research area, and we hope NDT2 and models like it can subsume this diversity once it is identified.\n\n\nWe thank the reviewer for highlighting these practical concerns, and hope our responses and modifications resolve these concerns.\n"}, {"review_id": "7Q54bgGPoz", "reviewer": "Reviewer_mmUN", "summary": "This paper introduces the 'Neural Data Transformer 2' (NDT2), which can seamlessly integrate data from neural intracortical recordings across sessions, tasks, and subjects, by applying attention across space and time instead of having to align potentially mismatched input channels. The authors proceed to demonstrate that such cross-session/task/subject pretraining can improve decoding performance in target sessions with only little available finetuning/training data. This has clear parallels to the language literature, where large pretrained transformers have proven immensely powerful for generalizing to new contexts in zero- or few-shot settings. The authors speculate that systems neuroscience could be entering a similar era of large pretrained models, with NDT2 serving as an early proof-of-principle of the necessary cross-context generalization.", "questions": "L136: I don't understand why the authors decide to use 'trials' to refer to seconds in the RTT task since, as the authors mention, this does not correspond to actual trials. If everything is measured in units of seconds, which is a perfectly natural unit to use, why not just say seconds and minutes throughout instead of arbitrarily redefining a trial to mean a second, and then sometimes saying minutes/seconds and sometimes trials? If the authors want to provide a reference point for what happens behaviourally in a given timeframe, they could compute the median trial time (i.e. time between targets) in the RTT and report that when describing the data. For the across-task comparison, they could quantify the amount of pretraining data in other tasks in seconds as well (which is probably a fairer comparison than trials anyways).\n\nFigure 2: it would be useful with a schematic illustrating the pretraining paradigm on the RTT data. It took me a while to figure out what the different baselines and symbols mean.\n\nL186-187 state that \"all models identically receive context tokens\", but does not specify what information is contained in these tokens. Figure 3 suggests that it is some kind of session information, but even after the section specifically about context tokens, it  was unclear what information is contained in these tokens. The authors may want to clarify this somewhere.\n\nL190-192: it is not clear that the NDT results are directly comparable to the rEFH results if the rEFH data uses contiguous training/test data and NDT does not, since any systematic drifts in electrodes/neural activity/behaviour/motivation would make the contiguous case substantially harder by requiring out-of-distribution generalization.\n\nFigure 2: Is the 'multi-task' data from the Maze task? or does it also include other tasks? It would be worth clarifying this in the legend.\n\nFigures 2 and 4: is it the same models being used for the supervised ($R^2$) and unsupervised (NLL) metrics?\n\nL237-238 \"This transfer is analogous, for example, to retained overlap in the first K PCs across two sessions, but generalizes nonlinearly and to many more sessions.\" What does this mean?\n\nFigure 4 (very minor): I would put the panel labels on the left instead of the right of each subplot.\n\nFigure 4: It's somewhat confusing that panels B and C use different y label notation - especially if they refer to the same quantity, which I'm not 100% sure whether they do given the inconsistent notation and minimal description in the legend.\n\nFigure 4: It's not entirely clear what's going on in Panel C, and it may be worth clarifying in the legend. Is this panel varying the amount of within-session 'fine-tuning' data while keeping the pretraining data constant? If so, how much data was used for pre-training in each case?\n\nFigure 2 & 4: I was surprised that pretraining on different sessions/subjects/tasks is done in an 'exclusive' manner, since a major motivation for foundation models is being able to use _more_ pretraining data. It would be interesting to see how performance changes after pretraining on either (i) just across-session data, (ii) the same amount of across-session data, but now adding across-subject data, and (iii) now also adding across-task data. Does the ability to incorporate more data benefit the model, even if the added data is in itself less 'useful'?\n\nFigure 5: The legend feels a bit too minimal. It took a few re-readings of the main text to understand what the different lines mean.\n\nL312 \"even if the ideal scenario of cross-species transfer seems unlikely\": this seems like a strong statement since transfer was attempted only with a single pre-training dataset. Cross-species transfer may still work if one were to pre-train on _many_ species to capture whatever is invariant across species rather than perhaps 'overfitting' the pretraining to monkey M1?\n\nThe above suggestions involve _adding_ many things and _removing_ nothing, which is not always easy given page limits. If space is a major limitation, it could be worthwhile to move Figure 3 and the associated text to the appendix in order to better explain the other analyses, which seem more critical to the primary message of the paper.\n", "limitations": "The authors have adequately addressed the limitations of their work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "Systems neuroscientists are collecting increasingly large datasets across a wealth of brain regions, tasks, and species, and developing methods for integrating these resources remains an immensely important challenge for the field. NDT2 represents an exciting step in this direction by showing how the tools used in the vision/language literature can also drive cross-context generalization and 'few-shot' learning in systems neuroscience. As more data becomes publicly available and we improve our computational infrastructure, these ideas may help pave the way towards neural 'foundation models' that could be useful both for BCI applications and for fundamental neuroscience research.\n\nThe authors demonstrate the utility of NDT2 through a range of interesting analyses of a previously published primate reaching dataset, and they finish with an exciting new analysis of a human BCI dataset that demonstrates the potential utility of cross-context pretraining in a clinical setting.", "weaknesses": "The major limitation of the present study is the relatively small amount of data being used, limited to just a few monkey reaching tasks and a single human dataset. This is of course a large amount of data in the context of systems neuroscience, but it remains very small in the context of generative pre-training, and it may result in some amount of overfitting on these specific contexts (illustrated e.g. by the drop in performance on the human data when pretraining on the monkey data). One would would imagine that generative pretraining for systems neuroscience becomes much more useful when aggregating data across _many_ experiments, e.g. by scraping CRCNS. This is of course way out of scope of the present paper, which represents an exciting starting point, proof of principle, and baseline for establishing the necessary methodology as discussed above.\n\nThe presentation was not always super clear. It was at times difficult to understand what was being plotted in the figures, how the model was constructed, and how the analyses were performed (see below for a few examples). I would particularly encourage the authors to expand their figure legends to help the reader understand what is being plotted.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are thankful for the reviewer’s enthusiastic and detailed feedback.\n\n### Weaknesses\nWe agree that an even wider pretraining is promising, but appreciate the reviewer’s recognition that NDT2 is already a substantial effort within systems neuroscience and serves as a necessary step before exploring the potential rewards of pretraining across multiple brain areas. With regards to volume, we’d like to highlight that our human models were pretrained with human data spanning over 6 years, with evaluation using data spanning over a month. This detail, along with changes for the reviewer’s other specific clarity concerns, have been added to the manuscript. We believe the added figure 2 overviewing model preparation should address the specific feedback provided, but we also provide more notes below.\n\n### Questions\nL136: Good point; we had adopted the terminology used in the Neural Latents Benchmark, where trials were identically defined for their O’Doherty dataset. In our study, trials are used over seconds as not all time is equally behaviorally interesting; for example, in Nir-Even Chen’s delayed reaching dataset, where we might average a reach every 8s, O’Doherty’s dataset for monkey Indy averages ~1s for a reach. However, we agree this choice requires care, particularly in the context of comparing pretraining efforts in the future. While even complex natural behavior could potentially be segmented, using wall-time is simple and therefore likely the most straightforward quantify pretraining scale. We revise the manuscript to use time in broad characterizations of the dataset but trials in specific experimental comparison, and add this discussion to the appendix.\n\nFigure 2, pretraining flow, data source, model metrics: Thank you for the suggestion, we have added a schematic in the global response. In addition to the precise description requested by reviewer p9Yw, we hope the training protocols are now clear.\n\nL186, context tokens: Yes, we clarify these context tokens are simply encodings of metadata IDs, the model learns a single vector per session, subject, or task.\n\nL190: Comparison to rEFH. Indeed, we ran a simple supplementary comparison to show that NDT2 decoding score goes from 0.54 to 0.52 when mirroring the Makin data splits. We add this in a footnote. We note that Makin et al appeared to tune early stopping of their rEFH model with test data, which may advantage their model.\n\nL237-8: We agree this statement is confusing and remove it.\n\nFigure 4, Panel C: Correct, we are varying the within-session tuning data. We directly take the largest pretrained model from multi-session and multi-task data, and annotate a line from the same points in panel B to make this clear (see global response).\n\nF2 & 4, Exclusive pretraining: Indeed, we ultimately do recommend pretraining broadly on the available data. However, in practical deployment scenarios, not all types of data will be available, e.g. multisession is not available early in subject’s use. Moreover, due to data privacy concerns or practical limits, we may not generally be able to deploy multi-subject data. Thus we wanted to emphasize the relative contribution of different types of data on a small volume of target data (with the implicit understanding that already well characterized data would benefit minorly, e.g. as in 2C). The suggested inclusive, marginal return on pretraining should provide similar insight if interpreted as such, but with smaller differences between curves. \n\nL312, cross species transfer: While cross-species over many species is an exciting proposal, it might be practically challenging as the volume of relevant human or monkey data likely surpasses those of other species, except perhaps rodents. We revise our statement to make it a practical limitation rather than a theoretical one.\n\n**Overall**, we are glad the reviewer provided such rich feedback and believe the manuscript has been greatly enriched with reviewer discussion.\n"}]

Author rebuttal:
We thank the reviewers for their time and specific feedback. We are pleased that the reviewers recognize that neural data pretraining is likely to be impactful (p9Yw, yEFS, mmUN), that the NDT2 approach is well-motivated yet original (p9Yw, yEFS), and that all reviewers appreciate the results achieved.

*Clarity*:
The primary criticism of all reviewers was clarity of presentation. We agree that important contextual details were omitted and that in other instances the material that was presented was difficult to follow. To address this, we have created a schematic to illustrate how different data sources were organized, how the model was prepared (p9Yw, mmUN), and also precisely describe this flow in the associated text (p9Yw). We have also updated all other figures. The first 3 figures in our attachment to this global response contains some of these updates:
- Fig 1 revises the manuscripts Fig 1, now illustrating the work’s conceptual motivation in examining the benefits of broader pretraining.
- Fig 2 is the new schematic overviewing how data was organized for model training.
- Fig 3A is a revision of our decoding figure with an elaborated legend and caption.

Further, by pooling the specific feedback provided by the reviewers, we have made substantial updates throughout the text. For example: 
1. In the section on related work, we more fully describe the cited works that relate to our method (e.g. BERT, asymmetric encoder-decoder). 
2. In our approach, we add concrete details to our tokenization scheme and model preparation.
3. We revise text in the description and presentation of experiments to improve clarity in tandem with revised figures.

*Open questions on data*: 
Reviewers yEFS, cSw7, and mmUN raised a concern that NDT2 leaves unanswered questions about the necessity and relevance of different types of data (i.e. sessions vs. subjects vs. tasks). We had designed our experiments to provide a coarse lens to compare the value of these three major types of data, but we recognize the importance of the question and hope future works can build on our findings. Cross-species and cortical area transfer remains to be investigated, but will be challenging given the limited volume of open system neuroscience data. Nonetheless, in the words of mmUN: NDT2 represents “exciting starting point, proof of principle, and baseline for establishing the necessary methodology.”

*Practical benefits*:
Reviewers yEFS and cSw7 asked for comments on the practical benefits of the NDT2 approach. We are also very interested in this topic, and are happy to add Figure 3B in our attachment on initial human experiments, where different capabilities of NDT2 are compared. Specifically, we evaluated how pretrained models, enabled by NDT2, perform without any same-day data (zero-shot), and with both unsupervised and supervised tuning data from the evaluation day. There are acknowledged challenges in the field with offline to online translation, and while decoding performance with NDT2 after supervised tuning did not exceed current linear methods, NDT2 for the first time enabled unsupervised recalibration and also allowed reasonable zero-shot decoding performance.

We also provide individual responses to address reviewer-specific feedback.


Meta-review: The paper proposes a Transformer-based architecture for modeling neural spike trains and show that pretraining on a variety of datasets improves accuracy. The reviewers appreciate the overall contributions of the paper. Their main criticism has been clarity, which the authors addressed in their rebuttal. Overall the reviewers are supportive of the paper and consider it a worthwhile contribution.
Decision: Accept (poster)