SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Balancing Memorization and Generalization in RNNs for High Performance Brain-Machine Interfaces
Joseph T. Costello, Hisham Temmar, Luis H. Cubillos, Matthew J. Mender, Dylan M. Wallace, Matthew S. Willsey, Parag G. Patil, Cynthia A. Chestek
Departments of Electrical and Computer Engineering, Biomedical Engineering, Robotics, and Neurosurgery, University of Michigan, Ann Arbor, MI, USA
Brain-machine interfaces (BMIs) can restore motor function to people with paralysis but are currently limited by the accuracy of real-time decoding algorithms. Recurrent neural networks (RNNs) using modern training techniques have shown promise in accurately predicting movements from neural signals but have yet to be rigorously evaluated against other decoding algorithms in a closed-loop setting. Here we compared RNNs to other neural network architectures in real-time, continuous decoding of finger movements using intracortical signals from nonhuman primates. Across one and two finger online tasks, LSTMs (a type of RNN) outperformed convolutional and transformer-based neural networks, averaging  $18\%$  higher throughput than the convolution network. On simplified tasks with a reduced movement set, RNN decoders were allowed to memorize movement patterns and matched able-bodied control. Performance gradually dropped as the number of distinct movements increased but did not go below fully continuous decoder performance. Finally, in a two-finger task where one degree-of-freedom had poor input signals, we recovered functional control using RNNs trained to act both like a movement classifier and continuous decoder. Our results suggest that RNNs can enable functional real-time BMI control by learning and generating accurate movement patterns.
Brain-machine interfaces (BMIs) have the potential to restore motor function to people with paralysis. While intracortical motor BMIs have successfully enabled paralyzed human patients to control computer cursors [1], robotic arms [2], or even move their own muscles through functional electrical stimulation [3], the decoding algorithm that predicts intended movement from neural signals still lacks accurate and naturalistic movement outputs. Recently, non-linear decoding algorithms using artificial neural networks have demonstrated higher offline kinematic prediction accuracy compared to linear approaches, which could enable more functional BMIs. Glaser et al. 2020 found that recurrent neural networks (RNNs), specifically the long short-term memory (LSTM, [4]) and gated-recurrent unit (GRU, [5]), outperformed other decoding architectures for predicting movement in several brain regions [6]. Other neural networks have improved decoding accuracy by de-noising neural signals using RNNs (LFADS, [7]) or transformers [8]. However, offline accuracy does not necessarily predict real-time closed-loop (online) performance [9, 10], where errors are cumulative and the user can adjust neural activity based on feedback. With goals of improving function of real-time BMIs, it is unclear if the high offline accuracy of RNNs translates to the online setting.
In online prediction settings, BMIs using RNNs for continuous movement prediction have demonstrated limited success but have not been evaluated against other neural network architectures. Sussillo et al. 2016 showed that an online RNN was more robust to electrode perturbations and slightly outperformed a linear Kalman filter in a reaching task [11]. In a closed-loop simulation, Hosman et
al. 2019 found that an LSTM decoder could outperform a Kalman filter in a cursor-control task [12]. Recently, Deo et al. 2023 found that RNNs tested with a human subject tended to significantly overfit to offline data and required novel dataset regularization for functional online use [10]. Alternatively, our group found that an ReFIT convolutional feedforward network outperformed a ReFIT Kalman filter in an online two finger group task by accurately predicting a large dynamic range of velocities [13]. Thus, there is a need to evaluate the online performance of RNNs against other high performing neural networks where all decoders have the possibility of overfitting to training data.
In addition to high prediction accuracy, RNNs have the capacity to memorize and generate realistic sequences with minimal input [14, 15, 16]. In the BMI context, an RNN decoder could memorize discrete postures and stereotyped movements between postures, and could help generate movement for degrees-of-freedom (DoFs) where little information is contained in the input (for example, if few cortical channels are tuned to one movement direction). RNNs have previously been used to generate smooth, multi-DoF limb movements [17, 14, 18], and other work has shown that RNNs can learn hundreds of separate simple tasks [19, 16]. However, it is unclear if a memorized RNN is controllable in a closed-loop setting where overfit decoders may easily fail [10].
Here, we evaluate the online closed-loop performance of RNNs against other neural network architectures in non-human primates performing a multi-DoF finger task. We find that offline accuracy generally predicts the relative ordering of online decoder performance, with LSTMs outperforming other architectures. Then, we show how reducing task complexity and allowing memorization of movements can enable the BMI to reach able-bodied performance, with a gradual transition between discrete and fully continuous movement outputs. We provide an example of using selective movement memorization to recover performance despite having poor neural inputs for a particular DoF. Finally, we find that RNNs trained to memorize a small set of movements still rely on user inputs and have some limited capacity to generalize when controlling velocity rather than position.
2.1 Microelectrode Array Implants
We implanted two male rhesus macaques (Monkeys N and W) with Utah microelectrode arrays (Blackrock Neurotech) in the motor cortex using the arcuate sulcus as an anatomic landmark for hand area, as described previously [20, 21]. In each animal, a subset of the 96 channels in M1, with threshold crossings morphologically consistent with action potentials, were used for offline recordings and closed-loop BMI control. Surgical procedures were performed in compliance with NIH guidelines as well as our institution's Institutional Animal Care & Use Committee and Unit for Laboratory Animal Medicine.
2.2 Signal Processing and Feature Extraction
We recorded 96-channel Utah array data using a Cerebus neural signal processor (Blackrock Neurotech). The Cerebus sampled data at  $30\mathrm{kHz}$ , applied a  $300 - 1000\mathrm{Hz}$  bandpass filter and downsampled to  $2\mathrm{kHz}$  before streaming the data to a computer running xPC Target version 2012b (Mathworks), which calculated spiking band power (SBP) by taking the magnitude of the signal and summing into  $32\mathrm{ms}$  bins. We used SBP since it is highly correlated with neural firing rate and a high-performance BMI input feature [22]. During online experiments with neural network decoders, the binned SBP of each channel was normalized to have zero mean and unit variance, based on the mean/variance of initial training trials.
2.3 Behavioral Task for Finger Movement Decoding
We trained monkeys N and W to acquire virtual targets with virtual fingers shown on a computer screen in front of the animal (Figure 1a). In hand control ("offline" trials), monkeys moved their fingers within a manipulandum that measured the angles of the index and middle-ring-small (MRS) finger groups using bend sensors, and controlled the virtual fingers (Figure 1b). During brain control (also known as closed-loop, "online" trials), monkeys controlled the virtual fingers using neural signals and a decoder. The finger task required placing the virtual index finger and MRS fingers on the respective targets and holding for  $750~\mathrm{ms}$  during offline trials or  $500~\mathrm{ms}$  during online testing. The target size was  $15\%$  of the active range of motion. The total number of possible target locations
Figure 1: BMI for decoding finger movement. (a) Experimental setup. During hand-control trials the monkey used a finger manipulandum to control a virtual hand and acquire targets, with visual feedback from the screen. Spiking-band power (SBP) was calculated in real-time from 96-channel Utah microelectrode arrays. During brain control, the decoded finger positions controlled the virtual hand. (b) Example movements for the 2-DoF task involving flexion and extension of the index finger and MRS fingers. In a random target task, the target is randomly chosen between extension and flexion for each DoF. (c) SBP was averaged into 32 ms bins and fed into an RNN (GRU or LSTM) decoder which predicted position and velocity for one or two finger groups. Other decoders not depicted here were also tested. (d) Example true and predicted positions and velocities from an LSTM decoder for Monkey N performing a 2-DoF random target task. The LSTM is able to predict both slow and fast velocities.
varied based on the task, ranging from 2 targets up to fully random (any position could be chosen), with 1-DoF (just index finger) or 2-DoF (index and MRS fingers). Target locations were randomly chosen at the start of each trial. We collected 500 offline trials for decoder training (although only 300 were collected for offline memorization, Figure 3a). We note that Monkey N typically continued to move his hand even during brain control mode. In this work we only used offline historical data from Monkey W who did not perform online control.
2.4 Simulated Datasets
For some analyses we created simulated offline datasets of a virtual user performing the same target acquisition task. The goal of these simulations was to test the relative impact of amount of training data, number of DoFs, and number of inputs on decoder performance, rather than measuring absolute performance. The simulated user moved with a velocity proportional to the distance to the target along each DoF. Artificial neural channels were generated such that each channel had a random relationship with the position, velocity, and acceleration at each timestep, as suggested in [23]. Further details and specifics can be found in Supplementary Methods.
2.5 Performance Metrics
Offline performance was measured by the Pearson correlation and mean-squared error (MSE) between predicted and true finger position and velocity. Offline velocities and positions were normalized to have unit standard deviation and zero mean (normalized using training data). We also calculated the mean-absolute error (MAE) between velocity predictions and the true velocity. To do so, true velocities were binned into 21 bins between -3 to 3 standard deviations, and the mean-absolute error was calculated for the predictions of each bin, where data across days were grouped together.
Online performance was measured using two sets of metrics: time and information throughput. Time metrics included the trial-time (the time for a trial with hold time subtracted), move-time (the time until all fingers touch the target for the first time), and orbiting time (trial-time minus move-time, representing extraneous movements where fingers left the target). Information throughput, measured in bits per second (bitrate), has previously been used to measure finger task performance [13] and was calculated using Fitt's law which accounts for both task difficulty and completion time:
$$
t h r o u g h p u t = \frac {\sum_ {k} \log_ {2} \left(1 + \frac {D _ {k} - \frac {R}{2}}{R}\right)}{t _ {a c q}} \tag {1}
$$
where  $D_{k}$  is the distance of the  $k$ th virtual finger to the center of the  $k$ th target,  $R$  is the target radius, and  $t_{acq}$  is time to reach the target. Better decoders (fast and precise) have a higher throughput. Before calculating throughput, we removed initial trials from the performance calculation to account for the decoder learning period before steady state performance, which took approximately 1-2 minutes. When more than 100 trials were performed we removed the first 50 trials, and when fewer than 100 trials were performed (typically when the performance was poor and the monkey was unable to complete many trials) we removed the first 20 trials.
2.6 Decoders
We tested five decoder architectures, chosen for their high offline or online decoding performance in other studies. Decoders took in binned SBP features (neural features) from up to 96 channels with 1-5 bins of time history for each channel and were trained to predict the position and velocity of 1 or 2 finger groups (Figure 1c). The neural network decoders were implemented and trained using PyTorch 1.12.1. Each decoder is detailed below; see Supplemental Figure 1 for more details on each network architecture. The total number of neural network decoder parameters varied from 479k for the LSTM to 923k for the TFM (see supplemental results).
Recurrent Neural Network (RNN) We tested two variants of RNNs: the LSTM [4] and GRU [5], which have previously shown high performance in offline neural decoding [6]. Decoders were implemented using the torch.nn.LSTM and torch(nn.GRU classes. As shown in Figure 1c, the RNN decoder takes in the current time bin of SBP features and the previous hidden state, updates the hidden state, and predicts position and velocities as a linear function of the hidden state. During training, a sequence length of 20 and zero-initialization for the hidden state were used for each sample. The final output of each sample sequence was used to compute the loss. During online decoding, the hidden state was stored in memory, such that the inputs at each timestep were the current neural bin and the previous hidden state.
Convolutional Feedforward Neural Network (FNN) We tested the FNN described in Willsey et al. 2022 [13], which previously outperformed a ReFIT Kalman Filter [24, 21] in online testing. The network uses an initial convolutional layer over time (for each channel) followed by several feedforward layers that use batch normalization, dropout, and ReLU. The network used bins of time history per channel (160 ms). Implementation details are described in [13] and Supplemental Figure 1, and hyperparameters were optimized as described below. Notably, we did not perform ReFIT recalibration as was done in [13].
Transformer (TFM) Transformers [25] have recently become popular for sequence processing tasks and have performed well in offline neural decoding (Ye 2021). The transformer tested here used a cosine positional encoding layer [25] followed by multi-head attention and feedforward layers, and a final linear layer to output kinematic predictions (see Supplemental Figure 1). The transformer layers were implemented using the torch.nn TransformersEncoderLayer. The transformer used five bins of time history per channel (160 ms). We note that unlike the transformer in Ye et al. 2021 [8] which was trained to predict smoothed firing rates, the transformer here was directly trained to predict kinematics.
Kalman Filter (KF) The kinematic Kalman Filter (KF) is a linear decoder widely used in BMI that optimally combines a state transition model (using the current position and velocity as a state) with an observation model (that relates the kinematics to neural activity) to estimate the current position and velocity. Like the other decoders, the KF takes in the current bin of neural features and the previous
state, and then outputs predicted finger velocities. The KF was implemented and trained as described in [20, 26] using the NumPy library, but here we did not use ReFIT recalibration.
2.7 Neural Network Decoder Training
For offline decoders, datasets were split into  $70\%$  for training,  $10\%$  validation, and  $20\%$  testing. For online decoders, datasets were split into  $80\%$  training (400 trials) and  $20\%$  validation (100 trials) since a test set was unnecessary. Neural features and finger positions/velocities were normalized based on the training set to have zero mean and unit variance. We first performed hyperparameter optimization for each decoder using two offline datasets for each monkey individually. Bayesian optimization was performed using the Optuna library, and decoder performance was generally robust to the specific choice of hyperparameters. See Supplemental Methods for details and Supplemental Table 1 for parameter values used. Decoders were trained using minibatch gradient descent with the PyTorch Adam optimizer. The loss function was the mean-squared error (MSE) loss with an additional term to penalize finger co-dependence (see Supplemental Methods; the additional term was only used for online decoders). The learning rate was held constant until validation loss plateaued, then halved and held constant until validation loss plateaued a second time (using the ReduceLROnPlateau scheduler). After training was complete, a linear regression was calculated to transform the normalized outputs back to the initial position and velocity scale. Training took less than 5 minutes for each network.
2.8 Online Decoder Testing
During online experiments, decoders were run in Python 3.7 using a dedicated linux computer with an RTX 2070 Super GPU (NVIDIA). After receiving binned neural features from the xPC, the decoder calculated position and velocity predictions, and a final position for each finger was sent back to the xPC to update the virtual hand (see [13] for more details). For the RNN decoders, we found that combining velocity and position predictions to generate the final displayed positions (98% integrated velocity, 2% position) was necessary to prevent the virtual fingers from becoming biased (similar to that of [27, 11]), and we used this for all neural network decoders. For tasks with less than 9 targets we increased the positional contribution up to 50%. During online comparisons, decoders were alternated in an A-B-A or A-B-A-B format, with performance averaged across sets. For the 2D decoder comparison with all 5 decoders, each decoder was only tested once (A-B) due to time. An online decoder test was stopped if the monkey was unable to acquire targets for more than 30 seconds or the monkey stopped attempting to acquire targets.
2.9 Training Optimizations
To optimize online performance, in early exploration we tested several machine learning techniques for improving accuracy and generalization (detailed in Supplemental Methods). When the initial online trial success rate was  $90\%$  or less, we found the following techniques improved individuated finger control and success rate (SR): modifying the loss function to encourage finger independence (11% increase in SR, 12% improvement in trial time), training on more single-finger movements (38% increase in SR, 27% improvement in trial time), and training on trials with added positional perturbations (11% increase in SR, 12% improvement in trial time). Adding a small amount of noise to the neural data during training (as suggested in [28]) improved average trial times by  $6\%$ . Finally, we found that training on more data improved offline accuracy (up through 2000 trials), but with only a small improvement after 500 trials. For online results detailed below, decoders used the modified loss function, additional single-finger movements, and added noise.
3.1 High Performance finger BMI with RNNs
Offline Decoder Comparison Previous work has shown that RNN decoders, specifically LSTMs and GRUs, can outperform other feedforward neural networks for predicting offline reaching movements from intracortical signals [6]. Here, we began by establishing a baseline for how well RNN decoders can decode dexterous finger movements, compared to other neural network architectures. One adult male rhesus macaque, Monkey N, was implanted with Utah arrays in the hand area of the
Figure 2: Offline and online decoder performance comparison. (a) Offline model performance measured by MSE (a.u) between predicted and actual movements. Values represent the average of ten models trained on ten separate days of the 2-DoF random task. (b) Mean absolute error of velocity predictions where the x-axis is the true velocity and the y-axis is the error in predicted velocity (lower is better). Units are normalized velocity and data is averaged across ten days. (c) Offline correlation between true and predicted movements for simulated datasets with varied channel count. As channel count increases, the difference in accuracy across decoders becomes smaller. (d) A representative online decode trace of Monkey N using an LSTM for a 2-DoF random target task. The fingers quickly move to the target and can hold with minimal overshoot. Boxes represent the target position for each finger and lines represent the finger positions. (e, f) Online model performance on exemplar days where all decoders were tested in the same day for 1-DoF (e) and 2-Dof (f) random tasks. Performance is measured by bitrate (higher is better). "Hand" is the performance during physical hand control (able-bodied performance). Across all figures, error bars denote one standard error of the mean (SEM), and data for a,b,d-f come from Monkey N.
primary motor cortex, and simultaneous neural activity and finger movements were recorded. For a continuous, random target task we compared the following decoders across ten days: LSTM, GRU, FNN, TFM, and KF. An example offline prediction on a representative day using an LSTM is shown in Figure 1d, where the predicted positions and velocities closely match the true finger movements, notably reaching both fast and slow velocities. To evaluate decoder accuracy, we calculated the MSE and correlation between predicted and true finger velocities and positions. As in Figure 2a, for Monkey N, LSTM decoders had the lowest positional and velocity error across all decoders and also the highest correlations of 0.86 and 0.77 for position and velocity respectively. The LSTM had a significant improvement in position and velocity MSE over all decoders (p<0.001 for all comparisons, one-sided paired t-test). All neural network decoders had significantly lower position and velocity MSE than the KF (p<1e-5 for all comparisons, one-sided paired t-test). The same relative ordering of decoder was found when the analysis was performed on historical data from a second monkey (Supplemental Figure 2), albeit with slightly higher errors overall.
To further reveal the difference in decoder accuracy with Monkey N, we calculated the mean-absolute error (MAE) for each ground-truth velocity (Figure 2b). All neural network decoders had similar low error at slow velocities; however, RNNs had significantly lower error for the fastest  $10\%$  of velocities than other neural networks  $(p < 1e - 5$ , one-sided t-test), suggesting they may enable faster real-time movements while maintaining slow-movement accuracy. Thus, RNNs more accurately predict offline finger movements than other architectures, agreeing with previous literature [6].
Online Decoder Comparison Since real-time, closed-loop control requires the user to actively react to and correct errors, the online and offline distributions of neural activity and movements may differ, preventing a high performing offline decoder from functioning well online. Therefore, we next compared decoders in a closed-loop setting, where only the monkey's neural activity (rather than his hand) controlled the virtual hand movement. To evaluate online performance, we compared
decoders within multiple sessions with Monkey N, and measured performance using information throughput (bitrate). Decoders were trained using only same-day data. Figure 2d shows example online movements using a 2-DoF LSTM, where the monkey quickly moved each finger group and then held at the target position, for a median time of 1.3 sec per trial. The LSTM tended to overshoot index-flexion targets, which Monkey N also tended to do during hand-control. On all tests, the LSTM decoder matched or had the highest information throughput, with a median bitrate  $15\%$  higher than the GRU across 2 comparison days, and  $18\%$  higher than the FNN across 3 days. Supplemental Video 1 shows example usage of each decoder.