Title: Neural Data Transformer 2: Multi-context Pretraining for Neural Spiking Activity
Venue: NeurIPS 2023
Accepted: yes
Keywords: Pretraining

Scaling Laws

Neuroscience

Brain-computer interfaces
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

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

1 Introduction

Intracortical neural spiking activity contains rich statistical structure reflecting the processing it subserves. For example, motor cortical activity during reaching is characterized with low-dimensional dynamical models [1, 2], and these models can predict behavior under external perturbation and provides an interpretive lens for motor learning [3-5]. However, new 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 [10, 11]. 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 we focus on one primary use case: neuroprosthetics powered by intracortical brain computer interfaces (iBCIs). With electrical recordings of just dozens to hundreds of channels of neuronal population spiking activity, today's iBCIs can relate this observed neural activity to behavioral intent, achieving impressive milestones such as high-speed speech decoding [12] and

high degree-of-freedom control of robotic arms [13]. Even so, these iBCIs currently require arduous supervised calibration in which neural activity on that day is mapped to behavioral intent. At best, cutting-edge decoders have included training data from across several days, producing thousands of trials, which is still modest by deep learning standards [12]. Single-session models still dominate the Neural Latents Benchmarks (NLB), a primary representation learning benchmark for spiking activity [14]. 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 is deep learning's seemingly robust recipe for leveraging heterogeneous data across domains: a generic model backbone (e.g. a Transformer [15]), unsupervised pretraining over broad data, and lightweight adaptation for a 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) [16], 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.

