AMAG: Additive, Multiplicative and Adaptive Graph Neural Network For Forecasting Neuron Activity
Authors: Jingyuan Li
Leo Scholl
Trung Le
Pavithra Rajeswaran
Amy L Orsborn
Eli Shlizerman
Venue: NeurIPS 2023 poster (2023)
Keywords: Neuroscience and Cognitive Science
Neural Activity Forecasting
Graph Neural Network
TL;DR: In this work, we emphasize the importance of modeling neural population dynamics via forecasting tasks and aim to improve the forecasting performance with Graph Neural Networks

Abstract: Latent Variable Models (LVMs) propose to model the dynamics of neural populations by capturing low-dimensional structures that represent features involved in neural activity. Recent LVMs are based on deep learning methodology where a deep neural network is trained to reconstruct the same neural activity given as input and as a result to build the latent representation. Without taking past or future activity into account such a task is non-causal. In contrast, the task of forecasting neural activity based on given input extends the reconstruction task. LVMs that are trained on such a task could potentially capture temporal causality constraints within its latent representation. Forecasting has received less attention than reconstruction due to recording challenges such as limited neural measurements and trials. In this work, we address modeling neural population dynamics via the forecasting task and improve forecasting performance by including a prior, which consists of pairwise neural unit interaction as a multivariate dynamic system. Our proposed model---Additive, Multiplicative, and Adaptive Graph Neural Network (AMAG)---leverages additive and multiplicative message-passing operations analogous to the interactions in neuronal systems and adaptively learns the interaction among neural units to forecast their future activity. We demonstrate the advantage of AMAG compared to non-GNN based methods on synthetic data and multiple modalities of neural recordings (field potentials from penetrating electrodes or surface-level micro-electrocorticography) from four rhesus macaques. Our results show the ability of AMAG to recover ground truth spatial interactions and yield estimation for future dynamics of the neural population.

AMAG: Additive, Multiplicative and Adaptive Graph Neural Network For Forecasting Neural Activity
Jingyuan Li', Leo Scholl', Trung Le', Pavithra Rajeswaran' Amy Orsborn' 1:2 Eli Shlizerman 1.3
1. Department of Electrical & Computer Engineering,
2. Department of Bioengineering.
3. Department of Applied Mathematics
University of Washington, Seattle, WA
{jingyli6, lscholl, tle45, pavir, aorsborn, shlizee Buw.eeeeuwwweee
Abstract
Latent Variable Models (LVMs) propose model the dynamics of neural popu- lations by capturing low-dimensional structures that represent features involved in neural activity. Recent LVMs are based deep learning methodology where deep neural network trained to reconstruct the same neural activity given as input and result to build the latent representation. Without taking past future activ- ity into account such a task is non-causal. In contrast, the task of forecasting neural activity based on given input extends the reconstruction task. LVMs that trained on such task could potentially capture temporal causality constraints within its latent representation. Forecasting has received less attention than reconstruction due to recording challenges such as limited neural measurements and trials. In this work, we address modeling neural population dynamics via the forecasting task and improve forecasting performance by including a prior, which consists of pairwise neural unit interaction as a multivariate dynamic system. Our proposed model-Additinn Multiplicative, and Adaptive Graph Neural Network (AMAG) leverages additive and multiplicative message-passing operations analogous to the interactions in neuronal systems and adaptively learns the interaction among neural units to forecast their future activity. We demonstrate the advantage AMAG compared to Non-GNN based methods synthetic data and multiple modalities neural recordings (field potentials from penetrating electrodes or surface-level from four rhesus macaques. Our results show the ability AMAG to recover ground truth spatial interactions and yield estimation for future dynamics of the neural population.
Introduction
The ability to encode, decode, and interpret neural activity could facilitate novel applications targeted
toward neural activity, such as Brain-Compuss Interfaces (BCIs) [5] 50 52 28 49 A promising
methodology for neural activity encoding and decoding is the of artificial Deep Neural Networks
(DNN) [53 84 78). These approaches encode neural activity into latent neural representations, facilitating decoding and interpretation of neural signals and correlating the activity with other modalities, e.g., brain-to-text translation (72..
Most existing DNN architectures, such as Latent Factor Analysis via Dynamical Systems (LFADS) and variants, infer latent dynamies from single-trial population activities by encoding the neural time series into summarized representation from which the original time series could be regenerated, a procedure known as reconstruction task (53 65 51]. While effective in capturing low-dimensional latent neural manifolds, the task does not consider causal aspects of neural dynamics. In contrast,
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

(A)
MLP
Adaptor
Add
add
Modulato
Neural Activity Forecasting
SI
Figure 1: AMAG Overview. (A) The forecasting task and AMAG architecture with temporal Encoding (TE), Spatial Interaction (SI), and Temporal Readout (TR) modules address the task. (B) Components of SI module.
the forecasting task, which aims to predict future activity based on past data, adhering to temporal causality constraints, could provide further insight into neural representations that closer to representations employed by the brain 82 practical advantages since the signals that are predicted
55 56 18). In addition, learning to forecast offers help to reduce latency in real-time neural decoding applications (e.g., BCIs) or serve as a reference to detect real-time deviations of current recordings compared to earlier ones, indicating dynamic changes in the neuronal system or equipment malfunction (3]-
Neural activity forecasting is inherently more challenging than reconstruction due the requirement of a larger number of recorded neural signals, which are typically limited, and sensitivity to noise present in the signals. Incorporating priors into model structure design can be beneficial to ensure that instead learning noisy information or overfitting the training data, the model learns features for forecasting future activity that generalizable. A well-known prior is the spatial interaction among neurons, involving both addition [27] 9] and multiplication processes [19, which could be modeled with Hodgkin-Huxles probabilistic graphical models to learn interactions [14
Leaky Integrate-and-Finn [27
11. Previous works have utilized
10 with a focus on capturing the neuron interaction itself rather than addressing downstream forecasting tasks. Models that learn the downstream task by considering spatial interaction have been proposed, such as Spatial Temporal Neural Data Transformer (STNDT) and Embedded Interaction Transformer (EIT), where the interaction between neurons is modeled using an attention mechanism [37] [4 68). Graph Neural Networks Networks (GNNs), designed handle graph-structured data, are an alternative to the attention mechanism to model neuron interactions 75).
42
69
36
In this study, propose GNN with additive and multiplicative peessae-passiingg operations, called AMAG, for neuron signal forecasting. As shown in Fig. (A), AMAG includes three modules: temporal encoding (TE), spatial interaction (SI), and temporal readout (TR). TE and TR model temporal dynamics for each recording channel independently. SI models the interaction between recording channels by leveraging the additive and multiplicative message-pane paaggge-pssnnggggg operations. We test the model on synthetic datasets and show the ability of AMAG to recover the ground truth interaction between channels. In addition, we demonstrate the effectiveness of the proposed model by testing it on two neural local field potential recording types (each with datasets from two animals). The first type is publicly available dataset recorded by Utah array (penetrating probes) and the second type is a novel dataset of cilcro-electrccortcoogoeo surface. These four recordings are from the motor cortices of rhesus macaque monkeys performing reaching. We compare our model with other approaches, including GNN or non-GNN based. These computational evaluations demonstrate that AMAG outperforms other models regarding forecasting accuracy on neural recordings. We summarize the main contributions of our work below:
(μECooG array recordings from the cortical
Modeling tentalltemmpoal relationship. We present graph neural network (GNN), called AMAG, that can forecast neuron activity by using sample-dependen additive and multiplicative message- passing operations with a learnable adjacency matrix.
Learning underlying spatial interaction. We demonstrate the ability of AMAG to learn the underlying spatial interactions on synthetic datasets and show the importance of learning forecast in discovering such interactions.

---

Generating future neural signals with AMAG. We apply AMAG to four neural recording datasets and demonstrate its reliability in generating future neural activity while recovering the channel spatial proximity in the learned adjacency matrix and aligning neural trajectories in the latent space.
2
Related work
Modeling Dynamics of Neural Signals. Studying the dynamics of neural signals provides insights
into the computation performed in the brain [70). Neural dynamics has been modeled by Gaussian
Process 773
79
81 and lincar dynamical systems 445
46
23
41 2]. More recently, series of
nonlinear models have been proposed, including Multi-Layer Perceptron (MLP) [43 83 7. RNNs
44 47 Latent representations of population dynamics are often extracted while these models optimize for a reconstruction objective, upon which interpretable features emerge naturally (54 58 4 A recent model that could forecast future neural activities is Preferential Subspace IDentification algorithm (PSID) [56] [55. PSID proposes to forecast future neural activity along with behavior variables in two stages. In the first stage, the model learns to predict behavior and the behaviorally-rellvant portion of neural signals. In the second stage, PSID learns to forecast both behaviorally- and behaviorally-irreleyy neural signals 55).
[22 34
57
32
60 Neural ODEs (35), and Transformers [77 37,
Time Series Forecasting. Time series forecasting. i.c., the prediction of next values in time series, a well-studied topic that encompasses various applications, such as traffic flow forecasting, motion prediction, and electric load forecasting. Examples of proposed methods include methods that CNN based [31 RNN based (24 13,26), Multi-layer perceptron (80) and so on. Transformers with additional components for capturing long-term dependencies have also been proposed for time series forecasting (68) 16, 74). Also, models such as Graph Neural Networks (GNNs) have been proposed for time series forecasting focusing on datasets with inherent structural organization, e.g.,
traffic flow (30 77 38). For example, these methods have been applied to functional Magnetic Resonance Imaging (IMRI) forecasting (7I) leveraging Diffusion Convolutional Recurrent Neural Network (DCRNN) and Graph WaveNet (GWNet) (76. 40). In these applications, GNN has been embedded in GRU gating mechanism to learn the spatial-temporal relationship (DCRNN), or CNN and GNN have been proposed to be intertwined to perform contrast, GraphS4mer, which combines graph neural networks and structured state space models, performs temporal embedding followed by spatial message- message-passing in sequential manner [66]. These methods perform only the additive message-pant paasingg which may not be sufficient for a complex system such as the brain.
modeling (GWNet). In
GNNs for Neural Signals. In recent years, GNNs have been proposed as plausible models for analyzing brain activity across different modalities 662 33 12 [2]. In non-invasive neural recordings, GNNs have been applied EEG and fMRI data to capture the relationships between different brain regions or features. For example, LGGNet learns local-global-grapp various cognitive classification tasks [17), and BrainGNN has been proposed to identify neurological biomarkers with GNN, leveraging the topological and functional information of fMRI (39 Recently, the method has been extended to analyze ECoG recordings as a graph diffusion process (59]. A few applications of GNNs to invasive neural measurements have been introduced, although these are less ubiquitous than GNNs for other recording types. The Two-Stream GNN model, which has been designed to classify anesthetized states based on ECoG recording by building two fixed graphs, is notable example of such a model [TT]. The first graph in this model is the functional connectivity matrix computed from the phase lag index, and the second is the dual version of the first one. The Two-Stream GNN is not suitable for forecasting since it combines temporal signals into a unified vector the model input.
representations of EEG for
In light of the related work, the proposed AMAG is a GNN based approach which explicitly models interactions between neural signals to learn the dynamic features and interactions that can generate future neural signals. This is different from earlier models of learning neural dynamics with the primary focus on reconstruction without explicit modeling neural signal interactions. AMAG incorporates novel additive and multiplicative message-passing operations, which distinguish from methods focusing time series forecasting to associate GNNs with neural dynamics.
3

---

3 Methods
Our objective is to develop a model that estimates future neural dynamics given past activities. We define the relationship between future and past as x41:+++ fol (x:),, where X 7 represents neural signals recorded from C' electrode channels over time window of size T, with each channel having 1) dimensional features. X denotes neural signals from 0 f, and X 141++T corresponds to predicted future neural signals for time window & + 1 to +T. T.
To estimate f.).), we formulate an optimization problem minimizing the loss L between predicted future signals X (41:147 and ground truth future signals X X+11+++ as follows
min Ex fo
o.)), X X(-))))))
(I)
Specifically, fo(-) is constructed from three sub-modules shown in Fig. Temporal Encoding (TE) Module, Spatial Interaction (SI) Module, and Temporal Readout (TR) Module. The TE module captures temporal features for each channel individually. Then SI facilitates information exchange between channels, followed by TR to generate future neural activities. We describe the three modules in detail in the sections below.
3.1 SI with Add and Modulator
ST constructed as GNN. We represent Graph 9 == (V,E, (.EEAA) A), where V denotes a set of nodes [V] = C, & is a set of edges, and weighted adjacency matrix A E RCxC contains the weights of edges. The neighborhood of node e is written as (v).. where (v) = {ulu E V, (u, E E}. We denote neural signal from channel v at time / as x( FC, MLP and o (-) in the following sections represent a one-layer neural network, multi-layer perceptron, and sigmoid function, respectively.
Considering two types of interaction between neurons, additive (27) and multiplicative interactions [9 25). we design the SI module incorporate two sub-graphs for additive and multiplicative operations, namely Add and Modulator modules, respectively. In addition, we introduce a sample- dependent matrix, called Adaptor, to adjust the Add for each input sample, accounting for the complexity and dynamic nature of the brain as system.
Add Module. The Add Module performs the additive interaction between channels with the message- passing function Ma(-), depending on the adjacency matrix A₃ E RCC Assuming we have d dimensional channel features h for node at timestep & (h, E Rd ), Maa+) updates h, with the additive message from neighbor node u as Ma .h. A.) = 4 (u.0) h, (w) where u E NN) Such that the updated channel feature of Add Module the weighted feature of neighborhood,
(e)
ag (v)
Σ
1 (1.0) h
(EN.(())
The element (u, v) in the adjacency matrix A. indicates the influence of channel on e. Since A, is shared for all the input sequences, while channel interaction can change across inputs, we introduce sample-dependent Adaptor Module. We treat A, as fundamental adjacency matrix, and then we further learn sample-deeenden Adaptor as a matrix S E RCCC to uniquely modulate A, depending on each sample. Then, the shared message-passen passing function Eq. 3 becomes a sample-dependent message-passing function,
(v)
s(u,e) (u.v) h (1)
(3)
(EN.(M)
g(w.)) represents interaction strength between channel u and relying temporal embeddings for each channel. c.g., (1)) [h ............ ......,,.))
h(,") where is the context window for future neural signal generation. Specifically, we compute S(u.1) as (MLP(((H)()) ()), such that 5(a.") ranges from 0 to similar matrix construction for learning interpretable GNNs (88.
Since the fundamental adjacency matrix for the Add Module is unknown, A. needs to be learned during the optimization of the model parameters. We observe that direct learning of A,, from random initialization could make the training process unstable and thus instead, we initialize the A, with the precomputed correlation matrix from the neural recordings stabilize the learning process.

---

Modulator Module. The Modulator Module is the multiplicative message passing operator, with function (.).). We incorporate the adjacency matrix A., E RCCC to encode the neighborhood information that performs a multiplicative operation. M. indicates how the feature of the neighborhood modulates the target node feature, which can be expressed as (h, hh Am)
A (( h ④ A, where @ represents Hadamard product. The Modulator Module includes all the multiplicative information from the neighbors update feature for target channel u at time m
Σ
Mm(h, (w) h AA)
m
NEN (v)
Similarly to the adjacency matrix for the Add Module, the adjacency matrix in the Modulator Module also trainable and is initialized from the correlation matrix.
In summary, at specific timestep I, the output of is
(v) h/ + F2FCCa( (")) + F3FC(m]"))
(5)
with 31, f2. and 8 controlling the contributions of the self-connection, Add Module, and Modulator Module respectively.
3.2 Temporal Processing with TE and TR Modules
TE and TR perform temporal encoding and decoding. In particular, TE embeds input neural signals x E RRxD from channel e into the embedding space H E RRxx xd, For multi-step forecasting, we mask neural signals after the context window with constant, i.e., ET constant. Then SI updates representation of channel v as (v) E RTxd followed by TR generating future neuron (v)
(v)
activity X T. taking (()) as inputs. TE and TR can be Transformer or GRU. For further details on TE and TR, see Appendix A.2
Parameter Learning. In summary, AMAG estimates future neural signals with fo implementing TE, SI, and TR sequentially to the past neural signals. The optimal parameters 0 are obtained through backpropagation minimizing the Mean Squared Error (MSE) loss,
L= == (I) LT Xxrll)).
(6)
4 Experiments
4.1 Learning The Adjacency Matrix of a Synthetic Dataset
Synthetic Dataset. To explore the ability of AMAG to capture channel interactions, we generate synthetic datasets using linear non-Gaussian dynamic systems similar to previous works 66 6). Specifically, with adjacency matrix A describing the spatial interactions among variables (channels), datasets X are generated recursively at each time step as X, = X,-1 + G(X,-1, A) + where uniformly distributed noise term. We consider two scenarios where the connection ratio of the adjacency matrix is 0.2 and 0.5, respectively. The resulting datasets consist of multichannel time series, where the value of each channel at a given time step depends on its past and its neighboring channels, as defined by the adjacency matrix.
A Adjacency Matrix Recovery with AMAG. We evaluate the ability of AMAG recover the ground truth adjacency matrix and its dependency on tasks. We do so by additionally incorporating reconstruction and masked input reconstruction tasks as supplementary training objectives together with one-step and multi-step and forecasting tasks. The learned adjacency matrices, along with the ground truth matrix for 20 channels, are visualized in Fig (Additional results with different numbers of channels in Appendix B. To assess the accuracy of recovering the underlying ground truth adjacency matrix, we employ the Fl Score as quantitative measure. Our results show that the one-step forecasting task achieves the highest recovery of ground truth edges (F1 0.88 at 0.5 connection ratio), followed by the multi-step forecasting task (/1 0.82). This is compared to the reconstruction task which yields lower Fl score of 0.63. The lower performance of reconstruction tasks could be due to the model's 8 access to future signals, which potentially alleviates the need to learn the interactions with neighbors that future signals rely While multi-step forecasting generates
5

---

Reconstruction
Reconsiruction with Mask Mulli-stop Forecasting
One-step Forecasting
Ground Truth
Connection Rato: 2.2
10
1.0
F1 0.79
F1 = 0.83
F1 = 0.92
F1 0.93
Connection Rabi: 0.5.
16
10
200
0.63
F1 0.66 0.60
F1 0.82
F1 0.88
Figure 2: Colormap visualization AMAG learned adjacency matrices arranged by Fl score (right better), with yellow indicating Ground Truth recovery while gree-blue indicating lower accuracy, as reconstruction task (first column), reconstruction of masked input task (second column), multi-step forecasting (third column) and one-step forecasting task (fourth column), alongside ground truth adjacency matrices generating synthetic data (last column), demonstrating AMAG's ability to recover the underlying spatial interactions doing forecasting tasks.
multiple future signals, some adjacency matrices may generate more accurate signals than others given moment, but considering all future moments collectively, different adjacency matrices may achieve similar performance, resulting in learned adjacency matrices being less accurate than one-step forecasting.
4.2 Evaluation on Forecasting with Neural Signals
Dataset and Preprocessing In this study, we collect neural signals from monkey Affogato (A), and Beignet (B) with μÆCoG 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 μ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 MI from Monkey (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 Hz, 4-8 Hz, 8-12 Hz, 12-25 Hz. 25-50 Hz, 50-100 Hz, 100-200 Hz, and 200-400 Hz. The resulting signal is downsampled to 30ms. Full broadband data from A and B were collected 25 kHz. We then applied a preprocessing procedure extract LMP and powerband signals from A and B datasets as used in M and C datasets [20 64). The powerband signals are included 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 ms before to 450 ms after movement onset. For the and datasets, trials are temporally aligned to the appearance of the peripheral target and include 600ms after the peripheral target onset.
Benchmark Methods, We compare AMAG with existing approaches, including variants of RNNs, Transformers, and GNNs. Specificallyy, we compare with the following methods: LRNN RNN model without non-linear activation function and gating mechanism 663] RNNF. RNN model with non-linear activation and gating mechanism [32] 47): LFADS, RNN based model with pre- computed initial states [65]. NDT, 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, 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-relates neural space and behavior-unrelate neural space combined to ensure accurate

