A Deep Probabilistic Spatiotemporal Framework for Dynamic Graph Representation Learning with Application to Brain Disorder Identification

1 Introduction

The human brain is a complex system that consists of numerous interconnected neuronal regions. As revealed by functional magnetic resonance imaging (fMRI), spontaneous spatiotemporal fluctuations in brain activity exist even during a resting state. Functional connectivity (FC) in brain networks is typically characterized via statistical dependence (such as correlations) between blood oxygen level dependent (BOLD) fMRI signals over spatially-distinct brain regions. Alterations in resting-state FC networks have been associated with neuropsychiatric or neurodevelopmental disorders [Filippi et al., 2019], for example in autism spectrum disorder (ASD) [Holiga et al., 2019; Valenti et al., 2020; Wang et al., 2021].

Machine learning (ML) techniques have been applied to the identification and prediction of neuropsychiatric disorders in the past decade [Tanveer et al., 2020; Hyde et al., 2019]. Traditional ML methods such as support vector machines (SVM), artificial neural networks, and deep learning methods are very widely used ML techniques [Moridian et al., 2022]. Deep learning has especially gained popularity in the studies of neuroscience as it offers an opportunity to discover and understand the underlying differences in the pattern of brain connectivity between individuals with psychiatric disorders and healthy controls using deep neural networks [RaviPrakash et al., 2019; Khodatars et al., 2021; Zhang et al., 2020]. Conventional approaches have modelled brain networks as 2D grid-like connectivity matrices, neglecting the non-Euclidean nature of brain networks. To address this, graph-based deep learning offers flexibility in modeling individual-level pairwise brain region interactions [Yan et al., 2019a], capturing group-level associations using population graphs [Jiang et al., 2020a], or combining both strategies [Li et al., 2022; Zhou and Zhang, 2021]. Yet, studies using graph neural networks (GNNs) in these contexts often focus on group-level network topology with phenotype data or supervised subject-level embedding using pre-computed population graphs for classifying static brain networks.

Dynamic spatiotemporal patterns in resting-state brain functional networks have been reported in several studies [Lee and Kim, 2019; Aedo-Jury et al., 2020]. They emphasize the consideration of not only the stationary brain network topology but also the dynamic temporal structure. Studies show that GNNs have potential in studying brain networks. However, the application of GNNs to model graphstructured brain network data has predominantly been limited to static representations, and there is a notable scarcity of GNN models designed to investigate the dynamic aspects of such networks. For instance, [Azevedo et al., 2020] employed 1-dimensional CNNs for temporal feature extraction and graph convolutional networks (GCNs) for spatial information sharing, and the simple framework showed great potential on both binary and continuous label prediction. [Kim et al., 2021] proposed self-attention and squeeze-excitation for readout, applied successively on GRU-encoded graphs to obtain spatiotemporal embeddings for downstream gender classification. Rather than focusing solely on the gen-

(a)

der classification problem, further works [Noman et al., 2022; Liu et al., 2023] incorporated GCNs to initially learn spatialaware graph embeddings of the dynamic FC. This was followed by a long short-term memory (LSTM) model to leverage the temporal dependencies of these learned graph embeddings for subject-level network classification of psychiatric disorders. These methods have predominantly relied on deterministic models, and their compartmentalized, taskspecific modules hinder the retention of full spatiotemporal information across layers.

In this paper, a deep spatiotemporal variational Bayes (DSVB) framework is proposed to learn the stochastic dynamic graph embeddings of FC networks, composed of sliding window segments of the resting-state fMRI time-series data. Figure 1 shows an overview of the proposed DSVB framework. Derived from general sequential variational Bayes (VB), our proposed framework provides a stochastic representation modelling solution of using dynamic FC networks for neuropsychiatric disorder diagnosis. Unlike existing GNN-based approaches that learn independent and identically distributed (i.i.d.) graph embeddings, we incorporate a spatial-aware recurrent neural network with attention-based message-passing scheme to model both the complex topological structure and temporal dependencies within a dynamic FC network sequence. Additionally, we introduce an adversarial training strategy to overcome the overfitting issue aggravated by limited training data. Our contributions are highlighted as follows:

1. An end-to-end probabilistic framework, integrating variational Bayes and recurrent GNNs, to learn the sequential stochastic graph latent embedding spaces in a purely unsupervised manner. This accommodates a wider range
2. of spatiotemporal variability across unseen brain FC networks.
3. A spatial-aware Gated Recurrent Unit (GRU) with attention-based message-passing scheme is incorporated to generate hierarchical latent embeddings to extract underlying non-Euclidean topological structure and rich temporal patterns of dynamic FC networks. The alterations in these latent embeddings are then exploited for downstream ASD identification.
4. An adversarial model perturbation strategy is introduced to apply regularization on the GNN-based graph embedding models during training. This adversarial regularization facilitates learning of a smooth latent embedding space that extrapolates well to unseen dynamic FC networks beyond the training dataset.
5. A nested-stratified 5-fold cross-validation shows that our proposed DSVB framework outperforms state-of-the-art methods in ASD identification from fMRI.

2 Methods

2.1 Dynamic Brain Network Construction

Dynamic brain graphs were constructed from fMRI data using the sliding-window technique. The regions of interest (ROI) fMRI time series was segmented into multiple overlapping blocks using sliding windows of size L = 20 with a shift of 10 time steps. To represent the dynamic FC networks, we computed the correlation matrix between ROI time series for each time window. The Ledoit-Wolf (LDW) regularized shrinkage covariance estimator is used to ensure wellconditioned FC estimates.

The dynamic brain network at each time-step t is represented by a graph Gt ≡ (V,E), where vi ∈ V represents a particular brain ROI and eij ∈ E is the connectivity edge between a pair of nodes vi and vj. The topological structure of dynamic brain networks Gt of N (the number of) nodes can be represented by a sequence of timeresolved adjacency matrices At = [at,ij] ∈ {0, 1}N×N (for each t), which were constructed by proportional thresholding of the sliding-windowed correlation-based FC matrices to keep the strongest 40% of connections based on absolute correlation value, and setting other connections to zero. Each dynamic graph Gt has associated node features Xt = [xt,1, . . . , xt,N]⊤ ∈ RN×Dx where xt,i ∈ RDx is the feature (column) vector of node vi. We used the connection weights of each node for xt,i.

2.2 Spatiotemporal Variational Bayes

To implicitly learn the spatiotemporal structure of dynamic FC, we consider the unsupervised learning task of reconstructing the sequential network-derived adjacency matrices A = {At}Tt=0 from node feature data X = {Xt}Tt=0 , where T + 1 is the sequence length. Under the sequential Bayesian framework, this is achieved via maximizing the variational evidence lower bound (ELBO) of conditional adjacency loglikelihood log pθ(A|X ).

Taking into account the difficulty of directly modelling the complex graph log-likelihood, we introduce a sequence of latent embeddings Z = {Zt}Tt=0 with Zt = [zt,1, . . . , zt,N]⊤ ∈ RN×Dz where zt,i ∈ RDz , and consider an importance decomposition of the ELBO:

where the subscripts θ and ϑ denote GNN parameters that model the generative (prior and decoder) distribution pθ(A,Z|X) and the approximate posterior (encoder) distribution qϑ(Z|X , A) which we can factorize, respectively as

Based on the ancestral factorization above, we can then expand (1) to obtain a sequential ELBO (sELBO) as follows:

where X≤t and X<t denote the partial sequences up to the tth and (t - 1)th time samples, respectively. DKL denotes the (positive-valued) Kullback-Leibler divergence (KLD), which measures the statistical discrepancy of the encoder from the prior. The conditional probabilities of (3) encapsulate the underlying topological structure and temporal coherence of the dynamic brain networks. This sELBO (3) then forms the basis of our proposed DSVB framework.

2.3 Recurrent Graph Neural Network