Contributions : Here, we contribute NDT2, a Transformer that pretrains over broad data sources of motor cortical spiking activity. NDT2 modifies NDT1 to improve scaling across heterogeneous contexts in three ways (Section 3): spatiotemporal attention, learned context embeddings, and asymmetric encode-decode [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 on offline evaluation on motor applications, demonstrating NDT2's value in decoding unstructured monkey reaching and human iBCI cursor intent. We also show proof-of-principle real-time cursor control using NDT2 in a human with an iBCI.

2 Related Work

Unsupervised neural data pretraining. Unsupervised pretraining'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 4 sampled in Table 1, 3 use masked autoencoding as a pretraining objective (EEG uses contrastive learning), 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 is data instability [18]. The high spatial resolution of iBCI microelectrode arrays enables recording of individual neurons and provides the signal needed for high-performance rehabilitation applications, but this fine resolution also causes high sensitivity to shifts in recording conditions [11, 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).

Data aggregation for iBCI. 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 a method called stitching [27]. For context, the extracellular spiking signals recorded on microelectrode arrays are sometimes 'sorted' by the shape of their electrical

A.

B.

ContextEmbed

Session 1Subject2

8

Unmasked

SessionSubject

Task

Encoder

Moredatarelevance

Neuron

Masked

Neural Decoder

Behavior Decoder

Time

Poisson NLL

waveforms, with different shapes attributed to spikes from distinct 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 to BCI and broader neuroscience applications over half a year [28-30, 11, 31]. However, even linear layers can add thousands of parameters, which risks overfitting in clinical iBCI data that comprise only dozens of trials.

Alternatively, many iBCI systems simply forgo spike sorting, which appears to have a minor performance impact [10, 18]. Then, multi-session data have consistent dimensions and can feed directly into a single model [10, 32, 33] (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, or dropout, emphasizing robustness as a design goal.

Relative to existing efforts, NDT2 aims to aggregate a larger volume of data by considering other subjects and motor tasks (Fig. 1B). 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.

Domain-adaptive vs. domain-robust decoding. Given little training data, special approaches help enable BCI decoder use in 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 a novel context with further data collection, which is reasonable especially if only unsupervised neural data are required. Several approaches align data from novel contexts by learning an input mapping that minimizes distributional distance between novel encodings and pretraining encodings explicitly [11, 37, 38]. NDT2 also allows unsupervised adaptation, simply by fine-tuning without additional objectives.

3 Approach

3.1 Designing Transformers for unsupervised scaling on neural data

Transformers prepared with masked autoencoding (MAE) [20] are a competitive model for representation 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 is reviewed in Fig. 1A.

The primary design choice is then data tokenization, i.e. how to decompose the full iBCI spatiotemporal 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 [1]. NDT1 [16] 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 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-wise spatial attention has prohibitive computational cost. We compromise and group K neurons to a token (padding if needed), akin to pixel patches in ViTs [41]. 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 50100Hz [42]; we adopt 50Hz (20ms bins) and consider temporal context up to 2.5s (i.e. 125 tokens in time). Given a total context budget of 2K tokens (limited by GPU memory), this leaves a few dozen tokens for spatial processing. It is currently common to record from 100-200 electrodes with an interelectrode spacing of 400 µm (Blackrock Utah Arrays) [29, 43], so 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 [44], which will warrant new spatial strategies.

NDT2 also receives learned context tokens. These tokens use known metadata to allow cheap specialization, 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.1. All datasets contains single-unit (sorted) or multiunit (unsorted) spiking activity recorded from either monkey or human primary motor cortex (M1) during motor tasks. In particular, we focus evaluation on a publicly available monkey dataset, where the subjects performed self-paced reaching to random targets generated on a 2D screen (Random Target Task, RTT) [47], and unpublished human clinical BCI datasets. RTT 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 it features long, near-hourly sessions of a challenging task where decoding accuracy increases with data [48]. 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 1s trials in keeping with NLB [14].

We also study M1 activity in 3 human participants with spinal cord injury (P2-P4). These participants have limited preserved motor function but can modulate neural activity in M1 using attempted movements of their own limbs. This activity can be decoded to control high degree-of-freedom behavior [13]; 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.gov (ID: NCT01894802) and informed consent was obtained before any experimental procedures were conducted. Details on the implants and clinical trial are described in [43, 13], and similar task data are described in [49].

4 Results

We demonstrate the three requirements of a pretrained spiking neural data model for a BCI: a multicontext capable architecture (Section 4.1), 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.

Obstacle-guidedreach

Weight

ExamplePretraining:20Ktrials

Other

Initialization

7%

26%

Session

Self-paced

Random

Human CursorBCI

13%

reach

(Sec3.4)

50%

A.

B.Per-Eval Tune + Probe:300 trials

Pretrained

UnsupervisedPretraining

Unsupervised Tuning

Supervised Probe

Learning?

Encoder

Encoder

Encoder

No

Yes

Neural

Behavior

Neural

Behavior

Neural

Behavior

C

Model preparation and evaluation. Most experiments used a 6-layer encoder ( ∼ 3M parameters). NDT2 adds a 2-layer decoder (0.7M parameters) over NDT1 [16]; we ran controls to ensure this extra capacity does not confer benefits to comparison models. To ensure that our models were not bottlenecked by compute or 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, as 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. 2B 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-context models. For unsupervised evaluation, we use the Poisson negative log-likelihood (NLL) objective, which measures reconstruction performance of randomly masked bins of test trials (Fig. 2C middle). For supervised evaluation, we report the R 2 of decoded kinematics, which for these experiments are 2D velocity of the reaching effector (Fig. 2C 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 a 2-layer Transformer probe (matching the reconstruction decoder). Here we provide the models with 5 minutes of data (300 training trials) for each evaluation session . This quantity is a good litmus test for transfer as it is sufficient to fit reasonable single-session models but remains practical for calibration. A 10% test split is used in 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 (Multi-Session), the sessions from the other monkey (Multi-Subject), or from non-RTT datasets entirely (Multi-Task, see Fig. 2).

A.Sorted

B.Unsorted

0.7

DataSource

Single-session

0.6

NDT2(Ours)

(Scratch)

Multi-Session

2

2

Velocity

R

rEFH (Makin 18)

Velocity

R

0.6

Multi-Subject

Multi-Task

rEFH (Makin 18)

NDT1

0.5

0.5

NDT1-Stitch

0.290

0.293

0.295

0.464

0.472

0.480

TestNLL(←)

Test NLL（)

Prior work in multi-session aggregation either use stitching layers or directly train on multi-day data with consistent unit count. Thus, we use NDT1 with stitching as a 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 in Fig. 3. For context, we show single-session performance achieved by NDT1 and NDT2, and the decoding performance of 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. [48]. 2 Single session performance for NDT1 and NDT2 is below this baseline. 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, in sorted data shown in Fig. 3A, all pretrained NDT2 models outperform rEFH and singlesession baselines, both in NLL and kinematic decoding. Surprisingly, multi-subject data work as well as multi-session 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-session 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 R 2 ) in these experiments (Fig. 3B, 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, we use 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

A

B

Pretraining data

C

Targetdata

0.62

Performance

→

0.62

Performance

N

0.6

77N

Test

0.61

0.61

0.4

Performance

0.60

0.60

>

Pretrainingdata

102

103

104

105

102

103

104

105

102

103

Pretraining trials

Pretraining trials

Targetcontexttrials

Single-session (Scratch)

Cross-Session

Cross-Subject

Cross-Task

efforts, we perform three analyses to coarsely estimate the data affinity [52] of the three different context classes (cross-session, cross-subject, and cross-task). Previously these relationships have been grounded in shared linear subspaces [6-8]; 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 in a novel context. We do this for both supervised (Velocity R 2 ) and unsupervised (NLL) metrics in Fig. 4A and B respectively. To contextualize performance of cross-context scaling, we measure in-distribution scaling of intra-session data (Scratch). For example, the largest cross-session model tuned with 100 trials achieves a NLL similar to 1K trials of intra-session data (Fig. 4B, Cross-session vs Scratch). This shows that cross-session data can capture a practically long tail of single-session neural variance (experiments rarely exceed 1K trials), with nonsaturated benefits in pretraining dataset size. Alternately, cross-session pretraining allows decoding performance similar to the largest single-session model performance (Fig. 4A, Cross-session vs Scratch), but this scaling is starting to saturate.

Shallower slopes for cross-subject and cross-task pretraining on both metrics indicate poorer transfer to evaluation data. In the unsupervised case (Fig. 4B), cross-subject and cross-task transfer never exceed the NLL achieved with 400 single-session trials. 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 [53]). For example, the decode R 2 achieved by the largest task-pretrained model compares with in-session models at 800 trials; the same model'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 represents about 50m of data collection in the monkey datasets, and coincidentally is 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 is ineffective at modeling the long tail of neural variance. Note that returns on supervised evaluation are likely similar or better based on Fig. 4A/B; we explore a related idea in Section 4.3.

Overall, the returns on using pretrained BCI models depends on the use case. If we are interested in best explaining neural variance, pretraining alone underperforms a moderately large in-day data collection effort (single-session model achieves lowest NLL in Fig. 4B). However, we do not see

A

P4 Implant Day

Day

CursorControlPilots

0.8-

0

61 →63→65

70

B

CO-ShotOLE>

0.7

Pretrained (O-Shot)

O-ShotNDT2

2

R

ocity

0.6-

C

5

X

OLE

Session

Broad

0.5

↓

4

>

0.4

Supervised tune

Time

0.3

Scratch

3

Unsupervisedtune

Scratch

0.2

(100SupervisedTrials)

Re

2

102

103

1

Targetsessiontrials

O-Shot

UnsupervisedSupervised

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 a decoder, and finally measure the decoding performance in a novel target session in Fig. 5A. We find that given either supervised or unsupervised tuning in our target session, beyond smoothly improving over the 0-Shot Pretrained model's performance, achieves decoding performance on par 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 are unlabeled (100 Supervised Trials) and in the optimistic case when >1K trials of supervised data are available (Scratch). As expected, though, gains against the Scratch models are 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 [11, 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 (e.g. 40) and intent labels are much noisier than movement recordings (intent labeling is described in Section A.1). 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 a 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 a minor performance impact (row 1 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; we use

linear decoding performance as a single-session reference linear decoding performance. Overall, the plentiful cross-session data likely occludes further gains from cross-subject and cross-task pretraining, but note that there is no negative transfer.

Next, given reports of monkey to human transfer [55], 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 to humans.

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 ( ∼ 2 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-specific 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 be used directly (0-shot) or updated with calibration data on the test day (Fig. 5B). 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.2), 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-alignment priors, as in [11, 35, 33], showing that pretraining and fine-tuning [20] may be a viable paradigm for closed loop motor BCIs.

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 to 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.

5 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 BCI, 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, is 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 a practical path forward: beyond explicit calibration phases, strategies for supervising BCIs are diverse, and can derive from user feedback [60], neural error signals [61, 62], or task-based estimation of user intent [63, 58, 64]. The ambition of pretrained BCI models, with broad paired coverage of the neural-behavior domain will be challenging given 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.

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 heldout 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. 4C), as in 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 spacetime 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 reconstruction 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 transfer [11, 38, 33, 35], evaluating more complex behavior decoding is a 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, 18]. 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-Carracedo 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 Karpowicz 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 R01NS121079 and UH3NS107714. 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, or usefulness of any information, apparatus, product, or process

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

