SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Dataset and Preprocessing In this study, we collect neural signals from monkey Affogato (A), and Beignet (B) with  $\mu$ ECoG implanted in motor cortices, covering five different subregions (Primary Motor Cortex (M1), Premotor Cortex, Frontal Eye Field (FEF), Supplementary Motor Area (SMA), and Dorsolateral Prefrontal Cortex (DLPFC)) with 239 effective  $\mu$ ECoG electrodes arrays which have been used in [67]. During data collection, the subjects performed a reaching task towards one of eight different directions. All procedures were approved by the University of Washington Institutional Animal Care and Use Committee. In the following experiments, we include all electrodes from Monkey A and electrodes in M1 from Monkey B (More in the Appendix B.2). We additionally apply AMAG on two public datasets with field potential recorded with Utah arrays from the M1 of Monkey Mihili (M) and Chewie (C) where subjects performed similar hand-reaching task [22, 21]. M and C datasets contain the computed Local Motor Potential (LMP) and power band signal from eight frequency bands:  $0.5 - 4\mathrm{Hz}$ ,  $4 - 8\mathrm{Hz}$ ,  $8 - 12\mathrm{Hz}$ ,  $12 - 25\mathrm{Hz}$ ,  $25 - 50\mathrm{Hz}$ ,  $50 - 100\mathrm{Hz}$ ,  $100 - 200\mathrm{Hz}$ , and  $200 - 400\mathrm{Hz}$ . The resulting signal is downsampled to 30ms. Full broadband data from A and B were collected at  $25\mathrm{kHz}$ . We then applied a preprocessing procedure to extract LMP and powerband signals from A and B datasets as used in M and C datasets [20, 64]. The powerband signals are included as additional feature dimensions. In experiments, we treat a target reaching trial as a sample for all four datasets. On M and C datasets, trials are temporally aligned to detect movement onset and include  $300\mathrm{ms}$  before to  $450\mathrm{ms}$  after movement onset. For the A and B datasets, trials are temporally aligned to the appearance of the peripheral target and include  $600\mathrm{ms}$  after the peripheral target onset.
Benchmark Methods, We compare AMAG with existing approaches, including variants of RNNs, Transformers, and GNNs. Specifically, we compare with the following methods:  $LRNN$ , RNN model without a non-linear activation function and gating mechanism [63];  $RNNf$ , RNN model with non-linear activation and gating mechanism [32, 47];  $LFADS$ , RNN based model with precomputed initial states [65].  $NDT$ , a decoder based transformer in temporal domain to estimate neural firing rate [78];  $STNDT$ , a model that is similar to NDT with additional attention based spatial interaction between neurons [37];  $TERN$ , a model that includes the GRU for encoding and temporal Transformer for decoding [32, 47].  $RNN PSID$ , RNN based model with 2-stage training for learning behavior-related neural space and behavior-unrelated neural space combined to ensure accurate
Figure 3: Panel A: Examples of predicted neuron trajectory for one-step (A.1) and multi-step (A.2) with representative methods DCRNN (orange), GWNet (green), GS4-G (purple), and AMAG (red). Ground Truth (GT) trajectories are shown in dashed blue. Panel B: Evaluation of the effect of attention dimension on NDT performance, with weight decay (w decay) in orange triangles, without weight decay (w/o decay) in purple circles. The larger size of the marker reflects the larger size of the model.
Table 1: One-step forecasting results on Monkey M, C, B, and A with benchmark methods and AMAG.
Monkey M	Monkey C	Monkey B	Monkey A
R2↑	Corr ↑	MSE ↓	R2↑	Corr ↑	MSE ↓	R2↑	Corr ↑	MSE ↓	R2↑	Corr ↑	MSE ↓
STNDT [37]	0.778±1e-4	0.883±4e-4	0.0333±1e-4	0.860±3e-3	0.928±1e-3	0.0111±2e-4	0.857±2e-3	0.932±8e-4	0.0089±2e-4	0.879±1e-3	0.939±7e-4	0.0074±7e-5
NDT [78]	0.837±2e-3	0.915±2e-3	0.0292±2e-3	0.908±4e-4	0.953±1e-4	0.0074±3e-5	0.897±3e-4	0.950±1e-4	0.0058±2e-5	0.924±6e-4	0.962±2e-4	0.0046±4e-5
LFADS [65]	0.847±3e-4	0.922±2e-4	0.0274±8e-5	0.905±7e-4	0.953±2e-4	0.0074±3e-5	0.846±1e-4	0.923±1e-4	0.0088±9e-5	0.903±7e-4	0.951±2e-4	0.0060±5e-5
RNNI [47]	0.823±2e-4	0.911±1e-4	0.0250±4e-5	0.886±7e-4	0.943±3e-4	0.0091±5e-5	0.909±6e-4	0.954±2e-4	0.0052±4e-5	0.926±7e-4	0.963±2e-4	0.0045±4e-5
RNN PSID [56]	0.883±1e-4	0.940±1e-4	0.0108±2e-5	0.915±3e-4	0.957±6e-5	0.0071±2e-5	0.915±7e-4	0.957±3e-4	0.0048±5e-5	0.908±3e-4	0.953±4e-4	0.0049±1e-5
LRNN [63]	0.879±7e-4	0.937±4e-4	0.0137±7e-5	0.916±2e-4	0.957±7e-5	0.0067±1e-5	0.916±8e-4	0.957±4e-4	0.0047±4e-5	0.927±3e-4	0.963±8e-5	0.0045±2e-5
TERN [47]	0.866±7e-4	0.932±4e-4	0.0247±3e-4	0.920±9e-4	0.960±2e-4	0.0062±4e-5	0.888±1e-3	0.945±7e-4	0.0067±2e-4	0.929±4e-4	0.964±2e-4	0.0043±2e-5
DCRRN [40]	0.956±3e-3	0.978±1e-3	0.0190±3e-4	0.965±3e-3	0.983±2e-3	0.0026±3e-4	0.964±1e-4	0.982±8e-5	0.0020±4e-6	0.977±2e-3	0.988±8e-4	0.0014±1e-4
GWNet [76]	0.971±5e-4	0.986±4e-4	0.0176±6e-4	0.985±3e-4	0.992±1e-4	0.0012±2e-5	0.942±2e-3	0.971±1e-3	0.0033±1e-4	0.949±2e-3	0.974±1e-3	0.0031±1e-4
AMAG	0.965±1e-3	0.982±1e-3	0.0209±1e-3	0.972±1e-3	0.986±6e-4	0.0021±8e-5	0.973±2e-3	0.986±1e-3	0.0015±1e-4	0.979±7e-4	0.990±4e-4	0.0013±4e-5
future forecasting accuracy [56]; Notably, we do not compare with RNN PSID [56] for multi-step forecasting comparison since learning behavior-related neural space in the first stage of RNN PSID requires access to all the previous steps at a given moment, which is not available in a multi-step forecasting scenario. In addition, we compared AMAG with GNN based models. In particular, GWNet: A model which intertwines CNN and GNN for spatial interaction with a learnable adjacency matrix [76], DCRNN: embedded GNN in each GRU processing step using the predefined adjacency matrix [40], GraphS4mer: Sequential learning of temporal embedding and spatial interaction with time-dependent adjacency matrices [66]. For GraphS4mer we consider two variants, GRU based temporal learning (GS4-G) and state space model based temporal learning (GS4-S). Both variants are not included in one-step forecasting comparison, considering GraphS4mer's design of learning dynamic graphs for temporal windows, which increases computation complexity and is unnecessary in one-step forecasting.
Experiment Setup. For all four datasets, we generate future neural signals given at least 5 steps as the context corresponding to the  $150\mathrm{ms}$  time window. For one-step prediction, we use GRU architecture for TE and TR, and for multi-step prediction, Transformers are used for TE and TR (multi-step forecasting performance with GRU in Appendix B.4). All the models are optimized by Adam Optimizer on Titan X GPU. The initial learning rates are  $1e - 4$  and  $5e - 4$  for non-GNN methods and GNN based methods. We include three evaluation metrics. R-squared  $(R^2)$  measures the proportion of variance in the future neuron recordings that the forecasted signal can explain. Correlation  $(Corr)$  focuses on the matching of the trend of the ground truth signal and the forecasted signal. Mean squared error  $(MSE)$  measures the L2 distance of the forecasted signal to the ground truth neuron signal. Further details are described in Appendix B.3.
Forecasting Performance. We evaluate the performance AMAG to forecast future neural activities and present the results in Table 1 and Table 2, for one-step and multi-step forecasting, respectively. The standard deviation is obtained from three runs. Compared with the best non-GNN methods, AMAG improves the  $R^2$  score by  $8.2\%$ ,  $5.0\%$ ,  $5.7\%$ , and  $5.0\%$  on Monkey M, C, B, and A datasets, respectively in Table 1. Transformer only methods, STNDT [37] and NDT [78] achieve less optimal performance on four datasets. For instance, on Monkey M, the  $R^2$  score is 0.883 for the best RNN based methods (RNN PSID) and 0.837 for the best Transformer based method (NDT). This could be due to the fact that the neural signals are continuous and future signals are primarily dependent on nearby precedents, which fits the RNN use cases better. Consequently, the advantage of the Transformer encoding long-term dependency may not be as prominent in this context. While Graph methods achieve comparable performance as shown in Table 1, AMAG shows consistent performance
(C)
Figure 4: (A): Performance evaluation of masking channels from High (H), Mid-high (M-H), Mid-low (M-L), and Low (L) groups categorized by weight of learned adjacency matrix, showing the importance of high weight channel for forecasting performance. (B): Visualization of interaction strength of target channel to the reaming channels arranged according to spatial proximity in microECoG array. (C): Examples of neural population trajectories in 2-dimensional PCA space projected from original neuron recording (left) and hidden embedding of AMAG (right).
Table 2: Multi-step forecasting results on Monkey M, C, B, and A with benchmark methods and AMAG.
Monkey M	Monkey C	Monkey B	Monkey A
R2↑	Corr ↑	MSE ↓	R2↑	Corr ↑	MSE ↓	R2↑	Corr ↑	MSE ↓	R2↑	Corr ↑	MSE ↓
LFADS [65]	0.233±4e-3	0.485±3e-3	0.0780±4e-4	0.514±3e-3	0.725±2e-3	0.0387±2e-4	0.427±4e-3	0.672±6e-3	0.0338±4e-4	0.731±2e-3	0.855±1e-3	0.0164±1e-4
STNDT [37]	0.227±3e-3	0.480±4e-3	0.0780±4e-4	0.518±1e-3	0.720±5e-4	0.0387±2e-4	0.525±4e-3	0.725±2e-3	0.0288±2e-4	0.726±9e-3	0.852±6e-3	0.0167±6e-4
LRNN [63]	0.250±3e-3	0.500±2e-3	0.0757±1e-4	0.483±4e-3	0.697±2e-3	0.0416±2e-4	0.507±6e-3	0.725±1e-3	0.0286±4e-4	0.696±6e-4	0.833±1e-3	0.0187±3e-5
RNNI [47]	0.269±6e-3	0.515±4e-3	0.0746±4e-4	0.517±9e-3	0.725±4e-3	0.0382±7e-4	0.472±9e-3	0.694±6e-3	0.0314±5e-4	0.733±2e-3	0.856±8e-4	0.0163±1e-4
TERN [47]	0.257±6e-3	0.510±4e-3	0.0729±6e-4	0.548±5e-3	0.746±6e-3	0.0358±5e-4	0.559±3e-3	0.752±2e-3	0.0265±2e-4	0.746±1e-3	0.865±4e-4	0.0154±8e-5
NDT [78]	0.287±3e-3	0.528±3e-3	0.0723±3e-4	0.565±2e-3	0.752±1e-3	0.0348±1e-4	0.575±4e-3	0.773±3e-3	0.0250±2e-4	0.756±3e-3	0.873±1e-3	0.0149±2e-4
GWNet [76]	0.272±8e-3	0.524±1e-2	0.0721±8e-4	0.606±4e-3	0.779±3e-3	0.0309±4e-4	0.588±2e-3	0.769±2e-3	0.0242±4e-4	0.724±1e-3	0.851±2e-4	0.0168±9e-5
GS4-S [66]	0.181±7e-3	0.463±2e-3	0.0792±7e-4	0.553±3e-3	0.749±3e-3	0.0350±3e-4	0.600±7e-3	0.782±1e-3	0.0236±3e-4	0.740±5e-3	0.861±3e-3	0.0158±3e-4
GS4-G [66]	0.268±3e-3	0.530±3e-3	0.0730±1e-4	0.586±7e-3	0.769±4e-3	0.0322±6e-4	0.659±3e-3	0.812±3e-3	0.0194±1e-4	0.753±8e-4	0.869±6e-4	0.0149±5e-5
DCRNN [40]	0.288±3e-3	0.545±4e-3	0.0707±7e-4	0.606±2e-3	0.782±2e-3	0.0302±2e-4	0.635±4e-3	0.797±2e-3	0.0208±3e-4	0.756±2e-3	0.870±9e-4	0.0148±1e-4
AMAG	0.331±4e-3	0.575±8e-4	0.0694±4e-4	0.657±2e-3	0.811±2e-3	0.0266±2e-4	0.665±2e-3	0.817±1e-3	0.0192±3e-4	0.763±4e-3	0.874±2e-3	0.0144±2e-4
over four datasets, compared to GWNet which has variable accuracy on different datasets. Similar trends can be observed in multi-step forecasting results. As in Table 2, Graph-based methods, including AMAG, demonstrate better learning of future neural dynamics, highlighting the potential of GNNs in this task. GS4-G performs better than GS4-S on all four datasets, indicating that GRU based temporal encoding could be more suitable for neural signal forecasting. We find that AMAG is more optimal by almost  $10\%$  in  $R^2$  on Monkey C and B when compared with non-GNN methods. Other GNN based methods, such as DCRNN, perform well, but AMAG exhibits additional improvement. For example, on Monkey C, AMAG achieves  $R^2$  score of 0.657, whereas DCRNN achieves around 0.606. Another advantage of AMAG is that it dynamically learns the pruning of unnecessary edges in the adjacency matrix, while DCRNN and variants of GraphS4mer rely on a predetermined correlation matrix to define the connected edges which is typically required to be sparse for optimal performance and handcrafted sparsification is needed. Thus, we prune the DCRNN, and GraphS4mer variants with K-Nearest Neighbour (KNN) pruning as mentioned in GraphS4mer such that each neuron is connected to half of the other neurons. Threshold pruning can be an alternative method for DCRNN, but we find that it is not optimal as the  $R^2$  on Monkey C is equal to 0.573 and 0.445 for similarity thresholds of 0.5 and 0.8, respectively. These are lower than  $R^2$  of 0.606 obtained with the KNN pruning.
In Fig. 3, we visualize the forecasted trajectory of the strongest baselines: DCRNN, GWNET, GS4-G, and AMAG. AMAG exhibits a higher closeness to the ground-truth signal (highlighted by black circles). This is consistent with behavior decoding from forecasted signals. Specifically, a linear behavior decoder is trained on ground truth neural and behavior signals (Monkey C). Then, the same decoder is used to decode behavior from AMAG and other GNN based baselines forecasted signals. Specifically, with the forecasted signals from the strongest baseline, we get decoding performance in terms of  $R^2$  as follows: 0.350 (GS4-G), 0.350 (DCRNN), 0.440 (GWNet), 0.555 (AMAG), compared to 0.656 when using ground truth neural signals, demonstrating that AMAG achieves the best behavior decoding performance (0.555) compared to other GNN based methods regardless of a gap compared to using ground truth signal. Overall, the results demonstrate the advantage of AMAG over other methods and highlight the significance of utilizing GNN based methods for forecasting neuronal signals.
Table 3: Computation complexity estimation for multi-step and one-step forecasting task. Bold marks AMAG estimates and those methods whose estimates are better than AMAG.
Mutli-step Forecasting
AMAG	LRNN	LFADS	STNDT	NDT	RNNf	TERN	DCRNN	GWNet	GS4-G
#P(M)	0.27	4.59	5.53	2.38	1.06	3.55	9.93	0.60	1.12	1.52
T(S)	9.74	2.18	3.44	4.70	5.96	7.22	8.48	11.01	12.27	13.53
M(GB)	5.74	0.12	0.15	0.19	0.04	0.11	0.34	1.69	1.53	4.48
One-step Forecasting
#P(M)	0.22	1.25	5.84	8.31	3.70	3.55	9.93	0.15	1.12	N/A
T(S)	9.90	1.97	3.29	4.61	5.94	7.26	8.58	11.22	12.54	N/A
M(GB)	5.36	0.04	0.15	0.43	0.10	0.11	0.34	1.05	1.53	N/A
Investigating Adjacency Matrix Learned by AMAG. We hypothesize that the learned adjacency matrix in AMAG indicates interaction strength between channels and channels with high interaction strength are vital in forecasting. To verify the assumption, we group the channels into High (H), Mid-high (M-H), Mid-low (M-L), and Low (L) according to connection strength in  $Aa$ . We estimate their importance as the degree of performance drop when channels in each group are masked in AMAG. As the testing of the generality of the channel importance, we performed the same procedure in NDT. The forecasting performance  $(R^2)$  of NDT and AMAG with masked channels as input are shown in Fig. 4 (A). Generally, masking of channels consistently leads to performance drop, particularly with more significant degradation of performance in both AMAG and NDT when masking happens in the H group. This confirms our assumption that the weights learned in adjacency matrices provide informative insights into channel importance. Additionally, we investigate the relationship between learned weights and channel proximity. Figure 4 (B) illustrates the projection of the interaction strength of a target channel (T) onto other channels in the space organized by the  $\mu$ ECoG arrangement. We observe that target channels exhibit stronger interactions with surrounding channels, which aligns with the properties of  $\mu$ ECoG arrays. These experiments collectively demonstrate that the learned adjacency matrix values are meaningful in terms of identifying important channels, and the learned strength of connections aligns with the spatial proximity among ECoG arrays.
Neuron Trajectory Alignment. Apart from learning meaningful adjacency matrices, it is also advantageous to study neural trajectories AMAG has learned during the forecasting task. For that purpose, we project the original neural signals and their hidden embeddings from AMAG into the 2D space with PCA in Fig. 4(C). We find that, while neural trajectories, projected from the original neural signals, are tangled with random start and endpoints, with the projection from the hidden embeddings of AMAG, the neural trajectories are disentangled for the corresponding time-step and exhibit a circular structure similar to [53]. This illustration indicates that AMAG learns informative features that can be visualized as coherent and interpretable neural trajectories in 2D space.
Computation Complexity. We additionally compare in Table 3 the number of parameters  $\# \mathrm{P}(\mathbf{M})$ , training time per epoch T(S), and maximum memory cost during training M(GB). We observe that GNN based methods use much fewer parameters than non-GNN methods, with AMAG consistently using a minimal or comparable (DCRNN for one-step forecasting) number of parameters when compared to graph model baselines (DCRNN, GWNet, GS4-G). This demonstrates that the source of the contribution to AMAG performance is its architecture rather than the number of parameters. The graph architecture of AMAG provides the topological constraint across channels. The absence of the topological constraint in non-GNN based methods could result in overfitting with limited training samples. Adding regularization terms, e.g., weight decay could help, but will also limit models' capacity. We demonstrate it in Fig. 3 (B). Adding weight decay to NDT improves performance (orange triangle vs purple circle), however, as the attention dimension increases, the effect of regularization diminishes, and for dim>1024, weight decay does not contribute to further improvement.
Here we examine how each component in AMAG contributes to forecasting. Specifically, we consider the following ablated versions of AMAG: (1) removing self-connection in AMAG (amAG), (2) removing both the "Add" and "Modulator" module (-AG), (3) Only sample dependent "Adaptor" is kept in "Add" Module ( $\tilde{a}\mathrm{MAG}$ ), (4) Adjacency matrix is not learnable (AMG), (5) The "Add" is ablated from the graph operation (-MAG), (6) "Modulator" is removed (AAG), (7) The "Adaptor" in excluded from the "Add" (aMAG). The results are shown in Table.4.
Figure 5: Visualization of the adjacency matrix of Add module  $(A_{a})$  and Modulator module  $(A_{m})$  learned by AMAG performing one-step forecasting (left) and multi-step forecasting (right) with correlation initialization (top) and random initialization (bottom).
We find that removing Add or Modulator modules can cause the most significant forecasting accuracy drop  $(-AG)$  as well as the self-connection part in the graph (amAG). When keeping sample-dependent Adaptor term only (aMAG) in the Add module, the performance is very close to ablating the Add module (-MAG). The sample-dependent Adaptor term appears to be less important compared to other components but still reduces the accuracy if being ablated from AMAG (aMAG).
Table 4: Comparison between the Ablated versions of AMAG.
R2↑	CORR↑	MSE↓
-AG	0.424±1e-3	0.650±1e-3	0.0456±1e-4
amAG	0.427±6e-3	0.652±4e-3	0.0453±5e-4
-MAG	0.647±2e-3	0.805±6e-4	0.0274±2e-4
ãMAG	0.648±1e-3	0.806±4e-4	0.0273±6e-5
A-AG	0.652±9e-4	0.807±3e-4	0.0268±1e-4
aMAG	0.655±2e-3	0.810±1e-3	0.0269±2e-4
AM-G Rand Init)	0.575±2e-2	0.767±7e-3	0.0329±1e-3
AM-G (Corr Init)	0.617±2e-4	0.786±7e-4	0.0296±2e-5
AMAG (Rand Init)	0.652±1e-3	0.807±8e-4	0.0270±1e-4
AMAG (Corr Init)	0.657±2e-3	0.811±2e-3	0.0266±2e-4
In addition, we investigate the effect of initialization (random versus correlation) when the adjacency matrices are not adaptable (AM-G), specifically, AM-G (Rand Init 1), AM-G (Rand Init 2), AM-G (Corr Init). The forecasting accuracy of AM-G compared to the other two variants demonstrates the importance of correlation initialization. In particular, random and correlation initialization achieve similar performance when the adjacency matrices are adaptable, while random initialization leads to a less stable training process, see Appendix B.4. In Fig. 5, we visualized the learned adjacency matrix with correlation initialization (top) and random initialization (bottom) for one-step (left) and multi-step forecasting (right). Globally, the learned adjacency matrices show similar connection patterns, while locally, the connection can be different with different initialization matrices. This indicates that regardless of the initialization method, AMAG can converge to graph patterns with globally similar connections while multiple local connections are available to achieve similar forecasting performance.
In conclusion, we propose AMAG for modeling spatial-temporal interaction between neural signals via forecasting future neuron activity. Experiments with both synthetic and neural signals show potential for discovering underlying neural signal interactions. The learned neural features by AMAG appear to disentangle in the PCA space. While adjacency matrix recovery is promising as training AMAG with one-step forecasting, more animal experiments could be done together with AMAG to understand the biological meaning of the learned adjacency matrix. It is worth noting that the memory requirement for GNN training increases as the number of neural recording channels increases. For future development, improvements to adapt the model for multi-region, multi-channel recordings could be considered.
Authors acknowledge the support in part by A3D3 National Science Foundation grant OAC-2117997 (ES,JL,TL,ALO), the Simons Foundation (LS,PR,ALO), the Washington Research Foundation Fund (ES), Departments of Electrical Computer Engineering (ES,JL,TL), Applied Mathematics (ES). Authors are thankful to the Center of Computational Neuroscience, and the eScience Center at the University of Washington. Authors would like to thank Rahul Biswas for discussions of synthetic dataset generation, Pan Li for discussions regarding GNNs, and the anonymous reviewers for insightful feedback.
7 Supplementary Information and Data
Supplementary information is available in the Appendix. Additional information and data will be made available through the project page https://github.com/shlizee/AMAG.