In this subsection, we introduce a model parameterization of the formulated sELBO based on a graph recurrent neural network. Here, the conditional latent prior and approximate posterior on each node vi are modelled as Gaussian distributions:

with isotropic covariances Σprior t,i = Diag(σprior t,i 2), Σenc t,i = Diag(σenc t,i 2 ), where Diag(·) denotes the diagonal function. The network-wide prior and approximate posterior mean and standard deviation pairs of Zt are then obtained as follows:

where the prior model φprior θ , data feature model φxθ , and latent feature model φzθ are modelled via fully-connected neural networks (FCNNs); the encoder model Φenc ϑ is modelled via the more expressive GNN. The memory-embedding recurrent states Ht = [ht,1, . . . , ht,N]⊤ ∈ RN×Dh are obtained as

and the recurrent model Φrnn θ is modelled as a spatial-aware Gated Recurrent Unit (GRU). By virtue of (6), Ht thus serves as the memory embeddings for the preceding graphs and embeddings (history path) Z<t, X<t, and A<t. Based on our model parameterization, a deterministic closed-form solution of the KLD in sELBO (3) can be analytically obtained as

where al denotes the lth element of the vector a.

Subsequently, we model the adjacency matrix decoder pθ(At|Z≤t,X<t,A<t) as a Bernoulli distribution, conditioned on the adjacency matrix reconstruction ˆAt obtained via the following inner product:

where σ denotes the sigmoid activation. In particular, we allow the adjacency matrix decoder to be conditioned on the entire history path, which suggested the incorporation of the memory-embedding Ht in the inner-product decoder, as opposed to the variational graph recurrent neural network [Hajiramezanali et al., 2019], which only considered Zt. In consideration of (8), the adjacency matrix decoder can be written in the form of a binary cross entropy (BCE) loss as follows:

where Aij and Ai denote the (i, j)th element and ith row of the matrix A, respectively. Hence, this BCE acts as the reconstruction loss between the estimated and the groundtruth adjacency graph edges. In particular, we approximate the first expectation term in (3) via Monte Carlo integration with Monte Carlo samples {Zkt , Hkt }Mk=1 , where k is the sample index and M is the number of samples. The approximate posterior latent samples Zkt ∼ qϑ(Zt|X≤t,A≤t,Z<t) are obtained via (5b) using the reparameterization Zkt = µenc,k t + σenc,k t ⊙ ϵkt , where ϵkt ∼ N(0, I) and ⊙ denotes the Hadamard (element-wise) product. The recurrent state samples Hkt are obtained via (6) given the previous time-step Zkt-1 and Hkt-1 . Substituting (9) and (7) into (3), we establish an unsupervised sELBO loss that underpins the proposed DSVB framework.

In sum, the KLD (7) compels the learnable encoder and prior model distributions (pθ, qϑ) to resemble each other, thus regularize the models from the brain network data (Xt,At) at current time-step. Furthermore, it promotes emphasis of the learned dynamic latent embeddings on the preceding history path (X<t, A<t) to incorporate strong temporal coherence, essential for propagating the learnable model distributions forward in time. On the other hand, the BCE (9) and the inner-product decoder (8) enforce the encoder model to construct latent embeddings that closely adhere to the nonEuclidean spatial characteristics of the dynamic brain networks. Together, these losses facilitate an unsupervised learning of generalized spatiotemporal embeddings that accommodate an extensive range of subject-level spatiotemporal variability across dynamic brain networks.

2.4 Attention-basedMessage Passing and Spatial-aware Gated Recurrent Unit

GNN is a class of message passing neural networks that uses aggregated topological information to construct and update node-level graph embeddings. In our proposed DSVB, we consider a GNN with message passing scheme inspired by the multi-head attention mechanism in the Transformer [Vaswani et al., 2017], of which the node embeddings are updated across layers as follows:

and αij are the attention coefficients computed as

with learnable weights {W1,W2, . . . ,W5} that project the embeddings from RDl to RDl+1 , and f (l) denotes the node embeddings at the lth GNN layer.