Saurabh Vyas, Matthew D Golub, David Sussillo, and Krishna V Shenoy. Computation through neural population dynamics. Annual review of neuroscience , 43:249-275, 2020.

Ege Altan, Sara A. Solla, Lee E. Miller, and Eric J. Perreault. Estimating the dimensionality of the manifold underlying multi-electrode neural recordings. PLOS Computational Biology , 17(11):1-23, 11 2021. doi: 10.1371/journal.pcbi.1008591. URL https://doi.org/10.1371/journal.pcbi.1008591 .

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 a dynamical mechanism for robust computation. bioRxiv , pages 2022-12, 2022.

Patrick T Sadtler, Kristin M Quick, Matthew D Golub, Steven M Chase, Stephen I Ryu, Elizabeth C Tyler-Kabara, Byron M Yu, and Aaron P Batista. Neural constraints on learning. Nature , 512(7515): 423-426, 2014.

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 , 97(5):1177-1186, 2018.

Juan A. Gallego, Matthew G. Perich, Stephanie N. Naufel, Christian Ethier, Sara A. Solla, and Lee E. Miller. Cortical population activity within a preserved neural manifold underlies multiple motor behaviors. Nature Communications , 9(1):4233, Oct 2018. ISSN 2041-1723. doi: 10.1038/s41467-018-06560-z. URL https://www.nature.com/articles/s41467-018-06560-z .

Juan A. Gallego, Matthew 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.1038/s41593-019-0555-4. URL https://www.nature.com/ articles/s41593-019-0555-4 .

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: 10.1101/2022.09.26.509498. URL https://www.biorxiv.org/content/10.1101/2022.09. 26.509498v1 .

MaxDabagia, Konrad P. Kording, and Eva L. Dyer. Aligning latent representations of neural activity. Nature Biomedical Engineering , 7(4):337-343, Apr 2023. ISSN 2157-846X. doi: 10.1038/s41551-022-00962-7. URL https://www.nature.com/articles/s41551-022-00962-7 .

David Sussillo, Sergey D. Stavisky, Jonathan C. Kao, Stephen I. Ryu, and Krishna V. Shenoy. Making brain-machine interfaces robust to future neural variability. 7:13749, Dec 2016. ISSN 2041-1723. doi: 10.1038/ncomms13749. URL https://www.nature.com/articles/ncomms13749 .

