SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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
1. In both pretraining and fine-tuning, we scale batch size (accumulating batches or using multi-GPU training when necessary) to be roughly proportional to full dataset so that each epoch requires 10-100 steps; we find performance is not too sensitive to batch size within an order of magnitude of this heuristic (especially in pretraining).
2. In pretraining we manually tuned LR to  $5e - 4$  in initial experiments and hold it fairly constant in pretraining. We swept learning rate in our hyperparameter comparisons below.
3. In pretraining, we use learning rate warm-up for 100 epochs, and decay to  $\epsilon$  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).
4. 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.
5. For human offline evaluation, we take the best of two evaluation hyperparameter settings:  $10\%$  or  $50\%$  masking during target-session tuning. We also report the  $R^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.
6. Hyperparameters were not swept further during realtime pilots.
Compute costs We estimate computational costs with respect to data volume, as model size is held relatively static (6-12 layers, 128-384 hidden size). Most analysis was run on SLURM clusters. Pilot realtime decoding used an NVIDIA 1060/2060, where loop time was  $< 15\mathrm{ms}$  (neural data was binned at 20ms). Model tuning was performed on a remote server and took  $< 5$  minutes.
1. Fitting datasets on the order of 1K trials typically requires 20m-1hr on 12G 1080/2080-series NVIDIA GPUs.
2. 10K-20K trial datasets require 2-8 32G-V100 hours.
3.  $100\mathrm{K}+$  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
Figure 6. Context embedding ablations. A. Multi-session model training curves, with varying session context token count (3 seeds). Learning is improved, but additional tokens do not have notable effect. The converged score in only modestly affected. B. Subject transfer training curves with similarly varied token budget for subject embedding. Models receive 1 session token. There is no clear additional improvement nor harm.
Figure 7. 32-neuron patches compared against 8-neuron patches. Smaller patches sometimes improve NLL but reliably decode more poorly. (Pool)-ed models average-pool all the neural data features for a given timestep before being given to the decoder, whereas regular decoders simply cross-attend all neural data features.
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  $1e - 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)
Sample HP sweep for Indy 06/27/2016 (Sorted)
Figure 8. Sweeps for regularization parameters. NDT2 requires lower dropout.
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 ([1e - 3, 5e - 2], and hidden size (128, 256). NDT2 does have higher variance, but the main sensitivity is to dropout. We run this sweep and test evaluation in one training stage, Our base NDT2 uses dropout 0.1, hidden size 256, weight decay  $1e - 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 initializes 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
Figure 9. Mask ratios over 5 datasets. At test time, the given ratio is held out during unsupervised evaluation.
prediction. We experimented with both and chose the latter for minor gains. For simplicity, we run the Transformer decoder in one forward pass for all timesteps, i.e. there 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.
Architecture comparisons - Velocity decode vs. NLL (Sorted)
Architecture comparisons - Velocity decode vs. NLL (Unsorted)
Figure 10. We breakout Fig. 3 into individual datasets (points indicate means on 3 seeds). NDT2 shows consistent improvements over stitching, single session baselines, and rEFH (in most cases), but the ranking between data sources shifts, particularly for decoding scores.
Figure 11. Session context tokens improve learning, as measured by unsupervised loss in model training curves, though the majority of benefit of realized with 1 token. We compare this for three data scales, annotated by S, where we scale the number of trials available per session. (The increments were  $100\%$ ,  $25\%$ , and  $6.25\%$  of the data). We hypothesized that more data per session would make session tokens less relevant, but the primary effect of increasing convergence appears unchanged at these scales.
Unsup. Transfer Scaling (100 Trial Calibration)
Sup. Transfer Scaling (100 Trial Calibration)
Pretrain-Positive Regimes
Figure 12. Individual dataset results for three scaling analyses. We presented the largest middle session in the primary text, here we also show results on the first and final sessions in the dataset. The unsupervised and supervised transfer scaling reiterate the previous conclusions: cross-subject and task transfer provides low returns on scaling for unsupervised reconstruction, and decoding results are much more optimistic than unsupervised results. For convergence analysis (Pretrain-positive regimes), all three trend lines suggest convergence beyond 1K trials.
Figure 13. The same breakout as above for pretrained decoder tuning. Again, we find that decoder tuning reliably outperforms non-adaptive decoders.