It follows from the scaled dot product (11) that the node embeddings fi and its neighbors fj can be interpreted as the 'query' and 'key' of a self-attention mechanism. The edge features are added to the key vectors in order to provide an additional source of information for the node. A similarity measure between the query and keys is then calculated via dot product to obtain the coefficients αij, which act as attention scores in the feature aggregation step of (10). The encoder model (5b) of the DSVB is thus modelled by a two-layered GNN with this message passing scheme, and we let f (0) = [φxθ (Xt), Ht, At ]⊤ and f (2) = [µenc t , Σenc t ]⊤ in (10).

To accurately model spatiotemporal dependencies, we parameterize the recurrent model (6) using a spatial-aware GRU, for which the recurrent states are updated across timesteps as follows:

with ¯Xt-1 = [φxθ (Xt-1), φzθ (Zt-1)]⊤. By virtue of the model equations (12), Ht thus serves as a memory embedding that retains graph-structured temporal information of the preceding latent state sequence Z<t. In comparison to conventional GRU, the FCNN therein is replaced by the set of GNNs Φxz,Φhz,Φxr,Φhr,Φxh,Φhh. Similarly, these single-layered GNNs adopt the message passing scheme (10), which allows the modified GRU to simultaneously exploit meaningful spatial structures and temporal dependencies of the dynamic graph-structured data.

2.5 Latent Embeddings for Graph Classification

Due to the recurrent nature of our proposed DSVB, the generated latent embeddings preserve non-Euclidean topological structure and rich temporal patterns which can be further exploited for dynamic graph classification. In combination, these node embeddings entail crucial spatiotemporal information of the entire dynamic graph sequence. To summarize the subject-level dynamic graphs, we apply a flattening operation on the mean of hierarchical embeddings ZT ,HT of the final time-step T to obtain the vectorized readout vec([µenc T ,µrec T ]⊤), where µrec T = 1M ∑Mk=1 HkT .

Subsequently, the flattened readout is fed into the following classification model:

modelled by multilayer FCNN. A softmax activation is then applied on the obtained logits ˆy ∈ R2 to get the predictive probability scores for the final subject-level brain network classification. The predicted class label is thus the one with highest predictive score. The classification loss is taken to be

the following multi-class cross entropy:

where c ∈ {0, 1}C is the class label and C is the number of classes. Incorporating (14) into the sELBO (3), we obtain the final DSVB-FCNN loss LDSVB in (15).

Unlike previous methods, the FCNN classifier (13) is trained jointly with the variational recurrent graph autoencoder (5)-(6) in an end-to-end fashion. Gradients of the classifier are allowed to back-propagate through time and update parameters {θ, ϑ} of the DSVB models.

2.6 Adversarial Model Regularization

Under a limited amount of training data, deep graph representation learning models are prone to generating node embeddings that aggravate overfitting. Such models typically construct a latent embedding space that overly- or grossly - overfits the limited dataset. Consequently, downstream taskspecific (e.g., graph classification) models that are adapted to this data-specific latent space perform poorly on unseen latent embeddings of general graphs. Taking this into consideration, a probabilistic adversarial training strategy is introduced to regularize the latent embedding space constructed by the DSVB models. With the adversarial regularization, the DSVB models are expected to overcome the data overfitting issue by realizing a more inclusive latent embedding space that can be readily extrapolated to unseen data beyond the limited training dataset.

Inspired by the model perturbation strategy in domain adversarial training [Ganin et al., 2016], we train the parameters {θ, ϑ, τ } of the DSVB-FCNN loss (15) in adversarial fashion:

On one hand, the DSVB parameters (θ, ϑ) are optimized to generate hierarchical embeddings ZT ,HT that fool the classifier φclassifier τ . On the other hand, the FCNN classification model parameters τ are optimized to accurately distinguish the class label of the generated latent embeddings. Such an adversarial competition is expected to achieve the Nash equilibrium that generalizes the latent embedding space and prevents model overfitting.

