A Generalist Intracortical Motor Decoder

Overall: We thank the reviewer for their detailed feedback (and catching of several typos). The two main critiques were the insufficiency of baselines and the poor writing quality. In the detailed response we comment on why it is difficult to compare with the baselines you have requested, but we are beginning experiments comparing with the public pretrained NDT2 model, which we hope addresses the concern on baselines. As for the poor writing, we have amended the typos you have pointed out, and we are happy to receive further feedback on any low and high level area where the reviewer finds the writing confusing so that we might improve the text and overcome this issue of writing quality.

Weakness - Contribution / Baselining: We would like to clarify the relationship of the listed models to NDT3. NDT and EIT are single-session models, trained on a single day’s data, and are not foundation models. EIT further does not provide a codebase. POYO introduces a pretraining method and trains a model with up to 100-200 hours of data, but POYO does not release pretrained weights, which makes it difficult to compare against their method without greatly increasing the scope of our work (evaluating scaling on a fairly distinct architecture). Given the challenges to comparing against other works in our own paper, we have reported NDT3 on the newly released FALCON benchmark makes it easy to compare against NDT3 in future works. See also our global response on baselining.

Weakness - Writing Quality: We apologize for the typos in the text and have corrected your examples. For L214-215, we are referring to benchmarks, the noun, not the action of evaluating different models. We are saying that there is no third party dataset with fixed preprocessing and private evaluation that many people in the field submit to for decoding, like ImageNet. The NLB benchmark is referenced in POYO, for example, but we believe its use is inappropriate (see response to USSQ for details). We are not claiming that NDT3’s architecture is superior to NDT2 in general, only that it is easier to scale to decoding over large amounts of heterogeneous data for pretraining. However, the fact that from scratch NDT3 achieves better decoding than from scratch NDT2 across our extensive evaluation does show that it may be easier to train high performance decoders with NDT3 than NDT2.

Question - Why Change NDT2?: We changed the architecture from NDT2 to the causal transformer in NDT3 for two reasons. First, it allows us to use one model to decode data of different dimensionalities without needing different readout layers. A similar approach is used in other generalist models like TDM. Second, FlashAttention2, which allows several-times faster pretraining, does not support unconventional masking strategies. The open source implementation of NDT2 relies on a shuffling strategy that requires such an unconventional attention mask and is not easily compatible with FlashAttention 2’s requirements. We have added text to address this question in C.3.

Question - Why broken lines?: Yes, this is attributed to an explosion in loss. This occurs in several of our pretrained models but is also not uncommon in pretraining overall (we believe), and generally is not known to degrade the quality of the outcome models at convergence.

Question - Dimensionality > 2: NDT3 is pretrained on data from many different behavioral tasks. While some tasks are 2D, as you mention, other tasks, like when decoding recordings from different muscles in EMG, have many more target dimensions to decode. You can see these examples in Fig 2B. There are also high-dimensional kinematics, as in FALCON H1, which is 7D robotic upper limb control.

We hope these answers address some confusions and again, welcome further critiques on writing quality in the new section on experiments.

Overall: We appreciate the reviewer’s desire for rigor in the evaluation of NDT3. We agree that further baselining provides important context for the ultimate benefit of the NDT3 model, and will work to include baselines from the pretrained NDT2 checkpoint in this discussion phase. Considering that full-scale evaluations of new methods (even basic DNNs) are fairly expensive, we agree that external benchmarks are the natural place to compare additional methods, and will discuss these in detailed response. We also agree that the technical novelty and empirical benefit in this work is limited, but our contributions advance beyond model advances and also provide analyses showing where we can expect benefits currently and directions for future improvement (see global response to limited benefit).