Brianna M. Karpowicz, 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.1101/2022.04.06.487388. URL https://www.biorxiv.org/content/10.1101/2022.04.06.487388v2 .

Francis Willett, Erin Kunz, Chaofei Fan, Donald Avansino, Guy Wilson, Eun Young Choi, Foram Kamdar, Leigh R. Hochberg, Shaul Druckmann, Krishna V. Shenoy, and Jaimie M. Henderson. A high-performance speech neuroprosthesis. Jan 2023. doi: 10.1101/2023.01.21.524489. URL https://www.biorxiv.org/ content/10.1101/2023.01.21.524489v1 .

B Wodlinger, J E Downey, E C Tyler-Kabara, A B Schwartz, M L Boninger, and J L Collinger. Tendimensional anthropomorphic arm control in a human brain-machine interface: difficulties, solutions, and limitations. Journal of Neural Engineering , 12(1):016011, Feb 2015. ISSN 1741-2560, 1741-2552. doi: 10.1088/1741-2560/12/1/016011. URL https://iopscience.iop.org/article/10.1088/1741-2560/12/1/ 016011 .

Felix Pei, Joel Ye, David Zoltowski, Anqi Wu, Raeed H. Chowdhury, Hansem Sohn, Joseph E. O'Doherty, Krishna V. Shenoy, Matthew T. Kaufman, Mark Churchland, Mehrdad Jazayeri, Lee E. Miller, Jonathan Pillow, Il Memming Park, Eva L. Dyer, and Chethan Pandarinath. Neural latents benchmark '21: Evaluating latent variable models of neural population activity, 2022.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In I. Guyon, U. V on Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper_files/paper/ 2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf .

Joel Ye and Chethan Pandarinath. Representation learning for neural population activity with neural data transformers, 2021. URL https://doi.org/10.51628/001c.27358 .

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. Dec 2021. doi: 10.48550/arXiv.2111.06377. URL http://arxiv.org/abs/ 2111.06377 . arXiv:2111.06377 [cs].

Chethan Pandarinath and Sliman J Bensmaia. The science and engineering behind sensitized braincontrolled bionic hands. Physiological Reviews , 102(2):551-604, 2022.

Daniel B. Silversmith, Reza Abiri, Nicholas F. Hardy, Nikhilesh Natraj, Adelyn Tu-Chan, Edward F. Chang, and Karunesh Ganguly. Nature Biotechnology , 39(3):326-335, Mar 2021. ISSN 1546-1696. doi: 10.1038/s41587-020-0662-5. URL https://www.nature.com/articles/s41587-020-0662-5 .

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, 2019.

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 NeurIPS 2022 Workshop on Learning from Time Series for Health , 2022. URL https://openreview.net/forum?id=c9qFg8UrIcn .

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.1038/s41597-022-01280-y. URL https://doi.org/10.1038/s41597-022-01280-y .

Christopher Wang, Vighnesh Subramaniam, Adam Uri Yaari, Gabriel Kreiman, Boris Katz, Ignacio Cases, and Andrei Barbu. BrainBERT: Self-supervised representation learning for intracranial recordings. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/ forum?id=xmcYx_reUn6 .

Armin W. Thomas, Christopher Ré, and Russell A. Poldrack. Self-supervised learning of brain dynamics from broad neuroimaging data, 2023.

Demetres Kostas, Stéphane Aroca-Ouellette, and Frank Rudzicz. BENDR: Using transformers and a contrastive self-supervised learning task to learn from massive amounts of EEG data. Frontiers in Human Neuroscience , 15, 2021. ISSN 1662-5161. doi: 10.3389/fnhum.2021.653659. URL https: //www.frontiersin.org/articles/10.3389/fnhum.2021.653659 .

Iyad Obeid and Joseph Picone. The temple university hospital EEG data corpus. Frontiers in Neuroscience , 10:196, 2016. ISSN 1662-453X. doi: 10.3389/fnins.2016.00196. URL https://www.frontiersin.org/ articles/10.3389/fnins.2016.00196 .

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 in Neural Information Processing Systems , volume 26. Curran Associates, Inc., 2013. URL https://proceedings. neurips.cc/paper_files/paper/2013/file/01386bd6d8e091c2ab4c7c7de644d37b-Paper.pdf .

Chethan Pandarinath, Daniel J. O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D. Stavisky, Jonathan C. Kao, Eric M. Trautmann, Matthew T. Kaufman, Stephen I. Ryu, Leigh R. Hochberg, Jaimie M. Henderson, Krishna V. Shenoy, L. F. Abbott, and David Sussillo. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature Methods , 15(10):805-815, Oct 2018. ISSN 1548-7105. doi: 10.1038/s41592-018-0109-9. URL https://www.nature.com/articles/s41592-018-0109-9 .

Francis R. Willett, Donald T. Avansino, Leigh R. Hochberg, Jaimie M. Henderson, and Krishna V. Shenoy. High-performance brain-to-text communication via handwriting. Nature , 593(7858):249-254, May 2021. ISSN 1476-4687. doi: 10.1038/s41586-021-03506-2. URL https://www.nature.com/articles/ s41586-021-03506-2 .