3 Application to Dynamic Brain Networks

3.1 fMRI Dataset &Preprocessing

We used a C-PAC pipeline [Cameron et al., 2013] preprocessed fMRI dataset from the Autism Brain Imaging Data Exchange (ABIDE I) open source database [Di Martino et al., 2009], with a sample of 144 subjects (70 ASD and 74 healthy controls (HC)) resting-state fMRI included in this case study. The inclusion criteria [Plitt et al., 2015] are: males with a fullscale IQ > 80; ages between 11 and 23; and the fMRI acquisition sites are New York University, University of California Los Angeles 1, and University of Utah, School of Medicine. We performed parcellation using Power et al.'s brain atlas [Power et al., 2011] to extract the mean time series for a total of N = 264 regions of interest (ROIs).

3.2 Graph Neural NetworkModel Training

The proposed DSVB framework is used to generate the dynamic node embeddings that correspond to the node feature and adjacency matrix sequences. The FCNN classifier is then used to discriminate between the ASD and HC subjects. We train the end-to-end DSVB-FCNN via an Adam optimizer using a learning rate of 1e-5 = 10-5 with exponential annealing, and a L2 regularization of 0.01 for 400 epochs. The output dimension of the graph-structure GRU is chosen to be 16. The output dimensions of each layer of the GNN encoder are chosen to be 32 and 16, respectively. The output dimensions of the single-layered FCNN data and latent feature models are chosen to be 64 and 8, respectively. The output dimensions of each layer of the FCNN classifier are chosen to be 32 and 2, respectively. These model hyperparameters are selected based on a 5-fold cross-validation.

3.3 Baseline Methods

SVM: The flattened upper triangle of the FC matrix is used as input to SVM. BrainNetCNN [Kawahara et al., 2017]: An extension of CNNs to handle brain graph-structured data using special kernels. GroupINN [Yan et al., 2019b]: An ensemble of GCNs to learn graph-level latent embedding representations. ASD-DiagNet [Eslami et al., 2019]: Dense autoencoder for embedding learning and single-layered perceptron for classification. Hi-GCN [Jiang et al., 2020b]: This method uses GroupINN to produce the embedding for all network instances, the learned embeddings are fed to a population-based GCN for classification. E-HI-GCN [Li et al., 2021]: An ensemble of HI-GCN, each of which is trained on different sparsity level brain networks. GAE-FCNN [Noman et al., 2022]: GAE is used to learn dynamic graph-level latent embeddings, which are then fed to a FCNN for classification. GAE-LSTM [Noman et al., 2022]: The learned embeddings of GAE are fed to a LSTM for classification.

4 Results and Discussion

To evaluate the effectiveness of our proposed method, a nested-stratified 5-fold cross-validation was applied for each experiment. The performance of ASD identification is measured using five metrics: accuracy, recall, precision, F1score, and area under the curve (AUC). Table 1 compares the average performances of the proposed DSVB-FCNN with the existing baseline classifiers based on both dynamic and static FC networks. It shows that our proposed DSVB-FCNN achieves substantial improvement over the baseline classifiers in each metric. The DSVB-FCNN also achieved low standard deviations in all metrics, demonstrating a strong consistency in its performances across the 5-fold cross-validation. In general, these results indicate that the framework facilitates robust brain disorder identification by generating graph latent embeddings that generalize well to unseen dynamic FC networks beyond the limited training dataset.

To provide insights into these performance gains of the DSVB-FCNN, we conduct an ablation study where certain components (i.e., variational Bayes, spatial-aware GRU and adversarial training) of the DSVB framework are removed to highlight their distinctive contributions to the improved dynamic FC classification result. Table 2 shows the average performances of the ablation study based on five DSVB model variants. Model I underperformed in comparison to other spatial-aware GRU-empowered models. This indicates the importance of the proposed spatial-aware GRU in endowing recurrent model's ability to capture rich spatiotemporal patterns crucial for dynamic network classification.