---

(A_1)
(A.2)
DCAAN:
One-step Forecasting
OCTRNN
Multi-step Forecasting
NDT R vs Scaling
0,07
AMM6
deccy
Gf
0.53
0.52
254
1928
2040 1656
Altention Dimension
Figure 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.
Moskey R²
Monkey g0
Monkey R
Moskey
MEE
MS51
STNDL 57
0.77
0.853:42-4
06635-11-4
0300=323
0.915=1c33
00111-3:44
0387422-3
80332:8:44
0.00002204
0.79ate33
0:3927c4
D0004427e5 5
08/71
0433 7--4
0607: Ri-5
cas 7e-1
0.55 72-4
0/054: Te-1
0566
0:33
0.00%
0506-17644
0991 1-364
RNO
2M1 1144
0296 2e-1
3e-4
00001 1--1
0 10-4
p001 33-4
8856179-1
0M611 3--1
Dicos 44-5
RNS 30
00108=3-5
0915=304
00871=3e-5
0.155764
5.00855655
0:032444
08004764
00375444
0.557=7655
4057
0.97745644
TERN
08862764
003224:44
0004783644
0.00099:4
0.90022-4
00062:463
0388elc3
0045=304
0.00772264
0.208444
0.6482244
0604552255 5
GWNe
10
0.91-55-4
8-866-444
66176-66-4
4.86-3e-4
0.912:1e-4
6-812:33-5
0500230
AMAG
099561633
3
80009-1e-3
0997-163
0.986-66-4
0.933-3-3
1e3
6.01F=11e4
0.999176-4
0.990:44-4
48015-4-8
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, GWNer: 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 '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 150ms 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 [.4). All the models are optimized by Adam Optimizer on Titan X GPU. The initial learning rates are 1e 4 and 5e 4 4 for non-GNN methods and GNN based methods. We include three evaluation metrics. R-squared (R²) 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 [B3
Forecasting Performance. We evaluate the performance AMAG to forecast future neural activities and present the results in Table [ 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 122 score by 8.2%. 5.0%, 5.7%, and 5.0% on Monkey M. C, B, and A datasets, respectively in Table Transformer only methods, STNDT (37] and NDT (78 achieve less optimal performance on four datasets. For instance, on Monkey M. the R² 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 shown in Table I AMAG shows consistent performance
7

---

(A)
(B)
(C)
0.60
Recording
Hidden Embedding
AMAG
0.55
0.0
0.50
NDT
0.45
PC2
Movement On 300ms Movement On Movementi
-10 PC1
M-L
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 micro- ECoG 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
Monkay 10
Mankey 29 1
Mankey ee"
Cwr
MEEE
Cer 1
MSE 1
Cor 1
M8B
Carz
MEEE
SINDY
012271 Se-1
D10781 4e-4
1030 5e-
00187. 2e-4
0525
07728 3e-3
010080 1-3-1
0.226 9e-3
01882 Ae-1
[10167
LRNN
(1281 Se-1
130013365
S001971 16-4
0011812-44
0.0 6-m7-4e-1
100000 1-6-1
0,00 5e-4
3.51a44-3
0.2575663
00729-16.4
0558856-3
60258-544
6.59-13-
0.70833-3
0.00693344
0.245
0.8054444
0015-28655
NDV
0.38723633
0.586333
30713636-4
0.668263
07532163
603452144
6575:463
0.736363
0.020122-4
0.750
0.873le33
0014022044
0
60309444
658-124
0.00122-3
0.24-1e-3
0
0.02663344
0.612333
0015823644
084-0
0.3685363
0530136-3
507302164
0.5627633
6699-163
08125363
0.051-8444
0.5992664
0.281
00100 3--
2011-177
3--1
0150
noun
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 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² on Monkey C and B when compared with non- 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² 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 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 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² 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² of 0.606 obtained with the KNN pruning.
In Fig. 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 used 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² 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.
8

---

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.
steli-step Forecasting RNNI
AMAG
LRNN
LFADS
STNDT
NDT
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
μР(М)
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 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 (st²) of NDT and AMAG with masked channels as input are shown in Fig. (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 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 target channel (T) onto other channels in the space organized by the 5ECoG arrangement. We observe that target channels exhibit stronger interactions with surrounding channels, which aligns with the properties of μ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. CC). 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 coherent and interpretable neural trajectories in 2D space.
Computation Complexity. We additionally compare in Table the number of parameters #P(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 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 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. (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.
4.3 Ablation Study
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 depen- dent "Adaptor" is kept in "Add" Module (@MAG), (4) Adjacency matrix is not learnable (AM-
G), (5) The "Add" is ablated from the graph operation (-MAG), (6) "Modulator" is removed (A-
AG), (7) The "Adaptor" in excluded from the "Add" (aMAG). The results shown in Table

---

One-step Forecasting
Multi-step Forecasting
Ave
Correlation Initaaiization
Correlation lilaaizztton
108
Random Initiaiizttion 12
Random Inttialization
Figure 5: Visualization of the adjacency matrix of Add module (A,) and Modulator module (Am) 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 Mod- ulator modules can cause the most significant forecasting accuracy drop (-AG) as well as the self-connection part in the graph (amAG). When keep- ing sample-dependent Adaptor term only (@MAG) in the Add module, the
Table 4: Comparison belween the Ablated versions of AMAG.
R²
CORR f
MSE +
-AG
0.424+1e-3
0.650+1e-3
0,0456+11cs lc-4
amAG
0-427+6e-3
0.652+44e-3
0.0453+55-4
-MAG
0.64742e-3
0.805+t6e44
0.0274+2e-4
āMAG
1.648:tle33
0.806:14e44
0.0273-i6e-5
A-AG
0.652.19e-4
0.80713e
0.026811e-4
performance is very close to ablating the Add module (-MAG). The sample- dependent Adaptor term appears be
aMAG
0.655112-33
0.81011e-3
0.0269112-4
AM-G Rand Init)
0.575++2e-2
0.767+
0.0329+1e-3
AM-G (Corr Init)
0.617:22e-4
0.786:1ee4
0.0296:22e55
less important compared to other com- ponents but still reduces the accuracy being ablated from AMAG (aMAG). 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 less stable training process, see Appendix [B.4] In Fig. [ 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 indi- cates that regardless of the initialization method, AMAG can converge to graph patterns with globally similar connections while multiple local connections are available achieve similar forecasting performance.
AMAG (Rand Init)
0.652+1e-3
0.807488-4
0.0270+1c-4
AMAG (Corr Init)
0.657+2e-3
0.811+2e-3
0.0266+22-4
Discussion
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 leamed neural features by AMAG appear to disentangle in the PCA space. While adjacency matrix recovery 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. requirement for GNN training increases the number of neural recording channels increases. For future development, improvements adapt the model for multi-region, multi-channel recordings could be considered.
worth noting that memory
10

---

Acknowledgements
Authors acknowledge the support in part by A3D3 National Science Foundation grant OAC-21 (ES,JL, the Simons Foundation (LS,PR,ALO)) the Washington Research Foundation Fund (ES), Departments of Electrical Computer Engineering (EL, 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.
Supplementary Information and Data
Supplementary information is available in the Appendix. Additional information and data will be made available through the project page https: htpps://githnn httpttttttuut/tt// com/shl
"'shhiieee (ezze///MG
References
[1] Larry F Abbott. "Lapicque's introduction of the integrate-and-fine model neuron (1907y". 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 1511107367 (2015).
[3] Navid Ayoobi and Elnaz. Banan Sadeghian. "A self-paces BCI system with low latency for motor imagery onset detection based on time series prediction paradigm". In: arXiv preprint arXiv:2204.0544 05450 (2022).
[4] Edoardo Balzani, Jean Paul Nocl, 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. 22661 (2022).
[5] Rahul Biswas and Eli Shlizerman. "Statistical perspective on functional and causal neural con- nectomics: A comparative study". " In: Frontiers in Systems Neuroscience 16 (2022), 817962. [6] Rahul Biswas and Eli Shlizerman. "Statistical perspective 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.0313 (2022).
[8] Anastasia Borovykh, Sander Bohte, and Cornelis W Oosterlee. "Conditional time series forecasting with convolutional neural networks". In: arXiv preprint arXiv: 1100.073399999 (2017).
[9] Nicolas Brunel. "Dynamics of sparsely connected networks of excitatory and inhibitory spiking neurons". In: Journal of computational neuroscience (2000). pp. 183-208.
[10] Luis Carrillo-Redi Reid, Shuting Han, Darik O 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: /EEE Transactions 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 O' Donnell. "TractGraphCNN: anatomically informed graph CNN for classification using diffusion MRI tractography". In: arXiv preprint arXiv:2301.019) (2023).
[13] Kyunghyun Cho, Bart Van Merriënboer, 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.

---

[15] Janis J Daly and Jonathan R Wolpaw. "Brain (Brain-computed interfaces in neurological rehabilita- tion". In: The Lancet Neurology 7.11 (2008), pp. 1032-1043.
[16] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. "Bert: Pre-training of deep bidirectional transformers for language understanding". In: arXiv preprint arXiv:1810.0488 (2018).
[17] Yi Ding, Neethu Robinson, Chengxuan Tong, Qiuhao Zeng, and Cuntai Guan. "LGGNet: Learning from gocal-global-graph representations for brain-computer interface". In: IEEE Transactions on Neural Networks and Learning Systems (2023).
[18) Matthew Dowling, Yuan Zhao, and II Memming Park. "Real-time variational method for learning neural trajectory and its dynamics". In: The Eleventh International Conference on Learning Representations.
[19] Katie A Ferguson and Jessica A Cardin. "Mechanisms underlying gain modulation in the cortex". In: Nature Reviews Neuroscience 21.2 (2020), pp. 80-92.
[20] Robert D) Flint, Christian Ethier, Hmily R Oby, Lee E. Miller, and Marc W Slutzky. "Local field potentials allow accurate decoding of muscle activity"". In: Journal of neurophysiology 108.1 (2012), pp. 18-24.
[21] Cecilia Gaalego-Carraceee Matthew G Perich, Raeed H Chowdhury, Lee E Miller, and Juan Álvaro Gallego. "Dataset from "Local field potentials reflect cortical population dynamics in region-speeific and frequency manner"'". In: (). DO1: 10. 5061 /drydddryddd.am xd2547dkt
URL:
10. /506//dryad xd2547dkt
[22] Cecilia Gallego-Carracedo, Matthew G Perich, Raeed H Chowdhury, Lee E Miller, and Juan Álvaro Gallego. "Local field potentials reflect cortical population dynamics in region-specific and frequency-dependent manner". In: Elife 11 (2022), e73155.
[23] Yuanjun Gao, Evan W Archer, Liam Paninski, and John P Cunningham. "Linear dynamical neural population models through nonlinear embeddings", In: Advances in neural information processing systems 29 (2016).
[24] Felix A Gers, Nicol N Schraudolph, and Jürgen Schmidhuber. "Learning precise timing with LSTM recurrent networks". In: Journal of machine learning research 3.Aug (2002), pp. 115- 143.
[25] Lukas N Groschner, Jonatan G Malis, Birte Zuidinga, and Alexander Borst. "A biophysical account of multiplication by a single neuron". In: Nature 603.7899 (2022), pp. 119-123.
[26] Sepp Hochreiter and Jürgen Schmidhuber. "Long short-term memory". In: Neural computation
9.8 (1997), pp. 1735-1780.
[27] Alan L Hodgkin and Andrew F Huxley. "A quantitative description of membrane current and its application to conduction and excitation in nerve". In: The Journal of physiology 117.4 (1952), p. 500.
[28) Kejia Hu, Mohsen Jamali, Ziev B Moses, Carlos A Ortega, Gabriel N Friedman, Wendong Xu, and Ziv M Williams. "Decoding unconstrained arm movements in primates using high-density electrocorticograply signals for brain-machine interface use". In: Scientific reports 8.1 (2018), pp. 1-13.
[29] Yicong Huang and Zhuliang Yu. "Representation Learning for Dynamic Functional Connec- tivities via Variational Dynamic Graph Latent Variable Models". In: Entropy 24.2 (2022),
p. 152.
[30] Weiwei Jiang, Jiayun Luo, Miao He, and Weixi Gu. "Graph Neural Network for Traffic Forecasting: The Research Progress". In: ISPRS International Journal of Informmtontnon 12.3 (2023), 100.
[31] Weiwei Jiang and Lin Zhang. "Geospatial data images: A deep-learning framework for traffic forecasting". In: Tsinghua Science and Technology 24.1 (2018), pp. 52-64.
[32] Bryan Jimenez. "Neural Network Models For Neurophysiology Data". PhD thesis. Purdue University Graduate School, 2022.
[33] Ming Jin, Enwei Zhu, Changde Du, Huiguang He, and Jinpeng Li. "PGCN: Pyramidal Graph Convolutional Network for BEG Emotion Recognition". In: arXiv preprint arXiv:2302.02520
(2023).
12

---

[34] Mohammad Reza Keshtkaran, Andrew R Sedler, Raeed H Chowdhury, Raghav Tandon, Diya Basrai, Sarah L Nguyen. Hansem Sohn, Mehrdad Jazayeri, Lee Miller, and Chethan Pandarinath. "A large-scale neural network training framework for generalized estimation of single-trial population dynamics".. In: bioRxiv (2021).
[35] Timothy Kim, Thomas Z Luo, Jonathan W Pillow, and Carlos Brody.
"Inferring latent dynamics underlying neural population activity via neural differential equations", In: Interna- tional Conference on Machine Learning. PMLR. 2021, pp. 5551-5561.
[36] Natalie Klein, Josue Orellana, Scott L Brincat, Earl K Miller, and Robert E Kass, "Torus graphs for multivariate phase coupling analysis". In: (2020).
[37] Trung Le and Eli Shlizerman. "STNDT: Modeling Neural Population Activity with a Spa- tiotemporal Transformer". In: arXiv preprint arXiv:2206.0477. (2022).
[38] Maosen Li, Siheng Chen, Yangheng Zhao, Ya Zhang, Yanfeng Wang, and Qi Tian. "Multiscale spatio-temporal graph neural networks for 3d skeleton-based motion prediction". In: IEEE
Transactions on Image Processing 30 (2021), pp. 7760-7775.
[39] Xiaoxiao Li, Yuan Zhou, Nicha Dvoreek, Muhan Zhang. Siyuan Gao, Juntang Zhuang, Dustin Scheinost, Lawrence H Staib, Pamela Ventola, and James S Duncan. "Braingnn: Interpretable brain graph neural network for fmri analysis". In: Medical Image Analysis 74 (2021), p. 102233, [40] Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan I Liu. "Diffusion convolutional recurrent neural network: Data-driven traffic forecasting". In: arXiv preprint arXiv: 1707 (2017).
[41] Scott W Linderman, Andrew Miller, Ryan P Adams, David M Blei, Liam Paninski, and Matthew Johnson. "Recurrent switching linear dynamical systems". In: arXiv preprint arXiv:1610.0880 066.046600 (2016).
[42] Hexuan Liu, Jimin Kim, and Eli Shlizerman. "Functional connectomics from neural dynamics: probabilistic graphical models for neuronal network of Cacnorhabditis elegans". In: Philosoph- ical Transactions of the Royal Society B: Biological Sciences 373.173 1758 (2018), p. 20170377.
[43] Ran Liu, Mehdi Azabou, Max Dabagia, Chi-Heng Lin, Mohammad Gheshlaghi Azar, Keith Hengen, Michal Valko, and Eva Dyer. "Drop, Swap, and Generate: A Self-Supervieed Ap- proach for Generating Neural Activity".. In: Advances in Neural Information Processing Systems 34 (2021).
[44] Ran Liu, Mehdi Azabou, Max Dabagia, Jingyun Xiao, and Eva Dyer. "Seeing the forest and the tree: Building representations of both individual and collective dynamics with transformers". In: Advances in Neural Information Processing Systems 35 (2022), pp. 2377-2391.
[45] Jakob H Macke, Lars Buesing, John P Cunningham, Byron M Yu, Krishna V Shenoy, and Maneesh Sahani. "Empirical models of spiking in neural populations". In: Advances in neural information processing systems 24 (2011).
[46] Jakob H Macke, Lars Buesing, and Maneesh Sahani. "Estimating state and parameters in state space models of spike trains", In: Advanced state space methods for neural and clinical data 137 (2015).
[47] Ganga Meghanath, Bryan Jimenez, and Joseph G Makin. "Inferring Population Dynamics in Macaque Cortex". In: arviv preprint arXiv:2304. 00000 (2023).
[48] Siqi Miao, Mia Liu, and Pan Li. "Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism". In: International Conference on Machine Learning (2022).
[49] Kai J Miller, Dora Hermes, and Nathan P Staff. "The current state of electrocorticography based compuer interfaces"' In: Neurosurgical focus 49.1 (2020), E2.
[50] Luis Fernando Nicolas-Alone and Jaime Gomez-Gil. "Brain computer interfaces, review".
In: sensors 12.2 (2012), pp. 121 1211-1279.
[51] Michael Nolan, Bijan Pesaran, Eli Shlizerman, and Amy L. Orsbomn. "Multi-block RNN Autoencoders Enable Broadband ECoG Signal Reconstruction". In: bioRxiv (2022), pp. 2022- 09.
[52] Chethan Pandarinath, Paul Nuyujukian, Christine H Blabe, Brittany I. Sorice, Jad Saab, Francis R Willett, Leigh R Hochberg, Krishna Shenoy, and Jaimie M Henderson. "High performance communication by people with paralysis using intracortical brain-computer interface". In: Elife 6 (2017) 17), e18554.
13

---

[53] Chethan Pandarinath, Daniel J O' Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D Stavisky, Jonathan C Kao, Eric M. Trautmann, Matthew T Kaufman, Stephen I Ryu, Leigh R Hochberg. et al. "Inferring single-trial neural population dynamics using sequential auto-encoders". In: Nature methods 15. 10 (2018), pp. 805-815.
[54] Felix Pei, Joel Ye, David Zoltowski, Anqi Wu, Raeed H Chowdhury, Hansem Sohn, Joseph O' O'Doherty, Krishna V Shenoy, Matthew T Kaufman, Mark Churchland, et al. "Neural Latents Benchmark 21: evaluating latent variable models of neural population activity". In: arXiv preprini www.i:v2v00....... 009,04463 (2021).
[55] Omid G Sani, Hamidreza Abbaspourazad, Yan T Wong, Bijan Pesaran, and Maryam M Shanechi. "Modeling behaviorally relevant neural dynamics enabled by preferential subspace identificationn. In: Nature Neuroscience 24.1 (2021), pp. 140-149.
[56] Omid G Sani, Bijan Pesaran, and Maryam M Shanechi. "Where is all the nonlinearity: flexible nonlinear modeling of behaviorally relevant neural dynamics using recurrent neural networks". In: bioRxiv (2021).
[57] Marine Schimel, Ta-Chu Kao, Kristopher T Jensen, and Guillaume Hennequin. "iLQR-VAN control-based learning of input-driven dynamics with applications to neural data". In: bioRxiv (2021).
[58]
Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis. "Leeamable latent embed- dings for joint behavioural and neural analysis". In: Nature (2023), pp. 1-9.
[59] Felix Schwock, Julien Bloch, Les Atlas, Shima Abadi, and Azadeh Yazzan-Shahmorad. "Estimating and Analyzing Neural Information flow using Signal Processing on Graphs" In: ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE. 2023, pp. 1-5.
[60] Qi She and Anqi Wu. "Neural dynamics discovery via gaussian process recurrent neural networks". In: Uncertainty in Artificial Intelligence. PMLR. 2020, pp. 454-464.
[61] Le Song, Mladen Kolar, and Eric Xing. "Time-varying dynamic bayesian networks". In: Advances in neural information processing systems 22 (2009).
[62] Tengfei Song, Wenming Zheng, Peng Song, and Zhen Cui. "EEG emotion recognition us- ing dynamical graph convolutional neural networks". In: IEEE Transactions on Affective Computing 11.3 (2018), pp. 532-541.
[63] Frieder Stolzenburg, Sandra Litz, Olivia Michael, and Oliver Obst. "The Power of Linear Recurrent Neural Networks". In: arXiv preprint arXiiv/802.000 033000 (2018).
[64] David Sussillo, Mark M Churchland, Matthew T Kaufman, and Krishna V Shenoy, "A neural network that finds a naturalistic solution for the production of muscle activity".. In: Nature neuroscience 18.7 (2015). pp. 1025-1033.
[65] David Sussillo, Rafal Jozefowicz, LF Abbott, and Chethan Pandarinath. "Lfads-latent factor analysis via dynamical systems". In: arXiv preprint arXiv: 1eo.1066806315 (2016).
[66] Siyi Tang. Jared A Dunnmon, Liangqiong Qu, Khaled K Saab, Christopher Lee-Messer, and Daniel L Rubin. "Spatiooemporal Modeling of Multivariate Signals With Graph Neural Networks and Structured State Space Models". In: arXiv preprint arXiv:2211.
(2022). [67] Michael Trumpis, Chia-Han Chiang, Amy I. Orsborn, Brinnae Bent, Jinghua Li, John A Rogers, Bijan Pesaran, Gregory Cogan, and Jonathan Viventi. "Sufficient sampling for kriging prediction of cortical potential in rat, monkey, and human μECoG", In: Journal of neural engineering 18.3 (2021), P. 036011.
[68] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. " "Attention is all you need". In: Advances in neural information processing systems 30 (2017).
[69] Giuseppe Vinci, Valérie Ventura, Matthew A Smith, and Robert E Kass. "Adjusted regular- ization in latent graphical models: Application to multiple-neuron spike count data". In: The annals of applied statistics 12.2 (2018), P. 1068.
[70] Saurabh Vyas, Matthew D Golub, David Sussillo, and Krishna V Shenoy. "Computation through neural population dynamics". In: Anmual review of neuroscience 43 (2020), pp. 249- 275.
[71] Simon Wein, Alina Schüller, Ana Maria Tomé, Wilhelm M Malloni, Mark W Greentee, and Elmar W Lang. "Forecasting brain activity based on models of spatiotemporal brain dynamics: A comparison of graph neural network architectures" In: Network Neuroscience 6.3 (2022),
pp. 665-701.
14

---

[72] Francis R Willett, Donald T Avansino, Leigh R Hochberg, Jaimie M Henderson, and Krishna brain-to-text communication via handwriting". In: Nature 593.7858 (2021), pp. 249-254.
V. Shenoy, "High-performant
[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.088 00817 (2020).
[75] Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip. "A comprehensive survey 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 preprini arXiv: 1906.0012 (2019).
[77] Jiexia Ye, Juanjuan Zhao, Kejiang Ye, and Chengzhong Xu. "How to build 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 In: Neurons, Behavior, Data analysis, and Theory (Aug. 2021). ISSN: 2690-2664, DOI: 10 51628 /01c 27358 URL: https nbdt neural
with Neural Data Transformers".
scholasticahg. com wrtpaatttecceeeeeeee 27358 representation - learning population iupulation-activiiyy
-atttaaal-----a--- -at.-dattta-ttoooo trrnsrrrrnerrrrr (visited on 08/29/2021).
[79] Byron M 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 Il Memming Park. "Variational latent gaussian process for recovering single- trial dynamies from population spike trains". In: Neural computation 29.5 (2017), pp. 1293- 1316.
[82] Yuan Zhao and II 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- Di-VAE".. In: Advances Neural Information Processing Systems 33 (2020). pp. 7234-7247.
[84] Feng Zhu, Andrew Sedler, Harrison A Grier, Nauman Ahad, Mark Davenport, Matthew Kaufmmn, 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.
15

---

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, define the input to the model as X RTxCxD, 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 or the time step l of input neural signals, we use the notations X RTTx or X; RCxD, respectively. To distinguish different data structures, matrices are denoted by bold capital letters, such X, X((),,, X,, while unbolded capital letters are used for vectors and scalars (e.g. x! E RD.. FC dentoes a single-layer fully connected neural network, LP denotes multi-layer fully connected neural network, and LN denotes the layer normalization technique. Aa and Am 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("), x(() E RTD, Then the correlation between channel u and channel e is
x(exT
Corr(u,
V V
We then initialize each element in the adjacency matrices A(".)) and ((r.)) 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 (TB) and readout (TR) modules. The TE module consists of single-layer GRU that takes input from
channel v and computes the corresponding temporal embedding features h! E Rd 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 and its neighboring channels (v), specified by the adjacency matrices A, and Am, the spatial interaction (SI) module updates the embedding features. This interaction among channels allows for modeling of spatial dependencies.
The TR module consists of single-layer GRU and fully connected (FC) layer. The module generates future neural signals. This pipeline of TE, SI and TR modules is shown below
h (e) GRUTE (x,", (1))) n
(0) SI(h (1 (uE(())) A., Am)
2,
r = GRUTn(Z (0) (0)
x(1] = FC(r(())
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 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 accumulation of errors. To address this issue and explore an alternative approach, we consider Transformers 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 up to the target time step T. based on the available neural recordings prior to
16

---

Before feeding neural signals into the Transformer model, several pre-proesssing steps are ap- plied. These steps include masking layer, a linear layer, and a positional embedding layer, which sequentially transform the neural signals into embedding space denoted as
E emb(X = FC(mask(XX) (X)) PE.
The masking layer (mask) masks out all the future inputs, such that X is constant during multi- step forecasting. The linear layer (FCC transforms the input neural signals from their original dimension D to an embedding space of dimension d using a linear transformation (FC' RD + R Additionally, positional encoding (PE) is applied to incorporate positional information into the embedding space, The positional embedding is defined as P PE( [pos, 2i) sin(pos/ 100002//4), PF(pos, 2i 1) sin(pos/ (0000%++1/mmood
described in [68]. Here pos corresponds the position in the temporal space, and 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.
e The embedded signal is then mapped to Q k("), and (( by weight matrix W Raxd as Query, Key and Value in Transformer model. Specifically, Q E q W W E W q. k!
n. El WR ko v(()= W.,.., The outputs from the temporal encoding can be written
(K(
ATT() Softmart
Vd )
H(') LL(MMP(EE 1)) ATT H(') [h ..... (T) H e Txd
Here we assume 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
= ST(h,) (v) (0) [=[") (v)
(w) I E N(v)), A.s Am):
zT
TR module is similar to TE, with the difference that instead of receiving embedded signal as input, TR receives (*) as input and generates x (*) with corresponding Q!") (*) W qr K(") z(() W v (z) W.. u. To clarify, we use the same notation for W k. W,, g. W.. but they are different
sets of parameters in TE and TR. While we describe the process for generating H H(v) and (()) 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 the input: 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 steps are considered as the context and not masked.
Multi-step forecasting: In this task, the model is trained to predict multiple future steps from given initial time step 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.
17

---

w/o Masking F1=0.78
Masking
Reconnttuuttoo
10
Ground Truth
08
Multi-step F1=0.87
One-step F1=0.94
00
40
02
Forecasting
28
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 fair comparison between these tasks, 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 & 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 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 leamed 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 0.2) connections the recovery of the ground truth connections for each channel. To evaluate the accuracy of the discovered adjacency matrices, we use the Fl score as a metric that measures the harmonic mean between precision and recall, providing comprehensive assessment of the matrix recovery. Fl 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 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 and Datasets. As described the main paper, collected two new datasets from monkey Affogato (A), and Beignet (B) with μ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 μÆCoG 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 reaching task towards one of eight different directions. The subject's hand position was tracked in
18

---

real-time using a marker- 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, 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 reward. For our study, we include the neural signals recorded during the 600ms following the appearance of the surrounding target as samples. Specifically, a preparatory context for forecasting the future 450ms of neural signals. Consequently, trials with 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 MI 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.
treat the first 150ms of neural signals as
Details of Monkey M and C Datasets. For dataset C and M [20 64], we 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 Hz, 4-8 Hz, 8-12 Hz, 12-25 Hz, 25-50 Hz, 50-100 Hz, 100-200 Hz, and 200-400 Hz. Full broadband data from A and B were collected at 25 kHz. We then applied a preprocessing procedure to extract LMP and powerband signals from A and B datasets 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 - std) and mean plus four times the standard deviation (mean + 4 x sid) are linearly scaled to interval [-1,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 hyperparameter search by exploring different values for the learning rate, hidden size, and weight decay. Specifically, we considered learning rates of 10 5 x 10 44 and 10 10-4, hidden sizes of 64, 128, 512, 1024, and 2048, and weight decay values of 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² correlation coefficient (Corr), and mean squared error (MSF) 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
The LFADS model 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
LFADS 665
19

---

of the encoder, controller, and generator in the model is 512. The model is trained with an initialized learning rate of 10 reducing by 0.95 every 50 epochs and weight decay rate of 10-4.
LRNN. We build a one-layer RNN model. For one-step forecasting, we the hidden dimension 1024 without weight decay on all four datasets. For multi-step forecasting, the hidden dimension is 2048 with a weight decay rate of 1 for M and C datasets and 10 for A and B datasets. All models are trained with the initial learning rate 10 decaying by 0.95 every 50 epochs.
RNNY 447] RNN__ is composed of a single-layer GRU. For one-step forecasting, the hidden dimension of GRU is 1024 with weight decay rate 10-1 on M and C datasets and weight decay for A and B datasets. For multi-step forecasting, we use hidden dimensions as 2048 on A and B datasets and 1024 on M and C datasets and set the weight decay rate to O. The initial learning rate is 10-4, which decays by 0.95 every 50 epochs in all cases.
RNN PSID [56]. We experiment with RNN PSID with the behavior-related hidden state of RNN PSID being 64 and the neuron-related hidden state being 512. The initial learning rate is 10-4, No weight decay is used in all cases.
The NDT model we use is a three-layer decoder-based Transformer. On A and B datasets, the attention dimension of the Transformer is 1024 for both multi-step and one-step forecasting. The attention dimension on M and datasets 512 and 256 for one and multi-step forecasting, respectively. Weight decay is set to 0 for one-step forecasting on four datasets. For multi-step forecasting, weight decay is 10 -4 on the M and C datasets and 10 5 on A and datasets. For all cases, set the initial learning rate 10tH which decays by 0.95 every 50 epochs in all cases.
NDT (78).
STNDT 37 STNDT model we use has 10 layers of spatial-temporal blocks, each with one spatial attention Transformer and one temporal attention Transformer. The spatial Transformer has hidden dimension as the temporal context length and the temporal attention has the dimension as the length of input neurons. We use the weight decay 10 for multi-step forecasting, 10 for one-step forecasting. We set the initial learning rate as 10 decaying by 0.95 every 50 epochs in all cases.
We experiment TERN as one-layer GRU followed by a Transformer layer, the hidden state both GRU and Transformer is 1024, and the model trained with an initial learning rate 10-1 decaying by 0.95 every 50 epochs. The weight decay rate is 10-6 for M and C datasets for both one-step and multi-step forecasting. Weight decay on and B datasets 0 and 10-4 and for one-step and multi-step forecasting.
TERN 47
GWNer (76). We use three layers GWNet where each layer has two blocks, including the three CNN kernels for dilation. residual, and skip convolution for and B datasets. Before generating the future neural signals, there is one end CNN layer and one read-out CNN layer. The channel size is 64 for residual and dilatation CNN, 128 for skip CNN, and 64 for end CNN on A and B datasets. For M and C datasets, the channel size for residual, dilation, skip, and end CNNs are 128, 128, 256, and 512, respectively. To be noticed, for M and datasets, each layer of WaveNet has three blocks. The initial learning rate is 5 x 10-4 decaying by 0.95 every 50 epochs. No weight decay is applied on and B datasets, and -7 weight decay rate on M and C datasets. Kernel size 2 for all CNN layers,
DCRNN [4).. DCRNN we use is built as two diffusion convolutional layers, and each layer performs 2 diffusion steps. The hidden dimension of each diffusion layer is 64 for A and datasets and 128 for M and C datasets, no weight decay is used. The learning rate is 5 x 10-1 decaying by 0.95 every 50 epochs.
Graphs4mer [66]. We use two variants of Graphs4mer, (i) with GRU (GS4-G) for temporal encoding, (ii) with structured state space model (GS4-S), to perform the temporal encoding for all four datasets. The hidden dimension size 256 on and B datasets with a weight decay rate of le 5 15- and 128 on M and C datasets with a weight decay rate of 0. The learning rate is 5 x 10 decaying by 0.95 every 50 epochs.
AMAG. In one-step forecasting scenario, we use single-layer GRU basic component for TR and TE with a 64-dimensional hidden size. For multi-step forecasting, we use Transformer layer as the basic component for TE and TR with hidden size 64. Sample-dependent matrix is computed using MLP taking the concatenated features from a pair of nodes as input, as used in
'TERN is also denoted as Neural RoBERTa, RoBERTa in short 32
20

---

Table 5: Multi-step forecasting performance of AMAG using GRU the components for TE and TR
modules (AMAG-G) versus AMAG using Transformer as the components for TE and TR modules (AMAG-T).
R² f
Corr t
MSE +
AMAG-G
0.653
0.807
0.0283
AMAG-T
0.658
0.811
0.0285
Corr
Tast w Corr liit
07
Test wi Corr Init
60
92
10
0.4
Epoch
Figure 7: Visualization of n22 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 x t, 64 x 2, 64 x 4, and 64 respectively where t is the temporal context length. There is no weight decay on M and datasets, 10-5 decay rate on A and datasets. The initial learning rate 5 10 decaying by 0.95 every 50 epochs.
B.4
Additional Results
GRU and Transformer Comparison as Base Components for TE and TR for Multi-step Fore- casting. 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.
Correlation Matrix as the Initialization of the Adjacency Matrix AMAG. well-initialized adjacency matrix in AMAG contributes to more stable training process. We investigate two initialization approaches: random and correlation matrix initialization. Both initialization methods result in similar performance in terms of the final 12 scores, i.e., 0.659 (randomly initialized) us. 0.658 (correlation matrix initialized). However, their learning curves exhibit distinet characteristics, as depicted in Fig. 7 When the adjacency matrix is initialized randomly (dashed curves), both training and testing curves display more fluctuations. In contrast, when the adjacency matrix is initialized with the correlation matrix (solid lines), the learning curves become smoother. This indicates that the correlation matrix initialization promotes a more stable and consistent learning process. Furthermore, examining the R²2 values in the initial epochs, we observe that the model initialized with the correlation matrix achieved higher performance compared to the random initialization. This finding suggests that careful initialization with the correlation matrix accelerate the training process, leading to improved performance during the early stages of training.
Adjacency Matrix Before and After Training. We visualize the learned adjacency matrix alongside the correlation matrix Fig. g We observe that the learned adjacency matrix more sparse compared to the original correlation matrix. This indicates that during the training process, AMAG prunes and
21

---

Correlation Matrix
Leamndd Adjacency Matrix
200
150
100
88
Figure 8: Illustration of correlation matrix as the initialization of the adjacency matrix in AMAG and the adjacency matrix after training. The learned adjacency matrix is more sparse compared with the correlation matrix (fewer values close to 1).
refines the correlation matrix, identifying the relevant and informative connections between neural channels for forecasting of neural signals.
Appendix C Analysis of Forecasting Error
We provide the following proofs understand what would be the upper bound error of the predicted neural activity. Considering the complexity of the AMAG model study, the theoretical guarantee of the temporal forecasting tasks could be potential research direction. the problem as follows: considering the one-step forecasting case and removing the nonlinearities in the model, the sample-dependent weight matrix, and the multiplicative weight matrix.
starting point, we simplify
Assume we have neural signal recorded from C' channels at time step t, X, E RCx1, output of TE layer, H, E (CxD)) (Ho = 0). The H, is updated with Linear RNN, such that H, = X,WY + H where WT and W are weights of linear RNN. Outputs of SI and TR are represented as Z, and RG respectively.
Z. AH,
R, = Z.WT + RA-1W
Ro 0
Z. R E (CDD)
We rewrite the He and Z, such that they only depend on X,
H. Σ XWW(
T=1
Z Σ AX, WXWi(WW)
R4 Σ Z, WW(
AX
The Predicted signal at the time /+ 8 1+ R.W. Since we assume the model stable, such that all the weight matrices should be bounded by M. i... ||WT]l2, [| ||WT] [[ll] < 1. Besides,
22

---

to make sure the RNN is stable, we need to constrain the eigenvalue of for WT. W WI, W2 to be smaller than 1, for example, AY, E wig(WTT) X; < With the assumption, we remove the term (W2)", (" > 2 which corresponds to the contribution of previous t-2 steps to t+1 step prediction. Thus we have
X = (AX,- + -X;- + AX,W1W])..
We rewrite the X, and X in form early step plus a variation term, such that
X, X X141 = X; +
- +
The prediction emror X1+1 X
e- -AXX-- + AX, -XX-IWWWWW +
- X
1 - -82-
Considering the first three terms that decide the e, we apply Schwar- incquality to get the upper bound,
[(AX +
& 3|/ + ||A||2|181|12
So we have:
|d]|2 3|| + || 1|X2--lll 1|61122 + ||x2||2
So if the signal is smooth enough that that ||, + ||52ll2 is small, then the prediction error is bounded by the adjacency matrix
References
[1] Larry F Abbott. "Lapicque's introduction of the integrate- and-fire model neuron (1907)". In: Brain research bulletin 50.5-6 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.0736 (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.060 05440 (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. 02866 (2022).
[5] Rahul Biswas and Eli Shlizerman. "Statistical perspective functional and causal neural con- nectomics: A comparative study". In: Frontiers in Systems Neuroscience 16 (2022), 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), c1010653.
[7] Martin Bjerke, Lukas Schott, Kristopher T Jensen, Claudia Battistin, David A Klindt, and Benjamin Dunn. "Understanding Neural Coding on Latent Manifolds by Sharing Features and Dividing Ensembles". In: arXiv preprint
(2022).
[8] Anastasia Borovykh, Sander Bohte, and Cornelis W Oosterlee. "Conditional time series forecasting with convolutional neural networks". In: arXiv preprint 1rXiv:1703.0mm 04031 (2017).
[9] Nicolas Brunel. "Dynamics of sparsely connected networks of excitatory and inhibitory spiking neurons". In: Journal of computational neuroscience (2000), pp. 183-208.
[10] Luis Carrillo-Reid, Shuting Han, Darik O' 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.
23

---

[1]] 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 O' Donnell. "TractGraphCNN: anatomically informed graph CNN for classification using diffusion MRI tractography". In: arXiv preprint vrXiv:2301.0191 (2023).
[13] Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. "On the properties of neural machine translation: Encoder-decoder approaches". In: arXiv preprint
arXiv:/409. 1259 (2014).
[14] Rainer Dahlhaus. "Graphical interaction models for multivariate time series". In: Metrika 51
(2000), pp. 157-172.
[15] Janis J Daly and Jonathan Wolpaw. "Brain-computer interfaces in neurological rehabilita- tion". In: The Lancel Neurology 7.11 (2008), pp. 1032-1043.
[16] Jacob Devlin, Ming-Wai Chang, Kenton Lee, and Kristina Toutanova. "Bert: Pre-training deep bidirectional transformers for language understanding". In: arXiv preprint arXiv: 04805 (2018).
[17] Yi Ding, Neethu Robinson, Chengxuan Tong, Qiuhao Zeng, and Cuntai Guan. "LGGNet: Learning from local-glohal-graph representations for brain-computer interface". In: IEEE Transactions on Neural Networks and Learning Systems (2023).
[18] Matthew Dowling, Yuan Zhao, and Memming Park. "Real-time variational method for learning neural trajectory and its dynamics". In: The Eleventh International Conference on Learning Representations.
[19] Katie A Ferguson and Jessica A Cardin. "Mechanisms underlying gain modulation in the cortex". In: Nature Reviews Neuroscience 21.2 (2020), pp. 80-92.
[20] Robert D Flint, Christian Ethier, Emily Oby, Lee Miller, and Marc W Slutzky. "Local field potentials allow accurate decoding of muscle activity". In: Journal of neurophysiology 108.1 (2012), pp. 18-24.
[21] Cecilia Gallego-Carracedo, Matthew Perich, Raeed Chowdhury, Lee Miller, and Juan Álvaro Gallego. "Dataset from "Local field potentials reflect cortical population dynamies in a region-specific and frequency- manner"" In: O). DOT: 10. 5001/ /061//rrydd xd2547dkt.
URL.: htt:httts:///aii
10. www.10661//raar xd2547dkt
[22] Cecilia Gallego-Caracedoo, Matthew G Perich, Raeed H Chowdhury, Lee E Miller, and Juan Álvaro Gallego. "Local field potentials reflect cortical population dynamics in a region-specific and frequency-dependem manner". In: Elife 11 (2022), e73155.
[23] Yuanjun Gao, Evan W Archer, Liam Paninski, and John P Cunningham. "Linear dynamical neural population models through nonlinear embeddings". In: Advances in neural information processing systems 29 (2016).
[24] Felix Gers, Nicol Schraudolph, and Jürgen Schmidhuber. "Learning precise timing with LSTM recurrent networks". In: Journal of machine learning research 3.Aug (2002), pp. 115- 143.
[25) Lukas N Grroshhnrr, Jonatan G Malis, Birte Zuidinga, and Alexander Borst. "A biophysical account of multiplication by a single neuron". In: Nature 603.7899 (2022), pp. 119-123 123.
[26]
Sepp Hochreiter and Jürgen Schmidhuber. "Long short-term memory". In: Neural computation 9.8 (1997). pp. 1735-1780.
[27] Alan L Hodgkin and Andrew F Huxley, "A quantitative description of membrane current and its application to conduction and excitation in nerve". In: The Journal of physiology 117.4 (1952), p. 500.
[28) Kejia Hu, Mohsen Jamali, Ziev B Moses, Carlos A Ortega, Gabriel N Friedman, Wendong Xu, and Ziv M Williams. "Decoding unconstrained arm movements in primates using high-density electrocorticography signals for brain-machine interface use". In: Scientific reports 8.1 (2018), pp. 1-13.
[29] Yicong Huang and Zhuliang Yu. "Representation Learning for Dynamic Functional Connec- tivities via Variational Dynamic Graph Latent Variable Models". In: Entropy 24.2 (2022),
p. 152.
24

---

[30] Weiwei Jiang, Jiayun Luo, Miao He, and Weixi Gu. "Graph Neural Network for Traffic Forecasting: The Research Progress". In: ISPRS International Journal of Geo- Geo-Information 12.3 (2023), 100.
[31] Weiwei Jiang and Lin Zhang. "Geospatial data images: A deep-learning framework for traffic forecasting". In: Tsinghua Science and Technology 24.1 (2018), pp. 52-64.
[32] Bryan Jimenez. "Neural Network Models For Neurophysiology Data". PhD thesis. Purdue University Graduate School, 2022.
[33] Ming Jin, Enwei Zhu, Changde Du, Huiguang He, and Jinpeng Li. "PGCN: Pyramidal Graph Convolutional Network for EEG Emotion Recognition". In: arXiv preprint arXiv:2302.02520
(2023).
[34] Mohammad Reza Keshtkaran, Andrew R Sedler, Raeed H Chowdhury. Raghav Tandon, Diya Basrai, Sarah L Nguyen, Hansem Sohn, Mehrdad Jazayeri, Lee E Miller, and Chethan Pandarinath. "A large-scale neural network training framework for generalized estimation of single-trial population dynamics".. In: bioRxiv (2021).
[35] Timothy D Kim, Thomas Z Luo. Jonathan W Pillow, and Carlos Brody. "Inferring latent dynamics underlying neural population activity via neural differential equations". In: Interna- tional Conference on Machine Learning. PMLR. 2021, pp. 5551-5561.
[36] Natalie Klein, Josue Orellana, Scott L Brineat, Earl K Miller, and Robert E Kass, "Torus graphs for multivariate phase coupling analysis". In: (2020).
[37] Trung Le and Eli Shlizerman. "STNDT: Modeling Neural Population Activity with Spa- tiotemporal Transformer". In: arXiv preprint (rXiv:2206.0772) (2022).
[38] Maosen Li, Siheng Chen, Yangheng Zhao. Ya Zhang. Yanfeng Wang, and Qi Tian. "Multiscale spatio-temporal graph neural networks for 3d skeleton-based motion prediction". In: IEEE
Transactions on Image Processing 30 (2021), pp. 7760-7775
[39] Xiaoxiao Li, Yuan Zhou, Nicha Dvorekk, Muhan Zhang, Siyuan Gao, Juntang Zhuang, Dustin Scheinost, Lawrence Staib, Pamela Ventola, and James S Duncan. "Braingnn: Interpretable brain graph neural network for fmri analysis". In: Medical Image Analysis 74 (2021), p. 102233.
[40] Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan Liu. "Diffusion convolutional recurrent neural network: Data-driven traffic forecasting". In: arXiv preprint arXiv: J 707707996 (2017).
[41] Scott W Linderman, Andrew C Miller, Ryan P Adams, David M Blei, Liam Paninski, and Matthew Johnson. "Recurrent switching linear dynamical systems". In: arXiv preprint arXiv:/610.0888 08466 (2016).
[42] Hexuan Liu, Jimin Kim, and Eli Shlizerman. "Functional connectomics from neural dynamics: probabilistic graphical models for neuronal network of Caenorhabditis elegans". In: Philosoph- ical Transactions of the Royal Society B: Biological Sciences 373. 1758 (2018), p. 20170377.
[43] Ran Liu, Mehdi Azabou, Max Dabagia, Chi-Heng Lin, Mohammad Gheshlaghi Azar, Keith Hengen, Michal Valko, and Eva Dyer. "Drop, Swap, and Generate: A Self-Supervised Ap- proach for Generating Neural Activity". In: Advances in Neural Information Processing Systems 34 (2021).
[44] Ran Liu, Mehdi Azabou, Max Dabagia, Jingyun Xiao, and Eva Dyer. "Seeing the forest and the tree: Building representations of both individual and collective dynamics with transformers". In: Advances in Neural Information Processing Systems 35 (2022), pp. 2377-2391.
[45] Jakob H Macke, Lars Buesing, John P Cunningham, Byron M Yu, Krishna V Shenoy, and Mancesh Sahani. "Empirical models of spiking in neural populations". " In: Advances in neural information processing systems 24 (2011)
[46] Jakob H Macke, Lars Buesing, and Maneesh Sahani. "Estimating state and parameters in state space models of spike trains". In: Advanced state space methods for neural and clinical data
137 (2015).
[47] Ganga Meghanath, Bryan Jimenez, and Joseph G Makin. "Inferring Population Dynamics in Macaque Cortex". In: arXiv preprint arXiv:2304 234.:23300000000 .60400 (2023).
[48] Siqi Miao, Mia Liu, and Pan Li. "Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism". In: International Conference on Machine Learning (2022).
[49] Kai J Miller, Dora Hermes, and Nathan P Staff. "The current state of electrocorticograpyy based brain brain- computer interfaces". In: Neurosurgical focus 49.1 (2020), E2.
[50] Luis Fernando Nicolas-A Alonso and Jaime Gomez-Gil. "Brain computer interfaces, a review".
In: sensors 12.2 (2012), pp. 1211-1279
25

---

[51] Michael Nolan, Bijan Pesaran, Eli Shlizerman, and Amy L Orsborn. "Multi-blor RNN Autoencoders Enable Broadband ECoG Signal Reconstruction". In: bioRxiv (2022), pp. 2022- 09.
[52] Chethan Pandarinath, Paul Nuyujukian, Christine H Blabe, Brittany L. Sorice, Jad Saab, Francis R Willett, Leigh R Hochberg, Krishna Shenoy, and Jaimie M Henderson. "High performance communication by people with paralysis using intracortical brain-computer interface". In: Elife 6 (2017), e18554.
[53] Chethan Pandarinath. Daniel J O' S'eaa, Jasmine Collins, Rafal Jozefowicz, Sergey D Stavisky, Jonathan C Kao, Eric M Trautmann, Matthew T Kaufman, Stephen I Ryu, Leigh R Hochberg, et al. "Inferring single-trial neural population dynamics using sequential auto-encoders". In: Nature methods 15.10 (2018), pp. 805-815.
[54] Felix Pei. Joel Ye, David Zoltowski, Anqi Wu, Raeed H Chowdhury, Hansem Sohn, Joseph E O' O'Doherty, Krishna V Shenoy, Matthew T Kaufman, Mark Churchland, et al. "Neural Latents Benchmark 21: evaluating latent variable models of neural population activity". In: arXiv preprint arXiv: 110.2409944444 (2021).
[55] Omid G Sani, Hamidreza Abbaspourazad, Yan T Wong, Bijan Pesaran, and Maryam M Shanechi. "Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification". In: Nature Neuroscience 24.1 1 (2021), pp. 140-149.
[56] Omid G Sani, Bijan Pesaran, and Maryam M Shanechi. "Where is all the nonlinearity: flexible nonlinear modeling of behaviorally relevant neural dynamics using recurrent neural networks".
In: bioRxiv (2021).
[57] Marine Schimel, Ta-Chu Kao, Kristopher T Jensen, and Guillaume Hennequin. "iLQR-VAE: control-based learning of input-driven dynamics with applications to neural data". In: bioRxiv (2021).
[58] Steffen Schneider, Jin Hwa Lee, and Mackenzie Weygandt Mathis. "Learnable latent embed- dings for joint behavioural and neural analysis". In: Nature (2023), pp. 1-9.
[59] Felix Schwock, Julien Bloch, Les Atlas, Shima Abadi, and Azadeh Yavdaan-Shahmorad. "Estimating and Analyzing Neural Information flow using Signal Processing on Graphs". In: ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP). IEEE. 2023, pp. 1-5.
[60] Qi She and Anqi Wu. "Neural dynamics discovery via gaussian process recurrent neural networks", In: Uncertainty in Artificial Intelligence. PMLR. 2020, pp. 454-464.
[61] Le Song, Mladen Kolar, and Eric Xing. "Time-varying dynamic bayesian networks". In: Advances in neural information processing systems 22 (2009).
[62] Tengfei Song, Wenming Zheng, Peng Song, and Zhen Cui. "EEG emotion recognition us- ing dynamical graph convolutional neural networks". In: IEEE Transactions on Affective Computing 11.3 (2018), pp. 532-541.
[63] Frieder Stolzenburg. Sandra Litz, Olivia Michael, and Oliver Obst. "The Power of Linear Recurrent Neural Networks". In: arXiv preprint frX.v:/802.003008 (2018).
[64] David Sussillo, Mark M Churchland, Matthew T Kaufman, and Krishna V Shenoy. "A network that finds naturalistic solution for the production of muscle activity".. In: Nature neuroscience 18.7 (2015). pp. 1025-1033.
"A neural
[65] David Sussillo, Rafal Jozefowicz, LF Abbott, and Chethan Pandarinath. "Lfads-latent factor analysis via dynamical systems". In: arXiv preprint arXiv:1608 063000006351111 (2016). [66] Siyi Tang, Jared A Dunnmon, Liangqiong Qu, Khaled K Saab, Christopher Lee-Messer, and Daniel L Rubin. "Spatiotemporal Modeling of Multivariate Signals With Graph Neural Networks and Structured State Space Models". In: arXiv preprint arXiv:2211.11176 (2022). [67] Michael Trumpis, Chia-Han Chiang, Amy L Orsborn, Brinnae Bent, Jinghua Li, John A Rogers, Bijan Pesaran, Gregory Cogan, and Jonathan Viventi. "Sufficient sampling for kriging prediction of cortical potential in rat, monkey, and human pECoG", In: Journal of neural engineering 18.3 (2021). p. 036011.
[68] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. " "Attention is all you need". In: Advances in neural information processing systems 30 (2017).
[69] Giuseppe Vinci, Valérie Ventura, Matthew A Smith, and Robert E Kass. "Adjusted regular- ization in latent graphical models: Application to multiple-neuron spike count data". In: The annals of applied statistics 12.2 (2018), p. 1068.
26

---

[70] Saurabh Vyas, Matthew D Golub. David Sussillo, and Krishna V Shenoy. "Computation through neural population dynamics".. In: Annual review of neuroscience 43 (2020), pp. 249- 275.
[71] Simon Wein, Alina Schüller, Ana Maria Tomé,
Wilhelm M Malloni, Mark W Greenlee, and Elmar W Lang. "Forecasting brain activity based on models of spatiotemporal brain dynamics: A comparison of graph neural network architectures" In: Network Neuroscience 6.3 (2022),
pp. 665-701.
[72] Francis R Willett, Donald T Avansino, Leigh R Hochberg, Jaimie M Henderson, and Krishna brain-to-text communication via handwriting". In: Nature 593.7858 (2021), pp. 249-20 254..
V Shenoy.
[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: /906.00/2 (2019).
[77] Jiexia Ye, Juanjuan Zhao, Kejiang Ye, and Chengzhong Xu. "How to build graph-base 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. DO1: 10 51628 001c 27358 URL: https nbåt scholasticaln cahq com wrtpccleeeleeeeeeeeee 27358 representationn learning for neural neural data- waw.eerrraaaaatttttaaaaaaooooooe twa.sttrrrssorssoooooaoeoo (visited on 08/29/2021).
population
[79] Byron M Yu, John P Cunningham, Gopal Santhanam, Stephen Ryu, Krishna V Shenoy, and Maneesh Sahani. "Gaussian- processss-ssessssese 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: 13504 (2022).
[81] Yuan Zhao and 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 I1 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-V."". 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 Giovannucei, 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.
27


Reviews:
[{"review_id": "8HIpwwelmh", "reviewer": "Reviewer_Yw6R", "summary": "The paper proposes a graph neural network-based model to forecast neural activities, which advances the DNN technology for neural understanding. It also proposes a method to leverage the causality structure of the signal into the moded design that generalizes the neural reconstruction task.", "questions": "It would be nice if the authors provided some theoretical analysis of why GNN based model can outperform the non-GNN-based model for neural forecasting.\n\nAlso, architecture complexity analysis is an interesting discussion.", "limitations": "limitations have been addressed", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper uses a graph neural network to do future prediction tasks, which is more complex than a reconstruction task. \n\nThe method has been verified in neural signals from monkeys and shown to outperform other state-of-the-art methods, including LFADS, and TERN.", "weaknesses": "Some analysis on the importance of each module, self-connection, add module, and modulator module need to be discussed. For example, why are they all necessary? \n\nFurthermore, the trade-off study between the model complexity and other sota methods is worth investigation. \n\nAlso, what is the computation bottleneck of the method, and how well does it scale?\n\nIs there any theoretical guarantee that the model can predict the neural activity if the trajectory satisfies some regularity?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1**. We thank the reviewer for suggesting additional analysis of the roles of Self-connection, Add and Modulator modules. These three modules are motivated by typical components in modeling neural activity, i.e. current activity, external additive input and gain modulation and according to our experiments are necessary to achieve AMAG accuracy as shown in the ablation study in Table 3 (main paper; Monkey C multi-step forecasting task). In particular, when Add or Modulator modules are ablated, R2=0.611 and R2=0.616, respectively, compared to R2=0.658 for full AMAG. \n\nWe also ablated both the Self-connection in AMAG (amAG), where, in this variant, R2 drops to 0.425, elucidating the importance of the Self-connection module. This analysis will be added to the final version of the paper.\n\n**W2**. We thank the reviewer for suggesting to estimate and compare AMAG complexity. We  estimated the number of parameters, training time, and maximum memory cost (Table R3 in the attached pdf). Graph-based methods use much fewer parameters than non-GNN methods, with AMAG consistently using minimal or comparable (DCRNN for one-step forecasting) number of parameters when compared to graph model baselines (DCRNN, GWNet, GraphS4mer). Less parameters comes at the expense of computation time and memory cost. We will include these estimates and the discussion in W3 below in the revised version of the paper.\n\n**W3**. As presented above in W2, AMAG generally requires fewer parameters and less training time (compared to graph-based methods). The major bottleneck of AMAG is the requirement for larger memory, thus, scaling AMAG could be limited by the availability of working memory.\n\n**W4**. We appreciate the reviewer raising the important point of analyzing how the proposed AMAG model predicts neural activity. Analysis of complex, non-linear models, such as GNNs / AMAG for time-series data is not well determined, though there are some initial interpretations and explorations made [1, 2]. \n\nWhile thorough analysis requires novel analytical tools and is outside of the scope of our work, under some simplifications analytical interpretation could be provided by simplifying AMAG to be linear and including only additive message passing and assuming that future signals rely on signals from two prior steps. Also, we constrain L2 norm of all weight matrices to be bounded by 1 for stability of linear RNN. Then upper bound error for $t+1$ can be expressed as\n\n$\\| \\hat{\\boldsymbol{X}}\\_{t+1} - \\hat{\\boldsymbol{X}}\\_{t+1} \\|\n\\leq 3\\|A\\|\\_2  \\| \\boldsymbol{X}\\_{t-1} \\|\\_2 + \\| A \\|\\_2\\| \\delta\\_1 \\|\\_2  + \\| \\boldsymbol{X}\\_{t-1} \\|\\_2 + \\| \\delta\\_t \\|\\_2  \\| + \\|\\delta\\_{t+1} \\|\\_2$\n\nHere, $A$ represents the adjacency matrix, $\\delta_1$ and $\\delta_2$ indicate variations in signals between consecutive time steps. Thus, if channel interactions remain small ($\\| A \\|_2$ value), and signals exhibit smooth transitions (small $\\| \\delta_t \\|_2 $) at each step, then the linear model will likely produce smaller prediction errors. \n\nWe will include such simplifying scenarios in the Appendix of the revised version.\n\n[1] Agarwal, Chirag, et al., 2022. \n\n[2] Li, Yiqiao, Jianlong Zhou et al., (2022).\n\n**W5** & **Q1**. We agree with the reviewer that such analysis could be instrumental to further interpret GNNs and AMAG. While rigorous analysis of such high-dimensional and nonlinear systems requires novel advanced analysis tools and is out of the scope of the current study, we identify empirically that GNNs advantage primarily lies in the ability to capture the topology of the activity. \nIndeed, ECoG recordings exhibit spatial-temporal interactions among recording channels which could be represented as a graph. GNNs, by design, explicitly model the electrode interactions, thus including the prior inherent topology of the dataset. This prior knowledge could constrain GNN to be closer to the optimum. Indeed, if we completely remove the graph structure in AMAG (amAG), the accuracy drops to R2= 0.423 (Monkey C; multi-step forecasting).\n\nIn non-GNN-based methods, channel signals are simply concatenated without topology information. The topology needs to be learned and thus typically these methods require more training samples. In datasets that we consider, the number of samples ranges from 874 to 1605. For such amounts of data, non-GNN methods, e.g., NDT, can overfit to the training data. Adding regularization terms, e.g., weight decay, could help, but will also limit the model's capacity. We demonstrate this in Fig. R2 (B) of the attached pdf by visualizing how R2 could be affected by the attention dimension and the weight decay. Adding weight decay to NDT training can improve (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 improvement. \nIn contrast, GNN-based methods can be viewed as a type of  ‘adaptive regularization’ which constrains the model learning trajectory. We will add this discussion to the revised version of the manuscript.\n\n**Q2**. We measured the complexity of the model in terms of the number of parameters, training efficiency, and memory cost. As discussed in reply to W2, graph-based models typically require fewer parameters but longer training time compared to non-graph-based methods. \nThe major computational bottleneck for AMAG is the requirement for larger working memory. One way to reduce memory cost and computation time would be to compress the original graph into a smaller hidden graph in the embedding space, with each node in the embedding representing multiple nodes in the original graph which is plausible considering that ECoG array records shared source input. Pruning could also reduce the graph size and reduce the memory cost, however, these methods will require manual setting of the pruning threshold and other hyperparameters.\n"}, {"review_id": "Fzf8Dx1cZm", "reviewer": "Reviewer_peT1", "summary": "The authors introduce a graph neural network in their study to predict neural activity. This network comprises a temporal encoding and decoding layer specific to each channel, with a spatial interaction layer positioned in between. The inclusion of the explicit spatial interaction layer aids in capturing the underlying spatial relationships, as demonstrated using synthetic data. The proposed model exhibits state-of-the-art performance in both one-step and multi-step forecasting tasks when evaluated on actual neural data.\nIn terms of the model-readout, it incorporates both additive and multiplicative operations, with both being necessary to attain the demonstrated accuracy. This observation is supported by an ablation study.", "questions": "- What is PSID, what  is DCRNN? After the introduction section, subsequent acronyms are not explicitly introduced or defined.\n- TE and TR can be transformer or GRU. Which one did you use for the results in the paper?\n- Why is the improvement in multi-step forecasting comparatively moderate despite the impressive enhancements in one-step forecasting? Intuitively, one would expect the performance gap between methods to increase with an increasing number of forecasting steps.\n- How many steps are used in multi-step forecasting?\n- RoBERTa in Table 2 is not mentioned in the caption", "limitations": "The authors have discussed limitations of their work.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- impressive enhancements in one-step forecasting performance\n- notable improvements are observed in multi-step forecasting, although to a lesser extent\n- explicit spatial interactions aid interpretability", "weaknesses": "- there is a presence of numerous non-standard acronyms in the paper that are not introduced upon their first mention\n- the utilization of graph neural networks for analyzing neural signals is not entirely novel, this study showcases somewhat incremental progress in this field ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1**. We appreciate the reviewer pointing out the use of non-standard acronyms in the paper. These are names of related methods introduced by their authors. We will make sure to include the full name of these methods at their first mention along with explanations of their origin, and provide citations in the revised version of our paper.\n\n**W2**. We agree with the reviewer that there have been previous works using graph neural networks for neural signals. But most of the previous works focus on neural signals recorded from EEG and fMRI where the temporal dynamics could be different from the field potential (Local Field Potential and ECoG) datasets that we study in this work. In addition, most of the previous works address the classification task, e.g., emotion detection or disease detection. Whereas we focus on the forecasting task. \n\nForecasting of neural signals is important for both scientific understanding and application. Indeed, as demonstrated in Section 4.1 of the main paper, through learning to forecast, AMAG learns to recover the underlying interaction between channels. Furthermore, learning to forecast can be used for anomaly detection monitoring of the neuronal recordings and for reducing the latency of the Brain-Computer Interfaces (BCI) especially when future behavior is not simply related to past behavior. We show in this work that AMAG can achieve SOTA forecasting performance with a novel design of GNN architecture with both additive and multiplicative messaging passing operations.\n\nIn addition to forecasting, we demonstrate that AMAG facilitates estimation of connectivity (on the synthetic dataset) and learned adjacency matrix on the experimental recordings dataset which shows meaningful ECoG arrangement (Figure 4 in main paper).\n\n**Q1**. We thank the reviewer for pointing out the need to define the acronyms early in the paper. \n\nPSID is the abbreviation of Preferential Subspace IDentification algorithm [54, 55], where the method identifies the behavioral relevant and irrelevant neuron subspaces by learning to predict the next step of neural signals and behavior signals in three stages. RNN PSID extends the algorithm with the RNN structure [55]. \nDCRNN refers to Diffusion Convolutional Recurrent Neural Network [39] . The model combines the Diffusion Graph with GRU to perform traffic forecasting tasks. These methods were introduced and discussed in the Related Work section, which follows the Introduction.  In the revised paper, we will make sure we define them earlier to avoid confusion.\n\nReferences from the main paper:\n[39] Yaguang Li, Rose Yu, Cyrus Shahabi, and Yan Liu. “Diffusion convolutional recurrent neural 458 network: Data-driven traffic forecasting”. In: arXiv preprint arXiv:1707.01926 (2017).\n[54] Omid G Sani, Hamidreza Abbaspourazad, Yan T Wong, Bijan Pesaran, and Maryam M 503 Shanechi. “Modeling behaviorally relevant neural dynamics enabled by preferential subspace 504 identification”. In: Nature Neuroscience 24.1 (2021), pp. 140–149.\n[55] Omid G Sani, Bijan Pesaran, and Maryam M Shanechi. “Where is all the nonlinearity: flexible 506 nonlinear modeling of behaviorally relevant neural dynamics using recurrent neural networks”. 507 In: bioRxiv (2021).\n\n**Q2**. For one-step prediction, the results reported in the main paper are when both TE and TR are GRUs. For multi-step prediction, TE and TR are Transformers. The results for multi-step with GRU is R2=0.658, as shown in Appendix Table A3. We additionally ran one-step forecasting using Transformer and obtained R2=0.969. These results led us to choose TE, TR differently in one-step and multi-step forecasting.\n\n**Q3**. We would like to clarify that the improvement of AMAG vs non-GNN methods on one-step forecasting range from 0.06 to 0.09 in terms of R2 and comparable to other GNN methods. The improvement for the multi-step forecasting is dataset dependent, ranging from 0.01 to 0.09 in terms of R2 (comparing the same set of methods as in one-step forecasting Table 1). \n\nIn most cases, the improvement is comparable between multi-step and one-step forecasting, except on Monkey A dataset, in which improvement of AMAG over other non-graph methods is more limited for multi-step forecasting. \n\n\nSuch a situation is possible since better one-step forecasting does not guarantee better multi-step forecasting. One can think of the multi-step forecasting as a greedy search where even if the optimal solution at each step is obtained, it does not guarantee a global optimal solution in the multi-step case, but in many situations it will reach close to it. \n\n**Q4**. For all four datasets, in multi-step forecasting we predict future 15 steps, the period when monkeys are moving the cursor from the center target to the surrounding target.\n\n**Q5**. We thank the reviewer for noticing this typo. The label should be TERN (RoBERTa is the name that was used in an earlier work [31], while in a later work a similar model was renamed to TERN [46].) The label will be corrected in the revised version of the paper.\n\n[31 Bryan Jimenez. “Neural Network Models For Neurophysiology Data”. PhD thesis. Purdue University Graduate School, 2022\n[46] Ganga Meghanath, Bryan Jimenez, and Joseph G Makin. “Inferring Population Dynamics in Macaque Cortex”. In: arXiv preprint arXiv:2304.06040(2023)."}, {"review_id": "bKnFetmfPf", "reviewer": "Reviewer_Byrz", "summary": "This paper proposes a graph neural network with additive and multiplicative message passing operations for neuron activity forecasting. The proposed model AMAG consists of a temporal encoder (TE), a spatial interaction (SI) module, and a temporal readout (TR) module. TE and TR modules are sequence models such as GRU and Transformer. The SI module consists of Add and Modulator modules, which are motivated by the additive and multiplicative interactions between neurons. Experiments on synthetic data and real-world neural recordings suggest the superiority of AMAG compared to non-GNN and GNN baselines.", "questions": "1. Experimental design: Model selection should be done on a separate validation set, and the test set should be held-out for reporting results only. Please follow this best practice for experiments. \n2. Baselines: Please have a more fair comparison for GNN baselines. e.g., use sparse graphs and use the original GraphS4mer architecture (S4 instead of GRU).\n3. Why is $A_a$ needed in the Add module? Would sampled-dependent $S$ be sufficient? It would be interesting to see an ablation experiment where $A_a$ is not included.\n4. Are the main results based on GRU or Transformer for TE/TR modules? Please clarify.\n5. Why aren’t GNN baselines included in Table 1 for one-step forecasting?\n6. Figure 4: Please explain how the trajectory in panel C is obtained.\n7. Please show comparisons to the strongest baseline(s) in Figure 3 and Figure 4A.", "limitations": "Yes, limitations are discussed in Discussion. The authors do not foresee any negative societal impact.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. There is some originality in the design of the additive and multiplicative message-passing operations in the SI module.\n2. The methods are technically sound.\n3. Overall this manuscript is easy to understand.", "weaknesses": "1. A major weakness of the manuscript is in the experimental design: the authors only split the data into train and test sets, and the hyperparameters and best model were chosen based on the test set (Appendix 2.3). The model could have been overfitted on the test set and made the reported results questionable, even though the authors showed results on a different train-test split in the Appendix. \n2. Comparisons to GNN baselines are not very fair. For instance, fully connected graphs are used for DCRNN and GraphS4mer. In the original papers for DCRNN (Li et al., 2017) and GraphS4mer (Tang et al., 2022), the graphs are pruned using a threshold and are therefore sparse. The authors also replace the S4 component in the baseline GraphS4mer with GRU, which is not the original GraphS4mer architecture.\n3. Given that there is a major issue in the experimental design, I would not consider the contribution significant.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1** & **Q1**. We thank the reviewer for pointing out this concern for experiments set up. As the reviewer noted, we provided additional experiments in the Appendix on a second train-test split for all four datasets. The test set of the first experiment is effectively the validation set since we train the model with the same set of hyperparameters of the first train-test split without additional hyperparameter search on the second train-test experiment. Since we did not search for hyperparameters when the model was trained on the second training and testing set, the results shown in the Appendix should not be overfitted to the second test set. \n\nThe results in the Appendix have the same trend as the results shown in the main paper, demonstrating that  AMAG successfully generalizes to a different set of training and testing data. \n\nAlso, we performed hyperparameter search similar to AMAG  for other methods and used the same set of hyperparameters when training the model on the second train-test split. This is to ensure that the results are comparable.\nWe will add an explanation of these details in the revised version of the paper.\n\n**W2** & **Q2**. We thank the reviewer for pointing out the need for further details regarding which graphs were used in baseline methods.  In both DCRNN and GraphS4mer we did use K-nearest neighbor (KNN) (as mentioned in GraphS4mer paper) pruned graphs. \nIn GraphS4mer we pruned the graph such that each neuron is connected to half of other neurons. We also tried threshold pruning (also mentioned in GraphS4mer) but setting the threshold is sensitive, especially with a large threshold (in terms of similarity), feature smoothing loss could become “NaN.” \n\nFor consistency, we pruned the graphs for DCRNN with KNN pruning. We additionally did experiments using threshold pruning with DCRNN on Monkey C with similarity thresholds of 0.5 and 0.8 where testing R2=0.573 and R2=0.445 (compared to DCRNN with KNN pruning being R2=0.618).\n\nWe will add these details and further discussion of graph pruning in the revised version of the paper.\n\nWe wish to point out that in regard to pruning, the advantage of AMAG is that AMAG automatically learns which connections should be kept instead of manually setting the sparseness of the model.\n\nThe reason that we did not use S4 is because S4 is suitable for long sequences; Based on the experiment, the performance in terms of R2 becomes 0.556 (Graphs4mer-S4) when using S4 instead of 0.579 (Graphs4mer-GRU). In the final version of the paper we will also include variations of baselines and their results and add further discussion of each baseline.  \n\n**Q3**. We thank the reviewer for the question. Aa is the primary weight matrix initialized with a correlation matrix that controls message passing between the channels without limitation in range. In contrast, each element in S is the output of the sigmoid function, ranging from 0 to 1 adapted for individual samples. S is viewed as the regularization term adjusting the weight in Aa matrix based on each input sample and is not sufficient by itself.\n\nWhen keeping S only we observed that the performance is very similar to removing both S and A - the included ablation results( approximately R2=0.605 on the Monkey C dataset). We will add this ablation variant to Table 3 in the revised version of the paper.\n\n**Q4**. We appreciate the reviewer's question regarding GNN baseline methods, such as DCRNN, GWNET, and GraphS4mer to one-step forecasting. These were originally designed for the multi-step forecasting scenario with the results in the original papers being for the multi-step forecasting.  Furthermore, since AMAG accuracy (R2, CORR) turned out to be close to 1, the one-step scenario would not be necessarily an ultimate test since other methods applied in this scenario could either be close to AMAG or lower than AMAG. We thus did not include the comparison of these methods with AMAG and focused on a more challenging (in terms of accuracy) scenario of multi-step forecasting. We will add these notes to the final version of the paper.\n\nAs a result of the reviewer's question, we tried to adapt GNN baselines to one-step forecasting. Since GraphS4mer learns a dynamic graph for each time window (several steps), performing one-step prediction requires learning the dynamic graph for each step which could be redundant. DCRNN and GWNET could be adapted (results in the attached pdf Table R1). We find that all three compared methods obtain approximately similar results. AMAG is better on Monkey A, and B, and GWNET is better on Monkey M and C. In addition, to achieve similar performance, AMAG used 0.12 million parameters while GWNet used 2.2 million parameters. This discussion and results will be incorporated into the final version of the paper. \n\n**Q5**. The results in the main paper are based on GRU for one step and transformer for multi-step.\n\n**Q6**. First, we get each step feature using either original 96-channel neuronal recordings (96 dimensions) or hidden states (64*96-dimension outputs of SI). Then, we perform PCA on the concatenated features of all timesteps and all samples. Figure 4C shows the first 2 dimensions in the PC space for original neuronal recording (left) and hidden states (right).\n\n**Q7**. Replot Figure 3 in attached pdf Fig.R1 (A).\nThe purpose of Figure 4A is to show that the SI module of AMAG is capable of selecting important channels. Specifically, testing if we mask channels with higher weight in the adjacency matrix (High importance) could cause a larger performance drop compared to masking channels with lower weight (Low importance). Thus comparisons are made per method. In addition, we examined the effect of masking the same set of channels of other methods, especially when the prior knowledge of channel interaction is not explicitly included in model design, i.e. NDT which was chosen since it has the same temporal structure as AMAG (using Transformer for temporal encoding).\n"}, {"review_id": "2Sw9u5DXa4", "reviewer": "Reviewer_BP1k", "summary": "This paper presents AMAG — a graph neural network for modeling and forecasting neural population dynamics. The graph neural network has mechanism to describe the additive and multiplicative interactions between neurons, and a sample-dependent matrix to adjust the additive component. Experiments are carried out on both synthetic data and real neural signals, in comparison with several baselines.", "questions": "In addition to address the comments raised in the weakness section above, please also considering clarifying the following questions. \n\n1. For the synthetic dataset presented in 4.1, what type of interaction mechanisms is used in the ground truth generation of data, i.e.., the function of G? Is it only additive, or multiplicative? The adjacency matrix A defines which interaction?\n\n2. Similarly, in the investigations based on the adjacency matrix in 4.2, which A's are being examined (A_a for additive interaction, or A_m for multiplicative interaction)? \n\n3. Overall, as mentioned in the weakness section, since the use of these two different interaction mechanisms is a key contribution in this paper, the authors should not only make it clear which A's are being considered (both in generation of data and in results), but it'd also be interesting if the authors add more analyses and results on how A_a and A_m may look like, and what are the underlying insights such differences or commonality could offer.\n", "limitations": "The authors discussed briefly the limitations associated with the study and future works necessary to address these limitations. \n", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "This paper is well written and easy to follow. The rationale behind AMAG design is clearly described, especially the basis of utilizing prior knowledge about additive and multiplicative processes. \n\nThe results overall show an improvement of forecasting performance by AMAG, and the experimental results especially those on actual monkey data are very interesting. \n\nThe selected results on visualization of interaction strength and the neural population trajectory embedding, i.e, those in Fig 4, are also very informative and interesting. \n", "weaknesses": "1. The quantitative results listed (e.g., Table 2) lack sufficient statistics, especially considering that the margin of improvements in some metrics are quite small. Some measure of std would be necessary.\n\n2. Other than the quantitative improvements (as shown in Table 1/Fig 3), the significance of such improvements can be better explained. E.g., in Fig 3, indeed it is evident that AMAG showed closer results to the true trajectory in the highlighted areas of the curve, however, in the remaining parts of the curve, it also falls short of capturing the true trajectories by a quite large margin (similar to the other methods), such as the big deflection in front of the highlighted area in the first curve of panel 2, or the deflection following the highlighted area in the second curve of the same panel -- Overall, it would seem that the improvements over other baselines are less significant or on par with the other major errors of the results. For signals that have a lot of temporal fluctuations, it would be good to understand what are the significance in minor to moderate level improvements in signal details when there are still substantial errors in the rest of the signals. \n\n 3. While the key innovation of the paper is motivated by leveraging prior knowledge about the additive and multiplicative interaction among neurons, the results provided limited insights into whether this modeling hypothesis could be verified from the results. There are ablation studies to show that both components are important (which is a strength of the paper), but there lack additional analyses into what are the different interactions being learned (e.g., by looking into the two adjacency matrices).\n\n4. Effect of the initialization on the two adjacency matrices should be elaborated with additional results (what happens when different correlation matrix is being used). What are the differences between the A_a and A_m after being initialized by the same matri?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1**. We thank the reviewer for recommending including statistics in Table 2.  We agree that adding multiple runs would elucidate the extent of variation and robustness of the results. Originally, we included runs as separate tables (one in the main paper and one in the Appendix). Following the reviewer’s recommendation, we included standard deviation of three runs for multi-step forecasting methods (As shown in Table R2 in the attached pdf). The improvement remains significant considering the standard deviation. We will update Table 2 and other tables with statistics in the revised paper.\n\n**W2**. The quantitative evaluation we examined, presented in Tables 1 and 3, is the collective evaluation comparing how predicted signals match the ground truth across multiple samples and multiple time points. Under the scenario, each sample and each time point equally contribute to the reported accuracy.\n \nThe examples shown in the right panel of Fig 3 are forecasted signals in the multi-step forecasting scenario (Monkey C), where AMAG results in R2=0.658. LRNN and NDT, in this case, result in 0.507 and 0.567, respectively. Considering that all three methods do not match the ground truth perfectly, there could be deflection in each of them. \n\nHowever, when comparing within the highlighted time window of the first curve of panel 2, AMAG is closer to the ground truth (GT) than other methods. Whereas, at an earlier time to the highlighted time window, there is indeed a deflection. \n\nThis deflection is for all three methods and close inspection shows that AMAG is slightly closer to GT than other methods. Similarly, for the second curve of panel two, AMAG is significantly better in the highlighted window, and in the window that follows it, all three methods appear to perform similarly. Overall, we observe that while there could be some windows in which all methods are similarly close or similarly deflect, there are windows in which AMAG has a significant advantage (these are the windows we highlighted). These improvements lead to an overall major improvement in the evaluation metric.\n\n**W3**. We thank the reviewer’s comment regarding further examination of the adjacency matrices for elucidation of primary components of AMAG. As suggested, we additionally visualized the learned adjacency matrices for both the Add module (Aa) and the Modulator module (Am) when AMAG is trained to perform one-step forecasting (Fig. R2 left panel) and multi-step forecasting (Fig. R2 right panel) on Monkey A. For each of the cases, we examined two scenarios to exclude initialization sensitivity: when both Aa and Am are initialized with the correlation matrix (Fig. R2 top) or both matrices are initialized randomly (Fig. R2 bottom). In the visualized matrix, the i-th row shows other channels’ contributions to the i-th channel. \n\nThese figures show that each matrix by itself (Aa or Am) appears consistent across initialization and forecasting length. However, when comparinge Aa with Am, while there is resemblance in their global pattern (which supports that a single module could generate reasonable prediction), each has distinct local patterns. Am appears to be more dense, while Aa is more sparse. Looking at specific target channels, some channels appear to contribute more in the additive module, while other channels could be more important in the modulator module.\n\n**W4**. We thank the reviewer for suggesting studying the effect of different initialization. As described in W3, Aa and Am are different after the learning regardless of initialization methods. For example, when both two matrices are initialized with the correlation matrix for one-step forecasting ( top row of Fig. R2 left panel), Aa and Am show different patterns (as discussed in W3). The conclusion generalizes to the random initialized weight matrix and multi-step forecasting cases.\n\nQuantitatively, we experimented using random initialization of Aa and Am in Appendix (Fig.A2), showing that AMAG achieves similar performance using different initialization (0.659 with random init vs. 0.658 using correlation init), but the learning process of correlation initialized AMAG can be more stable (Fig. A2 in Appendix). \nWith additional experiments on non-adaptable Aa and Am (AM-G) after initialization, correlation initialization can achieve better performance AM-G (Corr Init R2 = 0.623), compared to using random initialization AM-G (Rand Init 1 R2=0.585) (see below). \n\n| AMAG versions      | R2    | CORR  | MSE    |\n|--------------------|-------|-------|--------|\n| AM-G (Rand Init 1) | 0.585 | 0.774 | 0.0346 |\n| AM-G (Rand Init 2) | 0.596 | 0.773 | 0.0336 |\n| AM-G (Corr Init)   | 0.623 | 0.811 | 0.0285 |\n\nIn summary, the adaptable Aa and Am will converge to different matrices and can achieve similar performance with different initialization. However, when Aa and Am are non-adaptable (AM-G), the performance of the model can be affected by initialization methods.\n\n**Q1**. To generate the synthetic data, we only use the additive adjacency matrix (Aa) as described in Eq. 2 of the main paper. In this case, adjacency matrix A defines the additive interaction.\n\n**Q2**. For the analysis Investigating the adjacency matrix learned by AMAG,  we examined Aa, i.e., additive interaction.\n\n**Q3**. As discussed in Q2, the analysis of the adjacency matrix in 4.2 is based on Aa for additive interaction. We also visualized Aa and Am learned for one-step and multi-step forecasting with both random initialized matrices and correlation-initialized matrices in the attached pdf (Fig R2) (as discussed in W3 and W4). We find that regardless of the initialization method, in the one-step forecasting scenario, Aa matrix converges to similar patterns. Similarly, Am also converges to similar patterns, however Aa and Am have similar global patterns and different local patterns. This indicates that different channels contribute differently in terms of additive and multiplicative interaction."}, {"review_id": "c0g0nOuHBJ", "reviewer": "Reviewer_zdkz", "summary": "This manuscript looks at time series forecasting in neural multi-channel electrical signals.  They propose a new graph neural network based approach called AMAG that infers connectivity between channels to improve the ability to forecast over competing approaches, most of which are designed to infer latent dynamics or predict relevant task related behavior.  There are multiple steps in this approach, and the entire system is trained to maximize multi-step prediction.  Experimental results suggest the ability to capture relevant connectivity.", "questions": "What component is responsible for the major performance improvements compared to existing approaches?\n\nHow does this approach do at capturing other information such as neural decoding?  Is there evidence that improving at this forecasting task will really improve other downstream tasks?\n\nFor AM-G with the non-learnable adjacency, do you use the initialized correlation?  How sensitive is performance to that choice?\n", "limitations": "The only thing that I would add is that the authors should discuss more clearly that their improvements and results are focused on forecasting in the discussion, and that results are not shown for many of the other tasks where competing models have been effective.  Other than that, the comments seem fair.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "This approach appears to accurately capture connectivity in synthetic signals, which appears to largely hold up in real world data as well.\n\nPredictive performance shows large improvements in the forecasting (both one-step and multi-step) problem compared to competing approaches.\n\nAn ablation study implies that all components of the network are beneficial, motivating each component.", "weaknesses": "The major evaluation metric is forecasting in neural time series.  It is unclear how important this metric is by itself, whereas one often wants to evaluate how well we can extract relevant information (such as neural decoding) from the neural signals.  There is an implication that improving forecast of neural signals would improve these other problems, but it would be much stronger to explicitly show that.\n\nMy biggest hesitation is that it is unclear where the predictive improvements are coming from.  Most of the compared methods have performed well in many contexts, and the proposed structure is not so different than these existing methods.  The largest difference seems to be the GNN operating in the latent space of the system, whereas other methods capture these relationships in other ways, such as in the encoder function.  The ablation study does not clearly enough elucidate such differences, as all variants of the model seem to be extremely highly performing.  (I am assuming that Table 3 corresponds to Monkey C since the AMAG result matches up, but it doesn't explicitly state).  In fact, the ablated models still seem to be state-of-the-art, which is confusing.  I would encourage the authors to really explore the differences between their approach and existing methods and clearly show where the majority of the performance improvements happens.\n\nThere's a number of fuzzy details that need additional explanation throughout, such as the aforementioned ablation not clearly describing the data, and explicitly stating what the training of reconstruction and reconstruction with mask are in Figure 2. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**W1**. We thank the reviewer for pointing out the need for further information regarding the forecasting metric versus additional metrics. We focused on forecasting since the metric indicates the extent to which future signals, assumed unknown, could be predicted. For perfect forecasting accuracy, the unknown signals would be fully recovered. Note that in one-step forecasting, the accuracy in terms of R2 is close to such accuracy, while for multi-step, there is a significant gap still to overcome. There are also direct applications of estimating future activity, such as anomaly detection of neurons’ recording, reducing the latency of the BCI system, especially when the behavior signal itself is insufficient to predict future behavior, or there is no observable behavior. Examples include sleep spindles, epilepsy detection, and emotion detection, to name a few. We also investigated additional metrics, such as the reconstruction of connectivity. \n\nWe agree that behavior decoding could be another useful metric. We examined how the forecasted signal (multi-step) can be used for behavior decoding. We trained a behavior decoder using linear regression on ground truth neural and behavior signals applied to Monkey C. Then we applied the decoder to AMAG and other baselines' forecasted signals to examine how well their signals can be used for behavior decoding. Our results show improvement in decoding compared to the baselines (below). We will add these results and discussion to the revised version of the manuscript.  \n\n|    | GT    | GraphS4mer | DCRNN | GWNet | AMAG  |\n|----|-------|------------|-------|-------|-------|\n| R2 | 0.656 | 0.350      | 0.44  | 0.453 | 0.555 |\n\n**W2**.  As suggested by the reviewer, we included new variants in the ablation study by removing Self-connection (amAG) and removing both Add + Modulator (--AG) modules, i.e., the graph within AMAG, to elucidate the major component leading to SOTA forecasting (table below). The performance significantly drops for both amAG and –AG variants compared to full AMAG (R2=0.658). \n\nCombining the results with Table 3 in the main paper, we find that the two components: both the Self-connection module (amAG) and Add/Modulator (--AG) are essential for AMAG performance. Variants with only one module of the graph, i.e., only Add or only Modulator,  along with Self-connection lead to reasonable performance (-MAG=0.611, A-AG=0.616 ). The combination of all of them together provides the best accuracy. We will add these results and discussion to the revised version of the manuscript.  \n\n|      | R2    | CORR  | MSE    |\n|------|-------|-------|--------|\n| amAG | 0.425 | 0.652 | 0.0486 |\n| —AG  | 0.423 | 0.649 | 0.0488 |\n\n**W3**.  In the ablation study, all results were obtained for the multi-step forecasting task (Monkey C). \n\nThe reconstruction in Fig. 2 refers to: given an input neural signal, the model is trained to learn to reconstruct the exact input. \n\nReconstruction with masking indicates that the input into the model is masked (i.e., has missing data) such that the model learns to both reconstruct the input and fill in the missing information. Please see Appendix Section 2.1 for further details re. the experimental design.\n\n**Q1**. We thank the reviewer for W1 and this question to elucidate the components responsible for the major functionality of AMAG. As discussed in W2, both Self-connection and at least one module of the graph: Add or Modulator, are essential. The combination of all three provides the best performance as Add and Modulator modules supplement each other.\n\n**Q2**. Indeed, the forecasting task is a fundamental task and the recovery of unknown future signals can facilitate related applications and decoding, e.g., behavior decoding. As shown in reply to W1, AMAG forecasted signal facilitates a more accurate prediction of future behavior (R2=0.555) than other methods (R2~=0.45).\n\n**Q3**. We do use correlation initialization in AM-G, and it plays a role when the adjacency matrix is non-learnable. For example, when we substitute correlation initialization with random initialization, R2 drops from 0.623 to 0.585  (see below). Furthermore, fixed random initialization is more sensitive to the choice than fixed correlation initialization as we show for two random initialization variants in the table below (Rand Init 1 and Rand Init 2) compared to correlation initialization (Corr Init).\n\n| AMAG versions      | R2    | CORR  | MSE    |\n|--------------------|-------|-------|--------|\n| AM-G (Rand Init 1) | 0.585 | 0.774 | 0.0346 |\n| AM-G (Rand Init 2) | 0.596 | 0.773 | 0.0336 |\n| AM-G (Corr Init)   | 0.623 | 0.811 | 0.0285 |\n"}]

Author rebuttal:
We thank all reviewers for their insightful feedback. 
We tried to address all the questions for each reviewer in the rebuttal session below with references to weakness (**W**) and questions (**Q**).


Meta-review: This paper introduces AMAG, a graph neural network for predicting neuron activity, incorporating additive and multiplicative message passing operations. The model comprises a temporal encoder (TE), spatial interaction (SI) module, and temporal readout (TR) module. TE and TR leverage GRU and Transformer sequence models, respectively. The SI module integrates Add and Modulator modules inspired by neuron interactions. Comparative experiments using synthetic and real neural recordings demonstrate AMAG's superiority over non-GNN and GNN baselines.

The paper is well-structured and comprehensible. The rationale behind AMAG's design, including the incorporation of additive and multiplicative processes based on prior knowledge, is clearly elucidated. The results exhibit consistent enhancements in forecasting performance achieved by AMAG, particularly notable in the experiments involving actual monkey data. Overall, the paper contributes novel insights and advancements to the field.
Decision: Accept (poster)