Weakness - Limited benefit: As mentioned, we agree with the assessment of limited benefit, which motivated our framing of the model as a generalist that performs many tasks well rather than setting records on every task. Generalists are similarly limited in other fields (GATO, RT-X. We also want to emphasize that limited gains from scaling are not unprecedented. For example, Fig 4. In Entezari 21 shows that several downstream tasks don’t benefit significantly from orders of magnitude more images in pretraining. Given this context, and given the extent of our other experiments, we hope to convey that limited performance is more an interesting challenge for the field rather than a failing of our methodology.

Weakness - Practical cost: We agree that larger models are in general harder to use, but 350M parameter models are small for modern machine learning (BERT/GPT2 were this size in 2018) and represent the starting point for most works in foundation modeling. Their practicality can be demonstrated by the fact that consumer GPUs can tune models an order of magnitude larger, with the right tooling. While the necessity of scaling NDT3 specifically to 350M parameters is a concern, we suspect it is a concern for the field rather than a weakness of our method.

Weakness - Lack of baselines: We chose to sacrifice baselines used so as to instead maximize the number of datasets we evaluated on (we already tuned over 2000 models). Nonetheless, we agree with the reviewer that baselines provide important context, even if only partially evaluated. To this end, at least for basic baselines, RNNs are already evaluated on the FALCON benchmark, where they uniformly underperform multi-session NDT2 models (see appendix of the FALCON preprint and current leaderboard). We agree that the current scaling trends may not generalize to other foundation models, namely pretrained NDT2 or Azabou et al’s POYO. However, we believe that evaluating scaling trends on multiple candidate foundation models is out of scope, given that no model has been evaluated at large scale in this subfield. We are optimistic that our work can provide an important reference for future comparison, and have taken care to discuss the need for comparison with other methodologies in our discussion.

Weakness - NLB: We believe this is an important nuance that has been confused both in the field. NLB is a representation learning benchmark, and decoding performance on the leaderboards is evaluated by remote training of a linear probe on submitted neural data firing rate predictions. In contrast, our evaluations directly predict the behavior, without explicit neural data modeling. Therefore it is somewhat inappropriate to evaluate NDT3 with NLB. In contrast the used FALCON benchmark directly evaluates behavior, though it is new and other models have not been tested on this benchmark yet.

Further, we believe that Azabou et al. used the datasets from the NLB but preprocessed them for their own analysis and conducted private evaluation, which defeats the purpose of the benchmark. For example, there is no submission for POYO on the NLB leaderboards, even though the reported score is almost 5% higher than other methods. Since we cannot compare to POYO's reports in the official benchmark, and there is no pretrained weight available for POYO, we would have to train and tune POYO ourselves to make a comparison based on architectures. Beyond POYO, we are only aware of pretrained NDT2 and MtM (Zhang et al 2024)) for motor decoding. While MtM releases pretrained models, it is trained on mouse data outside the context of BCIs, and was released in July 24 (contemporary work).

11/24: We now report NLB evaluation in B.10, though as mentioned, these results are not comparable with POYO's reports.

(Response continued in thread)

Weakness - Pretrained NDT2 1, Documentation: Thank you for pointing this out. As you mention, NDT3 is a successor of NDT2, and is also a standard Transformer, so there is not much room or perhaps need for technical innovation in the architecture. The primary difference is that NDT2 uses MAE (Masked-autoencoder) dropout of its neural inputs, while NDT3 is a next-step model. A list of more specific technical advances over NDT2 is provided in C.6, but we agree that a detailed comparison of model preparation would be useful, and we have added this in C.3.

Weakness - Pretrained NDT2 2, Baselining: We did not collect pretrained NDT2 results as it was designed to show feasibility of scaling through unsupervised learning, not decoding, and from-scratch NDT2 under-performed from-scratch NDT3 on our evaluation datasets, providing no reason to believe that the pretrained NDT2 would perform better than pretrained NDT3. We agree though, that it provides important context, and have added a pretrained NDT2 evaluation. It did not outperform NDT3 from scratch.

Weaknesses - Novelty of Generalization: We agree, certainly, the topic of pretraining large models is active work in many subfields of neuroscience. However, different neural data modalities serve widely different communities at this point in time, and spiking data is central enough so as to have a significant audience if a foundation model were successful. Further, Zhang et al’s study of generalization appears most like Sec 3.1, showing broad generalization of the pretrained model when fine-tuning to different downstream tasks. However, Sec 3.2 and 3.3 differ from this general foundation model narrative in that we study generalization of fine-tuned representations, in that we evaluate generalization downstream.

Overall:

Weakness - Other factors: These are interesting points. We discuss them in their corresponding questions (Q1 and Q3) below. Weakness - Clarity: We appreciate this feedback. We have decreased the number of models reported in the main panels in Figure 3 and believe the colors are now easy to distinguish. The colors throughout Figure 3 now use the same palette, so they still refer to pretraining volume. We are happy to receive other feedback on writing and figures.

Q - Upper bounds: As you mention, there is no clear way to test the info-theoretic limit on possible decoding performance. Good deep networks are the current empirical strategy for estimating this limit. We want to distinguish two types of info-theoretic limits. First, there is the Bayes error due to variability in a given neural dataset, timepoint by timepoint, which we believe you are referring to. This is a limit similar to failing to classify an image because it is too blurry. We can estimate this upper bound by scaling training data in our target distribution. For example, this “task-inherent difficulty” cannot be the issue limiting our performance because scaling downstream data continues to improve performance on held-out timepoints in all models. We have calibrated the size of the fine-tuning datasets in our own evaluations so that we are not saturating them due to Bayes error. On the benchmark evaluation of FALCON M1, in contrast (Fig 11, rightmost panel), decoding performance in many models are essentially equivalent.

Second, there is a more subtle limit in how quickly a generalist can align to a test setting. Because the generalist must perform many tasks, we might not be able to be more data efficient in fine-tuning because of variability in the space of potential downstream tasks. This is why some tasks need lengthy prompts for language models to perform well, and is related to the task posed to toy Transformers performing in-context learning of linear regression (Garg 22). If the observations are noisy, 2 points may not be enough to fit the line. A final reference is a discussion of how Transformers can be shown to be computing this task recognition on this post by Adam Shai, which mentions computational mechanics as a formal framework for this problem. We have discussed this distinction qualitatively in the introduction in L68-78 and at the end of Sec 3.1.

If the reviewer is interested in a specific experiment to understand how different models overcome this second limit, we would recommend using the cleaner toy settings as studied in Garg 22. Neural data adds the additional complexity of unknown, temporally evolving noise.

Q - Normalization: We chose to max-norm over z-scoring because our data is not always unimodal. Our human data often has bimodal behavior (a cursor click signal being off or on). Further our data sometimes contain extreme outliers from sensor errors, which we want to ensure doesn’t disrupt model training. Normalizing by max value would effectively discard these points, while z-scoring is less robust to such outliers. We do not believe this choice impacts scaling, as normalizing data to a fixed small range is standard in machine learning, as long as the loss scale is reasonably adjusted. In practice the dynamic range of the normalized data is spanned by most points in pretraining and evaluation, so we also do not believe this choice of normalization impacts task balance. We are working on plots to demonstrate this point. Further, though this doesn’t exactly address your concern, NDT2 uses z-score normalization and our method compares favorably in direct comparison.

(Response continued in thread)

We appreciate the raising of the score. We agree that there is much to explore, possibly enough for several future works. We did not wish our paper to overclaim what scaling may look like for neural data, only to start the conversation. If there is a place where you find our language is too broad in claiming scalability for neural data overall, please let us know, and we will revise.

We'd also like to clarify your comments on NDT3's results vs POYO and NDT2. NDT3 evaluates scaling to the 2000 hour regime on 8 tasks, and we have highlighted trouble in a subset of tasks when moving from 200 to 2000 hours. POYO and NDT2 have both demonstrated successful scaling under the 200 hour regime, and specifically only for 2D movement, so we do not believe our work contradicts any existing result. Further, NDT3 handily outperforms NDT2 on the FALCON tasks - please see Figure 11. The NDT2 settings used in our evaluation largely match the publicly recommended settings, so we believe our evaluation to be fair.

Overall: We thank the reviewer for their responses. Given several comments on lack of clarity, we especially appreciate your suggestions for specific figures.

Weakness - Lack of clarity: We agree the writing is currently dense and perhaps overly technical. We have updated both figures according to your suggestions. We were concerned about technical completeness in Figure 4 (indeed the specific issue on joint vs sequential tuning is asked about by reviewer xNEo), but we appreciate that this was distracting in the main text. We believe these higher level comments are valuable feedback and would welcome pointers to other places where the flow is distracted.

Weakness - Limited utility: We agree that cross-subject use is a primary goal for NDT3. Our evaluations were designed entirely around cross-subject evaluation. Our primary evaluation shows that downstream benefit is variable but overall positive. We believe that this implies the potential for off-the-shelf improvements in at least some portion of the BCI community. On the other hand, we agree NDT3 may not be ready to be a general purpose foundation model for the field, but do believe our experiments highlight important challenges that might be informative for the Neuro-ML comunity.

Q - Non-neural tokens: In inference, behavior input tokens are zeroed and phase and return tokens are removed entirely. This removal doesn’t impact positional embeddings of the remaining tokens, as position embeddings are identifying only within each timestep (rotary embeddings are used across time), and to that end, specific ranges of position IDs are reserved for data of different modalities. For example, positions 10-20 are reserved for neural data, position 21 is reserved for return, and positions 21-36 are reserved for behavior. This factorized embedding strategy is described in Sec C.6. We have updated Fig 2 to distinguish the zero-ing vs dropping of tokens.

Q - Zero-padding: We concatenate segments to increase the throughput of real data seen in the model. This strategy is standard in language model training and we mention in C.2 that Geiping and Goldstein 2022 has shown this preprocessing choice has little performance impact in language. A full comparison in the neural data domain would be helpful but would require training with the more expensive padded preprocessing.

Q - Statistics: Thank you for pointing this out. We’ve now measured significance using paired t-tests with FDR correction, finding that the largest model has statistical significance over all other models except the 350M 200 hr model. Differences between other pretrained models are largely nonsignificant. We’ve highlighted some tests now in the main text and included a full significance matrix in the appendix (Fig 13). Since we are aggregating tasks of widely different performance (e.g. all data points in Fig 1B), we’ve omitted the resulting wide CIs to avoid confusion. We believe this omission is standard in multitask evaluations. For example, on page 46 of the PaliGemma vision-language model report, extensive experiments yield average performance benefits across tasks on the order of 2-3%. Nonetheless, CIs derived from these data also overlap much like our own. We would welcome feedback from the reviewer on how to present this issue if they have insights.

Q - Cross-entropy over Poisson: In early experiments, we found this choice did not greatly impact our evaluations and chose the cross-entropy since it is a more broadly common and expressive model, even if Poisson losses are common for neural data models. We mention our existing justification of this choice in C.6: an ablation in Fig 8C shows that ablating the neural loss entirely has minimal impact on evaluation performance. Due to this result, we did not specifically ablate the choice to use the cross-entropy loss. However, we have run a new scaling analysis to support that the cross-entropy loss is nontrivial, labeled Fig 9 in the revision. We tested pretraining with only the unsupervised cross-entropy loss up to 200 hours of data, and fine-tune on downstream tasks, still in a supervised fashion. The fact that there is scaling of downstream performance implies the neural objective is useful. The MSE behavioral objective may obscure differences in the neural objective, though. We’d recommend future work studying specifically neural representation learning, as opposed to decoding, explore this choice in more detail. Further, yes, there is a coefficient to balance MSE and cross-entropy loss to roughly equal orders of magnitude. We experimented with this briefly and found no consistent impact within an order of magnitude of changes on these coefficients.

Q - FEF/MT: Yes, throughout our work, monkey evaluations were conducted on subjects held out from pretraining entirely, with the intention of evaluating cross-subject capability. This includes the S1/FEF/MT experiments.