Darrel R. Deo, Francis R. Willett, Donald T. Avansino, Leigh R. Hochberg, Jaimie M. Henderson, and Krishna V. Shenoy. Translating deep learning to neuroprosthetic control. Apr 2023. doi: 10.1101/2023.04. 21.537581. URL https://www.biorxiv.org/content/10.1101/2023.04.21.537581v1 .

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.1038/s41586-023-06031-6. URL https://doi.org/10.1038/s41586-023-06031-6 .

Nur Ahmadi, Timothy G. Constandinou, and Christos-Savvas 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 1741-2552. doi: 10.1088/1741-2552/abde8a. URL https://iopscience.iop. org/article/10.1088/1741-2552/abde8a/meta .

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: 10.48550/arXiv.2202.06159. URL http://arxiv.org/abs/2202.06159 . arXiv:2202.06159 [cs, q-bio].

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 , 15(4):046016, Aug 2018. ISSN 1741-2552. doi: 10.1088/1741-2552/aab7a0.

Justin Jude, Matthew G. Perich, Lee E. Miller, and Matthias H. Hennig. Capturing cross-session neural population variability through self-supervised identification of consistent neuron ensembles. In Proceedings of the 1st NeurIPS Workshop on Symmetry and Geometry in Neural Representations , page 234-257. PMLR, Feb 2023. URL https://proceedings.mlr.press/v197/jude23a.html .

Ran Liu, Mehdi Azabou, Max Dabagia, Jingyun Xiao, and Eva L. Dyer. Seeing the forest and the tree: Building representations of both individual and collective dynamics with transformers. Oct 2022. URL https://openreview.net/forum?id=5aZ8umizItU .

Ali Farshchian, Juan A. Gallego, Joseph P. Cohen, Yoshua Bengio, Lee E. Miller, and Sara A. Solla. Adversarial domain adaptation for stable brain-machine interfaces. Jan 2019. URL https://openreview. net/forum?id=Hyx6Bi0qYm .

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/2022.08.26.504777. URL https://www.biorxiv.org/content/10.1101/2022.08.26.504777v1 .

Trung Le and Eli Shlizerman. Stndt: Modeling neural population activity with a spatiotemporal transformer, 2022.

Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Luˇ ci´ c, and Cordelia Schmid. Vivit: A video vision transformer, 2021.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations , 2021. URL https://openreview.net/forum?id= YicbFdNTTy .

Maryam M. Shanechi, Amy L. Orsborn, Helene G. Moorman, Suraj Gowda, Siddharth Dangi, and Jose M. Carmena. Rapid control and feedback rates enhance neuroprosthetic control. Nature Communications , 8(1):13825, Jan 2017. ISSN 2041-1723. doi: 10.1038/ncomms13825. URL https://www.nature.com/ articles/ncomms13825 .

Jennifer L Collinger, Brian Wodlinger, John E Downey, Wei Wang, Elizabeth C Tyler-Kabara, Douglas J Weber, Angus JC McMorland, Meel Velliste, Michael L Boninger, and Andrew B Schwartz. Highperformance neuroprosthetic control by an individual with tetraplegia. The Lancet , 381(9866):557-564, 2013.

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 C. 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.1101/2023.02.01.526664. URL https://www.biorxiv.org/ content/early/2023/05/04/2023.02.01.526664 .

Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning, 2021.

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.00784 , 2020.

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 .

Joseph G Makin, Joseph E O'Doherty, Mariana M B Cardoso, and Philip N Sabes. Superior arm-movement decoding from cortex with a new, unsupervised-learning algorithm. Journal of Neural Engineering , 15 (2):026010, Apr 2018. ISSN 1741-2560, 1741-2552. doi: 10.1088/1741-2552/aa9e95. URL https: //iopscience.iop.org/article/10.1088/1741-2552/aa9e95 .

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):0460e9, Aug 2021. ISSN 1741-2560, 1741-2552. doi: 10.1088/1741-2552/ac16b2. URL https://iopscience.iop.org/ article/10.1088/1741-2552/ac16b2 .

Kundan Krishna, Saurabh Garg, Jeffrey P. Bigham, and Zachary C. Lipton. Downstream datasets make surprisingly good pretraining corpora, 2022.

Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Joan Puigcerver, Jessica Yung, Sylvain Gelly, and Neil Houlsby. Big transfer (bit): General visual representation learning, 2020.

Amir Zamir, Alexander Sax, William Shen, Leonidas Guibas, Jitendra Malik, and Silvio Savarese. Taskonomy: Disentangling task transfer learning, 2018.

Ashish Teku Vaswani, Dani Yogatama, Don Metzler, Hyung Won Chung, Jinfeng Rao, Liam B. Fedus, Mostafa Dehghani, Samira Abnar, Sharan Narang, and Yi Tay. Scale efficiently: Insights from pre-training and fine-tuning transformers. 2022.

Danny Hernandez, Jared Kaplan, Tom Henighan, and Sam McCandlish. Scaling laws for transfer, 2021.