Figures 2a and 2b show the t-distributed Stochastic Neighbor Embedding (t-SNE) visualizations of respective latent state sequences {Zt}Tt=0 from models II and III of the ablation study. It shows that the latent state t-SNEs of model II are more concentrated towards the centre in comparison to the t-SNEs of model III, which is more uniform. This indicates that the VB elements (KLD and reparameterization) facilitate learning of a smooth latent embedding space to better accommodate subject-level spatiotemporal variability across unseen dynamic FC networks. Nevertheless, due to incorporation of the prior model via KLD loss (7), VB-based models generate latent embeddings that are highly accustomed to the spatiotemporal structure of the dynamic FC data. This aggravates model overfitting and increases the chance of false classification, resulting in an inferior performance of model III in Table 2 compared to model II.

Figures 2c and 2d show the t-SNE visualizations of respective readouts vec([µenc T ,µrec T ]⊤) from models III and V. In particular, Figure 2c shows that model III divided the readout t-SNEs into two predominant clusters based on the causal structure of each dynamic FC data. In contrast, t-SNEs of model V in Figure 2d are more regularized, which demonstrates the effectiveness of the proposed adversarial training strategy in learning indiscriminate latent embedding space that generalized well to unseen dynamic FC networks. The coordination between all components of the DSVB framework produces superior ASD identification performance, as shown in Tables 1 and 2.

In addition, we constructed higher-order dynamic FC by correlating the DSVB-learned embeddings zit between pairs of nodes, and examined the dynamic connectivity states in ASD by applying dynamic FC clustering of these dynamic networks. The optimal number of clusters = 3 is determined using the Elbow method [Bholowalia and Kumar, 2014]. Figure 3 shows the estimated connectivity matrices for three different states with the corresponding highest FC connections. By referring to [Power et al., 2011], the corresponding brain functional system can be traced back. The ASD group shows stronger FC connections in visual, default mode, salience networks; whereas higher connectivity is observed in HC for sensory and auditory networks. The difference in FC patterns between ASD and HC is statedependent. In ASD, visual ROIs show increased interconnections in state-1 and state-2 where default mode and salience are mostly activated in state-1. This is consistent with [Holiga et al., 2019], who also concluded that - compared to HC - ASD has hyperconnectivity in prefrontal and parietal cortices but hypoconnectivity in sensory-motor regions. Figure 4 shows that the latent embedding matrices derived from DSVB are more visually discriminative compared to the raw FC matrices. The ASD embeddings highlight hyperconnectivity within visual networks (ROIs 143 to 173), particularly on the occipital lobe, and this aligns with earlier neuroscience discoveries [Keehn et al., 2008; Clery et al., 2013; Matsuoka et al., 2020].

5 Conclusion

We have developed a deep probabilistic spatiotemporal framework based on sequential variational Bayes for dynamic graph representation learning. The proposed DSVB framework incorporates a spatial-aware GRU to capture both topological and temporal alterations across brain networks. A downstream FCNN then leverages the learned graph embeddings to reveal atypical neural connectivity patterns. Evaluation on resting-state fMRI data shows substantial improvements over the existing baseline methods, suggesting potential in neuropsychiatric disorder identification. Despite our current focus on brain disorders, this framework is domainagnostic and applicable to other dynamic graph-structured data. Moreover, the proposed framework could be extended to handle large graphs by incorporating sparse attention and heterogeneous graphs via cross-modality representations.

Acknowledgements

The work of Junn Yong Loo is supported by Monash University under the SIT Collaborative Research Seed Grants 2024 I-M010-SED-000242. The work of Chee-Ming Ting is supported by the Ministry of Higher Education, Malaysia under the Fundamental Research Grant Scheme FRGS/1/2023/ICT02/MUSM/02/1. Sin-Yee Yap has been based at Monash's Department of Data Science and AI for some of the duration of this work.

Contribution Statement

Sin-Yee Yap, Junn Yong Loo (corresponding author) and Chee-Ming Ting contributed equally to this work.