SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[1] Larry F Abbott. "Lapicque's introduction of the integrate-and-fire model neuron (1907)". In: Brain research bulletin 50.5-6 (1999), pp. 303-304.
[2] Evan Archer, Il Memming Park, Lars Buesing, John Cunningham, and Liam Paninski. "Black box variational inference for state space models". In: arXiv preprint arXiv:1511.07367 (2015).
[3] Navid Ayoobi and Elnaz Banan Sadeghian. "A self-paced BCI system with low latency for motor imagery onset detection based on time series prediction paradigm". In: arXiv preprint arXiv:2204.05450 (2022).
[4] Edoardo Balzani, Jean Paul Noel, Pedro Herrero-Vidal, Dora E Angelaki, and Cristina Savin. "A probabilistic framework for task-aligned intra- and inter-area neural manifold estimation". In: arXiv preprint arXiv:2209.02816 (2022).
[5] Rahul Biswas and Eli Shlizerman. "Statistical perspective on functional and causal neural connectomics: A comparative study". In: Frontiers in Systems Neuroscience 16 (2022), p. 817962.
[6] Rahul Biswas and Eli Shlizerman. "Statistical perspective on functional and causal neural connectomics: The Time-Aware PC algorithm". In: PLOS Computational Biology 18.11 (2022), e1010653.
[7] Martin Bjerke, Lukas Schott, Kristopher T Jensen, Claudia Battistin, David A Klindt, and Benjamin A Dunn. "Understanding Neural Coding on Latent Manifolds by Sharing Features and Dividing Ensembles". In: arXiv preprint arXiv:2210.03155 (2022).
[8] Anastasia Borovykh, Sander Bohte, and Cornelis W Oosterlee. "Conditional time series forecasting with convolutional neural networks". In: arXiv preprint arXiv:1703.04691 (2017).
[9] Nicolas Brunel. "Dynamics of sparsely connected networks of excitatory and inhibitory spiking neurons". In: Journal of computational neuroscience 8 (2000), pp. 183-208.
[10] Luis Carrillo-Reid, Shuting Han, Dark O'Neil, Ekaterina Taralova, Tony Jebara, and Rafael Yuste. "Identification of pattern completion neurons in neuronal ensembles using probabilistic graphical models". In: Journal of Neuroscience 41.41 (2021), pp. 8577-8588.
[11] Kun Chen, Tianli Xie, Li Ma, Andrew E Hudson, Qingsong Ai, and Quan Liu. "A Two-Stream Graph Convolutional Network Based on Brain Connectivity for Anesthetized States Analysis". In: IEEE Transactions on Neural Systems and Rehabilitation Engineering 30 (2022), pp. 2077-2087.
[12] Yuqian Chen, Fan Zhang, Leo R Zekelman, Tengfei Xue, Chaoyi Zhang, Yang Song, Nikos Makris, Yogesh Rathi, Weidong Cai, and Lauren J O'Donnell. "TractGraphCNN: anatomically informed graph CNN for classification using diffusion MRI tractography". In: arXiv preprint arXiv:2301.01911 (2023).
[13] Kyunghyun Cho, Bart Van Merrienboer, Dzmitry Bahdanau, and Yoshua Bengio. "On the properties of neural machine translation: Encoder-decoder approaches". In: arXiv preprint arXiv:1409.1259 (2014).
[14] Rainer Dahlhaus. "Graphical interaction models for multivariate time series". In: Metrika 51 (2000), pp. 157-172.
[72] Francis R Willett, Donald T Avansino, Leigh R Hochberg, Jaimie M Henderson, and Krishna V Shenoy. "High-performance brain-to-text communication via handwriting". In: Nature 593.7858 (2021), pp. 249-254.
[73] Anqi Wu, Nicholas A Roy, Stephen Keeley, and Jonathan W Pillow. "Gaussian process based nonlinear latent structure discovery in multivariate spike train data". In: Advances in neural information processing systems 30 (2017).
[74] Neo Wu, Bradley Green, Xue Ben, and Shawn O'Banion. "Deep transformer models for time series forecasting: The influenza prevalence case". In: arXiv preprint arXiv:2001.08317 (2020).
[75] Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip. “A comprehensive survey on graph neural networks”. In: IEEE transactions on neural networks and learning systems 32.1 (2020), pp. 4–24.
[76] Zonghan Wu, Shirui Pan, Guodong Long, Jing Jiang, and Chengqi Zhang. "Graph wavenet for deep spatial-temporal graph modeling". In: arXiv preprint arXiv:1906.00121 (2019).
[77] Jiexia Ye, Juanjuan Zhao, Kejiang Ye, and Chengzhong Xu. "How to build a graph-based deep learning architecture in traffic domain: A survey". In: IEEE Transactions on Intelligent Transportation Systems 23.5 (2020), pp. 3904-3924.
[78] Joel Ye and Chethan Pandarinath. “Representation learning for neural population activity with Neural Data Transformers”. en. In: Neurons, Behavior, Data analysis, and Theory (Aug. 2021). ISSN: 2690-2664. DOI: 10.51628/001c.27358. URL: https://nbdt.scholasticahq.com/article/27358-representation-learning-for-neural-population-activity-with-neural-data-transformers (visited on 08/29/2021).
[79] Byron M Yu, John P Cunningham, Gopal Santhanam, Stephen Ryu, Krishna V Shenoy, and Maneesh Sahani. "Gaussian-process factor analysis for low-dimensional single-trial analysis of neural population activity". In: Advances in neural information processing systems 21 (2008).
[80] Ailing Zeng, Muxi Chen, Lei Zhang, and Qiang Xu. "Are transformers effective for time series forecasting?" In: arXiv preprint arXiv:2205.13504 (2022).
[81] Yuan Zhao and II Memming Park. "Variational latent gaussian process for recovering single-trial dynamics from population spike trains". In: Neural computation 29.5 (2017), pp. 1293-1316.
[82] Yuan Zhao and Il Memming Park. "Variational online learning of neural dynamics". In: Frontiers in computational neuroscience 14 (2020), p. 71.
[83] Ding Zhou and Xue-Xin Wei. "Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-VAE". In: Advances in Neural Information Processing Systems 33 (2020), pp. 7234-7247.
[84] Feng Zhu, Andrew Sedler, Harrison A Grier, Nauman Ahad, Mark Davenport, Matthew Kaufman, Andrea Giovannucci, and Chethan Pandarinath. "Deep inference of latent dynamics with spatio-temporal super-resolution using selective backpropagation through time". In: Advances in Neural Information Processing Systems 34 (2021), pp. 2331-2345.
Appendix A Methods Details
In the following, we provide additional details regarding the initialization of the adjacency matrix, and temporal modeling (TE and TR components of AMAG).
Across the paper, we define the input to the model as  $\mathbf{X} \in \mathbb{R}^{T \times C \times D}$ , where  $T$  denotes the total number of time steps,  $C$  represents the number of channels in the neural recordings, and  $D$  signifies the feature dimension. To refer specifically to the channel  $v$  or the time step  $t$  of input neural signals, we use the notations  $\mathbf{X}^{(v)} \in \mathbb{R}^{T \times D}$  or  $\mathbf{X}_t \in \mathbb{R}^{C \times D}$ , respectively. To distinguish different data structures, matrices are denoted by bold capital letters, such as  $\mathbf{X}$ ,  $\mathbf{X}^{(v)}$ ,  $\mathbf{X}_t$ , while unbolded capital letters are used for vectors and scalars (e.g.,  $X_t^{(v)} \in \mathbb{R}^D$ ).  $FC$  dentoes a single-layer fully connected neural network,  $MLP$  denotes a multi-layer fully connected neural network, and  $LN$  denotes the layer normalization technique.  $A_a$  and  $A_m$  are the adjacency matrices in Add and Modulator modules.
A.1 Initialization of Adjacency Matrix
The adjacency matrices in AMAG are learnable parameters. We find that learning the adjacency matrix from scratch can be unstable. Hence we propose to initialize the adjacency matrix based on the correlation of each channel to stabilize the training. We flatten the features of channel  $v$  and  $u$ , such that  $X^{(u)}, X^{(v)} \in \mathbb{R}^{TD}$ . Then the correlation between channel  $u$  and channel  $v$  is
$$
C o r r (u, v) = \frac {X ^ {(u)} X ^ {(v) T}}{\sqrt {\left| X ^ {(u)} \right| \left| X ^ {(v)} \right|}}.
$$
We then initialize each element in the adjacency matrices  $A_{a}^{(u,v)}$  and  $A_{m}^{(u,v)}$  as  $Corr(u,v)$ .
A.2 Recursive Time Series Forecasting with TE and TR
Since Recurrent Neural Networks (RNN) are common in modeling continuous time-series signals, we employ a variant of RNN, Gated Recurrent Unit (GRU) [13], to facilitate the temporal encoding (TE) and readout (TR) modules. The TE module consists of a single-layer GRU that takes input from channel  $v$  and computes the corresponding temporal embedding features  $h_t^{(v)} \in \mathbb{R}^d$  for each step  $t$ . These embedding features capture the underlying temporal dynamics of neural activity. To enable the exchange of information between the target channel  $v$  and its neighboring channels  $N(v)$ , specified by the adjacency matrices  $A_a$  and  $A_m$ , the spatial interaction (SI) module updates the embedding features. This interaction among channels allows for modeling of spatial dependencies.
The TR module consists of a single-layer GRU and a fully connected  $(FC)$  layer. The module generates future neural signals. This pipeline of TE, SI and TR modules is shown below
$$
\boldsymbol {h} _ {t} ^ {(v)} = G R U _ {T E} \left(X _ {t} ^ {(v)}, \boldsymbol {h} _ {t - 1} ^ {(v)}\right), \quad \boldsymbol {h} _ {0} ^ {(v)} = \boldsymbol {0}
$$
$$
\boldsymbol {z} _ {t} ^ {(v)} = S I (\boldsymbol {h} _ {t} ^ {(v)}, \{\boldsymbol {h} _ {t} ^ {(u)} | u \in \boldsymbol {N} (v) \}, \boldsymbol {A} _ {a}, \boldsymbol {A} _ {m})
$$
$$
\boldsymbol {r} _ {t} ^ {(v)} = G R U _ {T R} (\boldsymbol {z} _ {t} ^ {(v)}, \boldsymbol {r} _ {t - 1} ^ {(v)})
$$
$$
X _ {t + 1} ^ {(v)} = F C (\pmb {r} _ {t} ^ {(v)})
$$
A.3 Simultaneous Temporal Forecasting with TE and TR
The GRUs employed in TE and TR modules are particularly well-suited for recursive forecasting of future neural signals since they leverage the previously generated one-step signals to predict the next steps. In the multi-step forecasting scenario, however, where predictions for multiple future time steps are required, using GRUs in an autoregressive manner could lead to accumulation of errors. To address this issue and explore an alternative approach, we consider Transformers as an alternative to GRU. Transformers [68], simultaneously generate all future neural signals and could avoid the accumulation of errors. In this case, both TE and TR are implemented as decoder-based Transformers. The TR module predicts multiple future steps from a given initial time step  $t$  up to the target time step  $T$ , based on the available neural recordings prior to  $t$ .
Before feeding neural signals into the Transformer model, several pre-processing steps are applied. These steps include a masking layer, a linear layer, and a positional embedding layer, which sequentially transform the neural signals into embedding space denoted as
$$
\boldsymbol {E} = e m b (\boldsymbol {X}) = F C (m a s k (\boldsymbol {X})) + P E.
$$
The masking layer (mask) masks out all the future inputs, such that  $X_{t:T}$  is constant during multi-step forecasting. The linear layer (FC) transforms the input neural signals from their original dimension  $D$  to an embedding space of dimension  $d$  using a linear transformation  $(FC: \mathbb{R}^D \Rightarrow \mathbb{R}^d)$ . Additionally, positional encoding (PE) is applied to incorporate positional information into the embedding space. The positional embedding is defined as  $PE(pos, 2i) = sin(pos/10000^{2i/d})$ ,  $PE(pos, 2i + 1) = sin(pos/10000^{2i + 1/d_{model}})$ , as described in [68]. Here pos corresponds to the position in the temporal space, and  $i$  represents the position in the feature space. It is worth noting that channels of neural recordings are treated as feature dimensions when incorporating the positional embedding, ensuring that the Transformer can distinguish between different channels.
The embedded signal  $\pmb{E}$  is then mapped to  $Q_{e}^{(v)}, K_{e}^{(v)}$ , and  $\pmb{V}_{e}^{(v)}$  by weight matrix  $\pmb{W}_{q}, \pmb{W}_{k}, \pmb{W}_{v} \in \mathbb{R}^{d \times d}$  as Query, Key and Value in Transformer model. Specifically,  $\pmb{Q}_{e}^{(v)} = \pmb{E}^{(v)}\pmb{W}_{q}, \pmb{K}_{e}^{(v)} = \pmb{E}^{(v)}\pmb{W}_{k}, \pmb{V}_{e}^{(v)} = \pmb{E}^{(v)}\pmb{W}_{v}$ . The outputs from the temporal encoding can be written as
$$
\boldsymbol {A} \boldsymbol {T} ^ {(v)} = S o f t m a x (\frac {\boldsymbol {Q} _ {e} ^ {(v)} (\boldsymbol {K} _ {e} ^ {(v)}) ^ {T}}{\sqrt {d}}) \boldsymbol {V} _ {e} ^ {(v)},
$$
$$
\boldsymbol {H} ^ {(v)} = L N \left(M L P \left(\boldsymbol {E} ^ {(v)} + \boldsymbol {A T T} ^ {(v)}\right)\right),
$$
$$
\boldsymbol {H} ^ {(v)} = \left[ \boldsymbol {h} _ {1} ^ {(v)}, \dots \boldsymbol {h} _ {T} ^ {(v)} \right], \boldsymbol {H} ^ {(v)} \in \mathbb {R} ^ {T \times d}
$$
Here we assume a single head in the Transformer, which could be extended to the multi-head situation as well. As in one-step forecasting, we update the feature with spatial interaction
$$
\boldsymbol {z} _ {t} ^ {(v)} = S I (\boldsymbol {h} _ {t} ^ {(v)}, \{\boldsymbol {h} _ {t} ^ {(u)} | u \in \boldsymbol {N} (v) \}, \boldsymbol {A} _ {a}, \boldsymbol {A} _ {m}),
$$
$$
\boldsymbol {Z} ^ {(v)} = [ \boldsymbol {z} _ {1} ^ {(v)}, \dots , \boldsymbol {z} _ {T} ^ {(v)} ].
$$
TR module is similar to TE, with the difference that instead of receiving embedded signal as input, TR receives  $Z^{(v)}$  as input and generates  $\hat{X}^{(v)}$  with corresponding  $Q_{z}^{(v)} = Z^{(v)}W_{q}$ ,  $K_{z}^{(v)} = Z^{(v)}W_{k}$ ,  $V_{z}^{(v)} = Z^{(v)}W_{v}$ . To clarify, we use the same notation for  $W_{k}, W_{q}, W_{v}$ , but they are different sets of parameters in TE and TR. While we describe the process for generating  $H^{(v)}$  and  $Z^{(v)}$  for current channel  $v$ , it can be adapted to other channels. The parameters in TE and TR are shared across channels.
Appendix B Experiments
B.1 Synthetic Dataset
Experiments with the Synthetic Dataset. In our experimentation with the synthetic dataset, we apply AMAG to learn four specific tasks:
- Reconstruction: In this task, the model aims to reconstruct neural signals without any masking or missing information. The entire input sequence is given as input to the model.
- Reconstruction with masking of the input: In this task, the model is trained to reconstruct neural signals while dealing with masked or missing information. The channels are masked in a blocked manner, where five consecutive channels at five consecutive time steps are randomly masked. The first  $t$  steps are considered as the context and are not masked.
- Multi-step forecasting: In this task, the model is trained to predict multiple future steps from a given initial time step  $t$  up to the target time step  $T$ .
- One-step forecasting: This task focuses on predicting the next step or a single future step based on the available input information.
Figure 6: Adjacency matrices learned by AMAG visualized as colormaps and arranged by their F1 score (from left to right is from lower to higher), with the rightmost map indicating Ground Truth. Yellow color indicates full recovery of ground truth value (100% accuracy) while green-blue colors indicate recovery with lower accuracy. In particular, we show outcome adjacency matrices that AMAG recovers when solving left-right: the reconstruction task, the reconstruction of masked input task, multi-step forecasting task, one-step forecasting task, and ground truth.
To ensure a fair comparison between these tasks, we employ the same model structure, specifically using GRUs for both the TE and TR modules in all four cases, and compute the loss term for each task over the time steps from  $t$  to the end of the sequence.
Additional Results with the Synthetic Dataset. In addition to the results in the main paper, we further assess AMAG ability to recover the learned ground truth adjacency matrix using the synthetic dataset, which consists of 40 channels with a connection ratio of 0.2. We visualize the learned adjacency matrices in Fig. 6, comparing the results obtained from AMAG computing various tasks: reconstruction, reconstruction with masking, multi-step forecasting, and one-step forecasting. The learned adjacency matrices have continuous values describing the strength of connections between channels. The color in the figure indicates the connection strength. We keep the top eight  $(40 \times 0.2)$  connections as the recovery of the ground truth connections for each channel. To evaluate the accuracy of the discovered adjacency matrices, we use the F1 score as a metric that measures the harmonic mean between precision and recall, providing a comprehensive assessment of the matrix recovery. F1 scores for the four considered tasks are ordered from low to high and are arranged from left to right in Fig. 6. A higher F1 score indicates a higher degree of accuracy in capturing the ground truth adjacency matrix, highlighting the effectiveness of each method in this discovery process.
B.2 Neural Recordings Datasets
Details of Monkey A and B Datasets. As described in the main paper, we collected two new datasets from monkey Affogato (A), and Beignet (B) with  $\mu$ ECoG implanted in motor cortices, apply and compare AMAG and other methods with these recordings. All procedures were approved by the University of Washington Institutional Animal Care and Use Committee. The  $\mu$ ECoG covers five different subregions (Primary Motor Cortex (M1), Premotor Cortex (PM), Frontal Eye Field (FEF), Supplementary Motor Area (SMA), and Dorsolateral Prefrontal Cortex (DLPFC)) with 239 effective electrodes which have been previously used in [67]. During data collection, the subjects performed a reaching task towards one of eight different directions. The subject's hand position was tracked in
real-time using a marker-based camera system (Optitrak) and mapped to control the position of a cursor with which they performed the point-to-point movements. The collected neural signals can be segmented into trials based on the stage of the behavioral task. Within each trial, the subjects were trained first to reach the center target, followed by the appearance of the surrounding target. Subsequently, a delay period ranging from approximately 50ms to 150ms was introduced. After the delay period, the subjects moved from the center target to the surrounding target. A trial is considered successful if the subjects successfully reach the surrounding target, resulting in a reward. For our study, we include the neural signals recorded during the 600ms following the appearance of the surrounding target as our samples. Specifically, we treat the first 150ms of neural signals as a preparatory context for forecasting the future 450ms of neural signals. Consequently, trials with a duration shorter than 600ms after the appearance of the surrounding target are excluded from the study. In our study, we use data collected from monkeys A and B across 13 different sessions. These sessions are combined to form a comprehensive dataset for analysis. The dataset is split into  $80\%$  for training,  $10\%$  for validation, and  $10\%$  for testing. Specifically, dataset A contains 985 training samples, 122 validation, and 122 testing samples, while dataset B contains 700 training samples, 87 validation samples, and 87 testing samples. In experiments, we include all 239 electrodes from Monkey A and 87 electrodes, specifically from the M1 region of Monkey B. The neural time series for all train and test samples from monkeys A and B used in our analyses, along with meta-data of relative electrode locations within the array, will be shared to facilitate benchmarking of alternative algorithms.
Details of Monkey M and C Datasets. For dataset C and M [20, 64], we use neural signals recorded from the M1 of the two monkeys, including 96 channels. In these two datasets, we align the trials based on the movement onset, i.e., we include neural signals from 300ms before the movement onset and 450ms after the movement onset. Similarly, we split the datasets into  $80\%$  training data,  $10\%$  validation data, and  $10\%$  test data, collected in 6 sessions for Monkey C and 7 sessions for Monkey M. The number of training, validation and testing samples for Monkey C is 1285, 160, and 160, respectively. For dataset M, we have 1023 training samples, 127 validation, and 127 testing samples. All the trials have at least 450ms after the target onset.
Details of Preprocessing. 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}$ . 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]. All four datasets follow the normalization procedures to ensure consistency and comparability. Specifically, we normalize the data such that values for each channel within the range defined by mean minus four times the standard deviation  $(mean - 4 \times std)$  and mean plus four times the standard deviation  $(mean + 4 \times std)$  are linearly scaled to interval  $[-1, 1]$ . The mean (mean) and standard deviation (std) are computed for each channel based on the training set.
B.3 Hyperparameters and Additional Model Details
To optimize the performance of models, we conducted a hyperparameter search by exploring different values for the learning rate, hidden size, and weight decay. Specifically, we considered learning rates of  $10^{-3}$ ,  $5 \times 10^{-4}$ , and  $10^{-4}$ , hidden sizes of 64, 128, 512, 1024, and 2048, and weight decay values of 0,  $10^{-4}$ ,  $10^{-5}$ ,  $10^{-6}$ , and  $10^{-7}$ . We used the Adam optimizer to update the model parameters during training. We evaluate the model performance in terms of the  $R^2$ , correlation coefficient (Corr), and mean squared error (MSE) metrics on the validation set every 10 training epochs. The reported values in the tables of the main paper represent the testing set performance over the course of 500 epochs for most cases when the validation set achieves best performance. However, for the one-step forecasting task using the AMAG model and the STNDT-based models for both one-step and multi-step forecasting, we trained the models for 1000 epochs to ensure sufficient convergence and performance evaluation.
Details of Models Used for Comparison
- LFADS [65]. The LFADS model is a non-causal bi-directional GRU. We constrain the RNNs used in the LFADS to be uni-directional. Furthermore, since forecasting tasks typically require higher-dimensional features compared to reconstruction tasks, we set the factor dimension and the inferred input dimension to 256 for datasets A and B, and 64 for datasets M and C. The dimension
of the encoder, controller, and generator in the model is 512. The model is trained with an initialized learning rate of  $10^{-4}$  reducing by 0.95 every 50 epochs and a weight decay rate of  $10^{-4}$ .
$^{1}$ TERN is also denoted as Neural RoBERTa, RoBERTa in short [32]
Table 5: Multi-step forecasting performance of AMAG using GRU as the components for TE and TR modules (AMAG-G) versus AMAG using Transformer as the components for TE and TR modules (AMAG-T).
Figure 7: Visualization of  $R^2$  score during the training and testing with the initialization of correlation matrix (solid curves) and random initialization (dashed curves). The correlation matrix initialized adjacency matrix stabilizes the training process.
[48]. The MLP encompasses four sequential fully connected layers whose input feature dimensions are  $64 \times t$ ,  $64 \times 2$ ,  $64 \times 4$ , and 64 respectively where  $t$  is the temporal context length. There is no weight decay on M and C datasets,  $10^{-5}$  decay rate on A and B datasets. The initial learning rate is  $5 \times 10^{-4}$  decaying by 0.95 every 50 epochs.
GRU and Transformer Comparison as Base Components for TE and TR for Multi-step Forecasting. We compare how different temporal learning components for TE and TR change model performance in multi-step forecasting. We show the results in Table 5. We find that by incorporating the Transformer for temporal learning, AMAG could potentially achieve slightly better performance. We thereby use Transformer as the main components of TE and TR for multi-step forecasting.