Fabio Rizzoglio, Ege Altan, Xuan Ma, Kevin L. Bodkin, Brian M. Dekleva, Sara A. Solla, Ann Kennedy, and Lee E. Miller. Monkey-to-human transfer of brain-computer interface decoders. bioRxiv , 2022. doi: 10.1101/2022.11.12.515040. URL https://www.biorxiv.org/content/early/2022/11/13/2022.11. 12.515040 .

Wei Wang, Sherwin S. Chan, Dustin A. Heldman, and Daniel W. Moran. Motor cortical representation of position and velocity during reaching. Journal of Neurophysiology , 97(6):4258-4270, Jun 2007. ISSN 0022-3077. doi: 10.1152/jn.01180.2006.

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 rnns for high performance brain-machine interfaces. bioRxiv , pages 2023-05, 2023.

Beata Jarosiewicz, Anish A Sarma, Jad Saab, Brian Franco, Sydney S Cash, Emad N Eskandar, and Leigh R Hochberg. Retrospectively supervised click decoder calibration for self-calibrating point-and-click brain-computer interfaces. Journal of Physiology-Paris , 110(4):382-391, 2016.

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.

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 arXiv:2203.02072 , 2022.

Nir Even-Chen, Sergey D Stavisky, Jonathan C Kao, Stephen I Ryu, and Krishna V Shenoy. Augmenting intracortical brain-machine interface with neurally driven error detectors. Journal of neural engineering , 14(6):066007, 2017.

Iñaki Iturrate, Ricardo Chavarriaga, Luis Montesano, Javier Minguez, and José del R. Millán. Teaching brain-machine interfaces as an alternative paradigm to neuroprosthetics control. Scientific Reports , 5(1): 13893, Sep 2015. ISSN 2045-2322. doi: 10.1038/srep13893. URL https://www.nature.com/articles/ srep13893 .

Guy H. Wilson, Francis R. Willett, Elias A. Stein, Foram Kamdar, Donald T. 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/2023.02.03.527022. URL https://www.biorxiv.org/content/10. 1101/2023.02.03.527022v1 .

Josh Merel, David Carlson, Liam Paninski, and John P Cunningham. Neuroprosthetic decoder training as imitation learning. PLoS computational biology , 12(5):e1004948, 2016.

Francis R. Willett, Daniel R. Young, Brian A. Murphy, William D. Memberg, Christine H. Blabe, Chethan Pandarinath, Sergey D. Stavisky, Paymon Rezaii, Jad Saab, Benjamin L. Walter, Jennifer A. Sweet, Jonathan P. Miller, Jaimie M. Henderson, Krishna V. Shenoy, John D. Simeral, Beata Jarosiewicz, Leigh R. Hochberg, Robert F. Kirsch, and A. Bolu Ajiboye. Principled bci decoder design and parameter selection using a feedback control model. Scientific Reports , 9(1):8881, Jun 2019. ISSN 2045-2322. doi: 10.1038/ s41598-019-44166-7. URL https://doi.org/10.1038/s41598-019-44166-7 .

Unknown. EvalAI leaderboard. https://eval.ai/web/challenges/challenge-page/1256/leaderboard/ 3184 , 2022. Accessed on May 16, 2023.

Christoph Feichtenhofer, Haoqi Fan, Yanghao Li, and Kaiming He. Masked autoencoders as spatiotemporal learners, 2022.

A Supplementary Material

A.1 Dataset Preparation

We aimed to perform minimal preprocessing on datasets. For pretraining, we 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 so as to not remove any potential cross-channel/session information. The one exception to this is that neurons with firing < 0.5Hz are removed in the sorted analysis of the 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.1Hz to upwards of 50Hz. We believe additional data curation is likely to improve model quality.

In total, the max number of pretraining trials or pseudo trials was on the order of 100K trials. Each trial lasted up to 2.5s (cropped or chunked if trials were longer), and used all recorded M1 activity, and PMd activity if available.

Decoding targets were either in a standard unit or in z-scores against the dataset mean and standard deviation (not a 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 z-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 on the function of the neural data; we are concerned comparisons on the 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 a semantic units of covariate data (experimental trials), and encourage documentation of pretraining composition.

Reaching datasets

Neural Latents Benchmark motor datasets ( MC_Maze , MC_Maze_small , MC_Maze_med , MC_Maze_large , MC_RTT ): ∼ 3.7K trials.

Churchland et al., obstacle-guided (maze) reaching, 2 monkeys, 9 sessions / ∼ 20K trials.

Nir-Even Chen et al., delayed reaching, 2 monkeys, 12 sessions / ∼ 80K trials total.

O'Doherty et al., self-paced reaching, 2 monkeys, 47 sessions / ∼ 40K seconds total.

Isometric manipulandum datasets

Gallego-Carracedo et al., isometric center-out and hold, 2 monkeys, 12 sessions, ∼ 2.7K trials total.

Dyer et al., 2 monkeys (same as above), 4 sessions/ ∼ 750 trials total.

Human BCI datasets [Private]

Human participant data from ongoing clinical trials. Subsetted to 2D cursor control activity, either under observation/attempted activity, partial, or full BCI control. During observation, participants observe a programmatically controlled cursors, which e.g. is performing center out at a steady pace in a 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 a trialized task where the cursor initializes on the center of a screen and the participant reaches to 1 of 8 radially

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 is 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 a 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, 0-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

The full, uncurated logs of all model training are available at https://wandb.ai/joelye9/context_general_bci?workspace=user-joelye9; interested users may compare specific hyperparameters logged there against those in the codebase during reproduction attempts.

Basic hyperparameters

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).

