SlowGuess/NeurIPS_2023_OCR huggingface dataset:
AMAG: Additive, Multiplicative and Adaptive Graph Neural Network For Forecasting Neural Activity
Jingyuan Li $^{1}$ , Leo Scholl $^{1}$ , Trung Le $^{1}$ , Pavithra Rajeswaran $^{2}$ , 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}@uw.edu
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.
The ability to encode, decode, and interpret neural activity could facilitate novel applications targeted toward neural activity, such as Brain-Computer Interfaces (BCIs) [15, 50, 52, 28, 49]. A promising methodology for neural activity encoding and decoding is the use 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 its variants, infer latent dynamics from single-trial population activities by encoding the neural time series into a 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,
Figure 1: AMAG Overview. (A) The forecasting task and AMAG architecture with temporal Encoding (TE), Spatial Interaction (SI), and Temporal Readout (TR) modules to 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 are closer to representations employed by the brain [82, 55, 56, 18]. In addition, learning to forecast offers practical advantages since the signals that are predicted can 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 to 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 of learning noisy information or overfitting the training data, the model learns features for forecasting future activity that are generalizable. A well-known prior is the spatial interaction among neurons, involving both addition [27, 9] and multiplication processes [19, 25] which could be modeled with Hodgkin-Huxley or Leaky Integrate-and-Fire [27, 1]. Previous works have utilized probabilistic graphical models to learn interactions [14, 42, 69, 36, 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, 44, 68]. Graph Neural Networks Networks (GNNs), designed to handle graph-structured data, are an alternative to the attention mechanism to model neuron interactions [75].
In this study, we propose a GNN with additive and multiplicative message-passing operations, called AMAG, for neuron signal forecasting. As shown in Fig. 1 (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-passing 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 a publicly available dataset recorded by Utah array (penetrating probes) and the second type is a novel dataset of micro-electrocorticography ( $\mu$ ECoG) array recordings from the cortical 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:
- Modeling spatial-temporal relationship. We present a graph neural network (GNN), called AMAG, that can forecast neuron activity by using sample-dependent 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 to 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.
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 [73, 79, 81] and linear dynamical systems [45, 46, 23, 41, 2]. More recently, a series of nonlinear models have been proposed, including Multi-Layer Perceptron (MLP) [43, 83, 7], RNNs [52, 34, 57, 32, 60], Neural ODEs [35], and Transformers [77, 37, 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-relevant portion of neural signals. In the second stage, PSID learns to forecast both behaviorally-relevant and behaviorally-irrelevant neural signals [56, 55].
Time Series Forecasting. Time series forecasting, i.e., the prediction of next values in a time series, is 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 are CNN based [31, 8], 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 (fMRI) forecasting [71] 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 spatial-temporal modeling (GWNet). In contrast, GraphS4mer, which combines graph neural networks and structured state space models, performs temporal embedding followed by spatial message-passing in a sequential manner [66]. These methods perform only the additive message-passing, which may not be sufficient for a complex system such as the brain.
GNNs for Neural Signals. In recent years, GNNs have been proposed as plausible models for analyzing brain activity across different modalities [62, 33, 12, 29]. In non-invasive neural recordings, GNNs have been applied to EEG and fMRI data to capture the relationships between different brain regions or features. For example, LGGNet learns local-global-graph representations of EEG for 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 a notable example of such a model [11]. 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 as the model input.
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 it from methods focusing on time series forecasting to associate GNNs with neural dynamics.
Our objective is to develop a model that estimates future neural dynamics given past activities. We define the relationship between future and past as  $\hat{X}_{t+1:t+\tau} = f_{\theta}(\mathbf{X}_{0:t})$ , where  $\mathbf{X} \in \mathbb{R}^{T \times C \times D}$  represents neural signals recorded from  $C$  electrode channels over a time window of size  $T$ , with each channel having  $D$  dimensional features.  $\mathbf{X}_{0:t}$  denotes neural signals from 0 to  $t$ , and  $\hat{\mathbf{X}}_{t+1:t+\tau}$  corresponds to predicted future neural signals for time window  $t+1$  to  $t+\tau$ .
To estimate  $f_{\theta}(\cdot)$ , we formulate an optimization problem minimizing the loss  $\mathcal{L}$  between predicted future signals  $\hat{\pmb{X}}_{t + 1:t + \tau}$  and ground truth future signals  $\pmb{X}_{t + 1:t + \tau}$  as follows
$$
\min  _ {f _ {\theta}} \mathbb {E} _ {\boldsymbol {X}} \left[ \mathcal {L} \left(f _ {\theta} \left(\boldsymbol {X} _ {0: t}\right), \boldsymbol {X} _ {t + 1: t + \tau}\right) \right]. \tag {1}
$$
Specifically,  $f_{\theta}(\cdot)$  is constructed from three sub-modules as shown in Fig. 1: 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
SI is constructed as a GNN. We represent a Graph  $\mathcal{G} = (\mathcal{V},\mathcal{E},\mathbf{A})$ , where  $\mathcal{V}$  denotes a set of nodes  $|\mathcal{V}| = C$ ,  $\mathcal{E}$  is a set of edges, and weighted adjacency matrix  $\mathbf{A} \in \mathbb{R}^{C \times C}$  contains the weights of edges. The neighborhood of a node  $v$  is written as  $N(v)$ , where  $N(v) = \{u|u \in V, (u,v) \in \mathcal{E}\}$ . We denote neural signal from channel  $v$  at time  $t$  as  $X_{t}^{(v)}$ .  $FC$ ,  $MLP$  and  $\sigma(\cdot)$  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 [19, 25], we design the SI module to 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 a system.
Add Module. The Add Module performs the additive interaction between channels with the message-passing function  $\mathcal{M}_a(\cdot)$ , depending on the adjacency matrix  $A_{a} \in \mathbb{R}^{C \times C}$ . Assuming we have  $d$  dimensional channel features  $\pmb{h}_{t}^{(v)}$  for node  $v$  at timestep  $t$  ( $\pmb{h}_{t}^{(v)} \in \mathbb{R}^{d}$ ),  $\mathcal{M}_a(\cdot)$  updates  $\pmb{h}_{t}^{(v)}$  with the additive message from neighbor node  $u$  as  $\mathcal{M}_a(\pmb{h}_{t}^{(u)}, \pmb{h}_{t}^{(v)}, \pmb{A}_a) = A_a^{(u,v)}\pmb{h}_{t}^{(u)}$ , where  $u \in N_a(v)$ . Such that the updated channel feature of Add Module is the weighted feature of neighborhood,
$$
\boldsymbol {a} _ {t} ^ {(v)} = \sum_ {u \in \boldsymbol {N} _ {a} (v)} A _ {a} ^ {(u, v)} \boldsymbol {h} _ {t} ^ {(u)}. \tag {2}
$$
The element  $(u,v)$  in the adjacency matrix  $A_{a}$  indicates the influence of channel  $u$  on  $v$ . Since  $A_{a}$  is shared for all the input sequences, while channel interaction can change across inputs, we introduce a sample-dependent Adaptor Module. We treat  $A_{a}$  as a fundamental adjacency matrix, and then we further learn a sample-dependent Adaptor as a matrix  $S \in \mathbb{R}^{C \times C}$  to uniquely modulate  $A_{a}$  depending on each sample. Then, the shared message-passing function Eq. 3 becomes a sample-dependent message-passing function,
$$
\boldsymbol {a} _ {t} ^ {(v)} = \sum_ {u \in N _ {a} (v)} S ^ {(u, v)} A _ {a} ^ {(u, v)} \boldsymbol {h} _ {t} ^ {(u)}. \tag {3}
$$
$S^{(u,v)}$  represents interaction strength between channel  $u$  and  $v$  relying on temporal embeddings for each channel, e.g.,  $\pmb{H}^{(v)} = [\pmb{h}_1^{(v)},\dots,\pmb{h}_t^{(v)}]$  where  $t$  is the context window for future neural signal generation. Specifically, we compute  $S^{(u,v)}$  as  $\sigma (MLP([H^{(u)},H^{(v)}])$ ), such that  $S^{(u,v)}$  ranges from 0 to 1, similar to matrix construction for learning interpretable GNNs [48].
Since the fundamental adjacency matrix for the Add Module is unknown,  $\mathbf{A}_a$  needs to be learned during the optimization of the model parameters. We observe that direct learning of  $\mathbf{A}_a$  from random initialization could make the training process unstable and thus instead, we initialize the  $\mathbf{A}_a$  with the precomputed correlation matrix from the neural recordings to stabilize the learning process.
Modulator Module. The Modulator Module is the multiplicative message-passing operator, with function  $\mathcal{M}_m(\cdot)$ . We incorporate the adjacency matrix  $A_{m} \in \mathbf{R}^{C \times C}$  to encode the neighborhood information that performs a multiplicative operation.  $\mathcal{M}_m(\cdot)$  indicates how the feature of the neighborhood modulates the target node feature, which can be expressed as  $\mathcal{M}_m(\pmb{h}_t^{(u)}, \pmb{h}_t^{(v)}, \pmb{A}_m) = A_m^{(u,v)}\pmb{h}_t^{(u)} \odot \pmb{h}_t^{(v)}$ , where  $\odot$  represents Hadamard product. The Modulator Module includes all the multiplicative information from the neighbors to update feature for target channel  $u$  at time  $t$  as  $\mathbf{m}_t^{(u)}$
$$
\boldsymbol {m} _ {t} ^ {(v)} = \sum_ {u \in \boldsymbol {N} _ {m} (v)} \mathcal {M} _ {m} \left(\boldsymbol {h} _ {t} ^ {(u)}, \boldsymbol {h} _ {t} ^ {(v)}, \boldsymbol {A} _ {m}\right) \tag {4}
$$
Similarly to the adjacency matrix for the Add Module, the adjacency matrix in the Modulator Module is also trainable and is initialized from the correlation matrix.
In summary, at a specific timestep  $t$ , the output of SI is
$$
\boldsymbol {z} _ {t} ^ {(v)} = \beta_ {1} \boldsymbol {h} _ {t} ^ {(v)} + \beta_ {2} F C \left(\boldsymbol {a} _ {t} ^ {(v)}\right) + \beta_ {3} F C \left(\boldsymbol {m} _ {t} ^ {(v)}\right), \tag {5}
$$
with  $\beta_{1},\beta_{2}$ , and  $\beta_{3}$  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  $\mathbf{X}^{(v)} \in \mathbb{R}^{T \times D}$  from channel  $v$  into the embedding space  $\mathbf{H}^{(v)} \in \mathbb{R}^{T \times d}$ . For multi-step forecasting, we mask neural signals after the context window  $t$  with constant, i.e.,  $\mathbf{X}_{t:T}^{(v)} = \text{constant}$ . Then SI updates representation of channel  $v$  as  $\mathbf{Z}^{(v)} \in \mathbb{R}^{T \times d}$ , followed by TR generating future neuron activity  $\hat{\mathbf{X}}_{t:T}^{(v)}$  taking  $\mathbf{Z}^{(v)}$  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  $f_{\theta}$  implementing TE, SI, and TR sequentially to the past neural signals. The optimal parameters  $\theta$  are obtained through backpropagation minimizing the Mean Squared Error (MSE) loss,
$$
\mathcal {L} = \mathbb {E} _ {\boldsymbol {X}} \left(\left\| \hat {\boldsymbol {X}} _ {t: T} - \boldsymbol {X} _ {t: T} \right\| _ {2}\right). \tag {6}
$$
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 [61, 5, 6]. Specifically, with adjacency matrix  $\mathbf{A}$  describing the spatial interactions among variables (channels), datasets  $\mathbf{X}$  are generated recursively at each time step as  $X_{t} = X_{t - 1} + \mathcal{G}(X_{t - 1},\mathbf{A}) + \mu$ , where  $\mu$  is 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.
Adjacency Matrix Recovery with AMAG. We evaluate the ability of AMAG to 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.2 (Additional results with different numbers of channels in Appendix B.1). To assess the accuracy of recovering the underlying ground truth adjacency matrix, we employ the F1 Score as a quantitative measure. Our results show that the one-step forecasting task achieves the highest recovery of ground truth edges ( $F1 = \mathbf{0.88}$  at 0.5 connection ratio), followed by the multi-step forecasting task ( $F1 = \mathbf{0.82}$ ). This is compared to the reconstruction task which yields a lower F1 score of 0.63. The lower performance of reconstruction tasks could be due to the model's access to future signals, which potentially alleviates the need to learn the interactions with neighbors that future signals rely on. While multi-step forecasting generates
Figure 2: Colormap visualization of AMAG learned adjacency matrices arranged by F1 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 at a 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