In pretraining we manually tuned LR to 5 e -4 in initial experiments and hold it fairly constant in pretraining. We swept learning rate in our hyperparameter comparisons below.

In pretraining, we use learning rate warm-up for 100 epochs, and decay to ϵ by 2500 epochs. This is a high threshold that is typically not reached: training converges within 100-1K epochs for our manually tuned LR. In fine-tuning, we experimented with similar ramping schedule but settled on fixed small LR (which are typically grid-searched).

For RTT, we swept and found that a decoding lag of 120ms worked reasonably well. (This is similar to reports in [14]. For human BCI, we do not use decoding lag.

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 2 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.

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 SLURM 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 a remote server and took <5 minutes.

Fitting datasets on the order of 1K trials typically requires 20m-1hr on 12G 1080/2080-series NVIDIA GPUs.

10K-20K trial datasets require 2-8 32G-V100 hours.

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

0.4

A

B

0.39

0.318

0.324

Test loss

Test loss

0.36

0.36

0.312

0.321

SessionTokens

SubjectTokens

0

0

1

0.33

1

0.32

8

8

102

103

104

103

104

Training Steps

Training Steps

Sorted

0.6

Intra-session

）

Makin18(rEFHEst)

Cross-Session

（

口

Cross-Subject

R2

Velocity

×

Cross-Task

0.5

Patch 32 (Pool)

Patch32 (Nopool)

Patch 8 (Pool)

Patch 8 (No pool)

NDT1

0.289

0.291

0.294

Test NLL（↓)

and adding them to the data token sequence (i.e. in-context aggregation). Our pilot experiments found no 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 to converged NLL (some 1 e -3 ) is modest but non-negligible, considering the NLL resolution in Fig. 3. This trend replicates at smaller data scales (Fig. 11). 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 in a similar architecture in the image domain, ViT. This is because units within a patch are not meaningfully related, especially across datasets; adjacent units often correspond to different physical locations on an electrode array, and this physical meaning is typically underspecified in open data releases. Smaller patch sizes, ideally at the single neuron level, thus appear to be a better design.

However, in early experiments we explored patch size but quickly found a tradeoff. 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 we mean pool across the population at each timestep (Pool) or not. Smaller patches may be worth revisiting if 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

Sample HP sweep for Indy 06/27/2016 (Sorted)

Sample HP sweep for Indy 06/27/2016 (Unsorted)

nodop

0.365

0.6150 -

（↑）

0.360

Poisson NLL(↓)

0.6125 -

Poisson NLL (

0.6100

·

0.355

0.6075

：

0.6050-

?

0.6025

nodop

：

0.345 -

0.6000-

0.5975

0.4

NDT1

NDT1

NDT1 + Stitch

NDT1 + Stitch

NDT2 (Patch 32)

NDT2 (Patch 32)

·

weight_decay

Poisson NLL (↓)

0.365

0.360 -

0.355

0.016

0.008

0.024

0.032

0.040

Poisson NLL (↓)

0.6125

0.6100

0.6150

0.6075

weight_decay

0.008

0.016

0.040

0.024

0.032

：：

0.6050

0.350-

0.6025

0.6000

0.345 -

0.5975

NDT1

NDT1

NDT1 + Stitch

NDT1 + Stitch

NDT2 (Patch 32)

NDT2 (Patch 32)

encoder and decoder layers, and the final layer's output is 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] ), weight decay ( [1 e -3 , 5 e -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 1 e -2 . In the code, this experiment is configured in exp/arch/tune_hp, exp/arch/tune_hp_unsort .

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 to 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 a 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, or only including stitching at 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 at each timestep, or we could use a thin Transformer decoder to allow information from multiple timesteps to aid the

Sorted

0.64

1

（

0.62

Mask Ratio

2

R

50%

0.60

25%

C

75%

0.58

0.56

0.290

0.291

0.292

Test NLL (↓ )

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 is 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 (e.g. 100K), and report with that setting. The primary difference is that cross-attention restricts neural data tokens from attending to the kinematic query tokens, while in-context strategies do not distinguish the two. For decoding probes, where the decoder is prepared on only a few hundred trials, we find it beneficial to mean-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 to make conclusions on design choices given variability in results.

Archictecture comparisons - Velocity decode vs. NLL (Sorted)

dataset=odoherty_rtt-lndy-20160407_02

dataset = odoherty_rtt-lndy-20160627_01

dataset = odoherty_rtt-lndy-20161005_06

0.65-rEFH

■

0.68

■

0.550 -

■

0.66 -

0.525

0.60

0.64

0.500

Velocity

VelocityR

0.62

Velocity

0.475

0.55

0.60

0.450

0.58

0.50-

0.56

0.425

0.54

rEFH

0.400

0.45

0.3745 0.3750 0.3755 0.3760 0.3765

0.3400.341

0.3420.343

0.344

0.345

0.254 0.255 0.256 0.257 0.258 0.259

Test NLL

Test NLL

Test NLL

dataset = odoherty_rtt-lndy-20161026_03

dataset = odoherty_rtt-lndy-20170131_02

Architecture

0.625

0.70 -

NDT2 (Patch 32)

0.600

NDT1

0.68

0.575

NDT1-Stitch

Velocity

0.66

忍

Velocity

0.550

TransferSource

0.525

rEFH

Session

0.64

0.500

Scratch (Single context)

0.62

0.475

Subject

0.450 -

+

Task

0.60-rEFH

0.425

0.244

0.245

0.246

0.247

0.248

0.238

0.240

0.242

0.244

Test NLL

Test NLL

Archictecture comparisons - Velocity decode vs. NLL (Unsorted)

dataset = odoherty_rtt-lndy-20160407_02

dataset=odoherty_rtt-lndy-20160627_01

dataset = odoherty_rtt-lndy-20161005_06

0.65-

0.75 -

0.65-rEFH

0.60

0.70

0.60

2

+

0.55

Velocity

0.65

Velocity

0.55

0.50

rEFH

0.50-

0.60

0.45

0.55

0.40

0.45

0.568

0.570

0.5720.5740.576

0.578

0.5950 0.5975 0.6000 0.6025 0.6050 0.6075

0.372 0.3740.3760.3780.3800.382

Test NLL

Test NLL

Test NLL

dataset = odoherty_rtt-lndy-20161026_03

dataset = odoherty_rtt-lndy-20170131_02

Architecture

0.70-

NDT2 (Patch 32)

0.75

0.65

NDT1

0.70

Velocity

0.60

NDT1-Stitch

Transfer Source

0.65

0.55

Session

Scratch (Single context)

0.60-rEFH

0.50

Subject

+

Task

0.55

0.45

+

0.4150.4200.4250.4300.435

Test NLL

Test NLL

0

Test loss

0.5

S =1.4e + 03 trials

0.4

0

Test loss

0.5

S=5.5e+03trials

ContextTokens

0.4

0

1

8

0

Test loss

0.5

S =2.2e + 04 trials

0.4

101

102

103

104

Training Steps

Unsup. Transfer Scaling (100 Trial Calibration)

Initial

Middle

Final

5.68 x 10-1.

6.2 x 10-1

4.44 × 10-1

Series

NLL

5.64 × 10-1

Test NLL

6.1 × 10-1.

NLL

Test

Intra

Session

Task

4.4 × 101

Subject

5.6 ×10-1

6 ×10-1

4.36 × 10-1

102

103

104

105

102

103

104

105

102

103

104

105

Pretraining trials

Pretraining trials

Pretraining trials

Sup. Transfer Scaling (1o0 Trial Calibration)

0.60-

0.6 -

0.60 -

0.45

Series

#-Intra

Session

0.4 -

0.45

0.30 -

102

103

104

105

102

103

104

105

102

103

104

105

Pretraining trials

Pretraining trials

Pretraining trials

Pretrain-PositiveRegimes

Initial

Middle

Final

5.7 × 10-1

0.0074

0.0113

4.46 × 10-1 ,

0.0097

5.68 × 10-1,

6.2 × 10-1

4.44 × 10-1

5.66 × 10-1,

0.0071

6.15 × 10-1.

0.0073

4.42 × 10-1

_0.0066

Pretraining

5.64x10-1

6.1x10-1

4.4 × 10-1

Pretrain Task

Pretrain Session

Scratch

5.62 × 10-1,

4.38 ×10-1

6.05 ×10-1

x-0.0033

5.6 ×10-1

x_0.0020

4.36 × 10-1

_0.0029

6×10-1

5.58 × 10-1,

4.34 × 10-1

102

2×10²3×1a²×10²6×102

102

103

102

2×10²3×1a²×10²6×102

Target context trials

Targetcontext trials

Target context trials

TuningaDecoder(Unsorted)

dataset=odoherty_rtt-lndy-20160627_01

dataset=odoherty_rtt-lndy-20170131_02

dataset=odoherty_rtt-lndy-20160407_02

0.8

0.8 -

0.7

0.7

UntunedPretrained

0.7

0.6

0.6

Variant

0.5

2

0.6

-

Scratch

Unsup tune

Sup tune

0.5

Scratch (100 Trial Sup)

0.4

0.5 -

UntunedPretrained

0.3

0.4-UntunedPretrained

0.2

0.4

102

103

102

102

Target session trials

Target session trials

Target session trials

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nSystems 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nL136: 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\nLIMITATIONS:\nThe authors have adequately addressed the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\n1. 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nSee weaknesses", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nOriginality: 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\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\ni.\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\n\nLIMITATIONS:\nThe 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?", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nMy 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.\n\nQUESTIONS:\nTo 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?\n\nLIMITATIONS:\nI see no specific limitation to be reported.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}