Learning Dynamic Graph Representation of Brain Connectome with Spatio-Temporal Attention

1 Introduction . Neuroimaging modalities provide measurements of brain activity by capturing the signals of neural activity . Functional magnetic resonance imaging ( fMRI ) is a non-invasive imaging method that measures the blood-oxygen level dependence ( BOLD ) in order to estimate the neural activity of the whole brain over time [ 20 ] . Functional connectivity ( FC ) is defined as the degree of temporal ∗This work was conducted while the first author was at KAIST 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . correlation between regions of the brain . Based on the fact that these connectivities form networks that change over time , graph-based network analysis of brain connectome has been one of the key approaches to understanding how the brain works [ 7 , 4 , 43 ] . Graph neural networks ( GNNs ) are a type of deep neural networks that have recently been successful in learning the representation of graph-structured data [ 50 ] . The graph-structured nature of the brain has led to an increased interest in learning the reperesentation of the brain FC network with the GNNs . Learning the representation of the brain connectome can be linked to decoding trait or state from human brain signal measurements . Accordingly , the current trend in studies attempting to apply GNN to the brain connectome is to input the FC graph from either resting-state [ 25 , 26 , 2 , 33 , 23 , 48 , 49 ] or task fMRI data [ 29 , 30 , 32 ] and predict a particular phenotype of the subjects , such as gender [ 26 , 2 , 23 , 22 ] or presence of a specific disease [ 26 , 33 , 29 , 30 , 32 , 22 ] . While these studies have shown potential strengths and opportunities for learning the network representation of the brain , they also suggest limitations of current GNN-based methods . One of the most common limitations with previous GNN-based FC network analysis methods is that most of them fail to take advantage of the dynamic properties of the FC network , which fluctuates over time . Incorporating the dynamic features of the FC network into the neuroimaging analysis has been an important direction in the field of functional neuroimaging [ 21 , 37 ] . A work by [ 15 ] tried to address this issue by using the Spatial Temporal Graph Convolutional Network ( ST-GCN ) [ 54 ] model to incorporate dynamic features of the FC network . However , [ 15 ] reported lower accuracy than other non-dynamic GNN-based FC methods [ 23 , 2 ] in the gender classification experiment , leaving a question about the effectiveness of the dynamic FC method . In addition , another limitation of the method is that no temporal explainability is provided from the model . This is a major drawback considering that the goal of applying GNNs to functional neuroimaging methods is not only to achieve high classification accuracy , but also to uncover the functional basis of the brain [ 23 , 32 ] . Another recent work by [ 3 ] , using GraphNets [ 5 ] and DiffPool [ 58 ] for the dynamic FC analysis , also suffers from the same limitations in terms of poor classification accuracy and lack of temporal explainability . Here , we propose Spatio-Temporal Attention Graph Isomorphism Network ( STAGIN ) for learning the dynamic graph representation of the brain connectome with spatio-temporal attention . The proposed method exploits the temporal features of the dynamic FC network graphs to improve the classification accuracy of the model . In particular , we address the issue that the node features of the input dynamic graph should contain temporal information and concatenate encoded timestamp with the node features ( Section 4.1 ) . In addition , the proposed method includes novel attention-based READOUT modules ( Section 4.2 ) and the Transformer encoder [ 46 ] ( Section 4.3 ) in order to further improve the classification performance and provide spatial-temporal explainability at the same time . STAGIN achieves state-of-the-art performance with the Human Connectome Project ( HCP ) dataset [ 45 ] in gender classification for resting-state fMRI and task decoding for task fMRI . We inherit k-means clustering analysis of the resting-state dynamic FC [ 1 ] and general linear model ( GLM ) statistical mapping of task fMRI [ 14 ] for interpreting the spatio-temporal attention learned from STAGIN , which are widely accepted analysis methods for the fMRI data . The interpretation of the learned spatio-temporal attention replicates neuroscientific findings from previous large-scale fMRI studies in both resting-state and task fMRI , which further validates our proposed method . Our work holds potential societal impact in that brain decoding methods can be linked to finding neural biomarkers of important phenotypes or diseases . However , potential negative impact related to privacy concerns that arise from abuse or misuse of accurate decoding methods should also be noted . Although our method is yet behind the decoding capability that can be abused or misused , our research can not still be free from these ethical considerations . 2 Related works . 2.1 Graph Neural Network on Dynamic Graphs . Many networks that arise around us are inherently dynamic , with changes in the existence of nodes and edges over time . Learning the representation of dynamic graphs has piqued the interest of researchers and has led to development of methods that can embed dynamic graphs using their time information [ 35 ] . Methods that incorporate attention for learning the representation of dynamic graphs have also been proposed [ 51 , 40 ] . However , it is not easy to apply these techniques directly to the dynamic brain graphs because of the different inherent properties of the dynamic brain graphs that do not include any addition or deletion of nodes and are sampled uniformly over time . Nonetheless , our work is inspired by these earlier studies , particularly for the encoding of temporal information and their concatenation to the node features , as proposed in Section 4.1 [ 51 , 40 ] . 2.2 Attention in Graph Neural Networks . Bringing attention to the GNNs is a topic that is being actively studied in the field of geometric deep learning [ 27 ] . One of the most successful uses of attention is to compute the attention at edges of the graph and scale the importance of the links when the features of the neighborhood node are aggregated [ 47 , 6 ] , often providing performance gain in learning the representation of input graphs . Another stream of applying attention to the GNNs comes with the motivation to define a pooling function on the graph domain . Since it is not straightforward to decide on what basis the coarsening should be carried out for graph structured data , works such as [ 16 , 28 , 38 ] have addressed this problem by selecting the nodes with top scores computed from projecting the node feature vectors into a learnable parameter vector , or from a GNN layer aggregated local graph features . Although the motivation may have been different , these graph pooling methods are closely related to the spatial attention modules that we propose in Section 4.2 in that they exploit learned relative scores across the vertices of the graph . While some works have already been aware that the appropriate use of node-wise attention can improve performance of downstream tasks [ 55 , 11 ] , we note that previous methods tend to score attention based on randomly initialized parameters or local graph structures which may be suboptimal for graph classification tasks that require taking the whole graph feature into account . 3 Theory . 3.1 Problem definition . The goal of our study is to train a neural network f : Gdyn → hGdyn , where Gdyn = ( G ( 1 ) , ... , G ( T ) ) is the sequence of brain graphs with T timepoints and hGdyn ∈ RD is the vector representation of the dynamic graph G ( t ) with length D. The graph G ( t ) = ( V ( t ) , E ( t ) ) at time t is a pair of vertex set V ( t ) = { x1 ( t ) , ... , xN ( t ) } of N nodes and edge set E ( t ) = { { xi ( t ) , xj ( t ) } | j ∈ N ( i ) , i ∈ { 1 , ... N } } where N ( i ) denotes the neighborhood of the vertex i . If f learns to extract a disentangled representation of the dynamic brain graph Gdyn , then the classification of a certain phenotypic characteristic ( e.g . gender ) from hGdyn can be performed with a linear mapping as a downstream task . Another important consideration in this work is to ensure the explainability of the model f , being able to inform us which part of the brain at which timepoint was considered important when extracting the meaningful representation hGdyn . Specifically , we formulate f = q ◦ g as a composition of the GNN g and the Transformer encoder q , where g outputs the set of graph representations hG ( t ) from each timepoint and q exploits self-attention to integrate hG ( t ) into the final representation hGdyn : g : Gdyn → ( hG ( 1 ) , ... , hG ( T ) ) , ( 1 ) q : ( hG ( 1 ) , ... , hG ( T ) ) → hGdyn . ( 2 ) We will omit timepoint notation ( t ) for brevity , whenever it is not of contextual importance . 3.2 Graph Isomorphism Network . The GNNs are generally composed of functions that ( i ) integrate the node features from its neighbors , and ( ii ) embed the integrated information with a nonlinear transformation to obtain the next layer node features . These functions are called AGGREGATE , and COMBINE functions , respectively , and the choice of these functions define many variants of the GNN , a ( k ) v = AGGREGATE ( k ) ( { h ( k−1 ) u : u ∈ N ( v ) } ) , ( 3 ) h ( k ) v = COMBINE ( k ) ( h ( k−1 ) v , a ( k ) v ) , ( 4 ) where h ( k ) v denotes the feature vector of node v at layer k and h ( 0 ) v : = xv . The Graph Isomorphism Network ( GIN ) is a variant of the GNN suitable for graph classification tasks , which is known to be as powerful as the WL-test under certain assumptions of injectivity [ 52 ] . The GIN typically defines sum as the AGGREGATE and a multi-layer perceptron ( MLP ) with two layers as the COMBINE updating the node representation h ( k ) v at layer k [ 52 ] by : h ( k ) v = MLP ( k ) ( ( 1 + ( k ) ) · h ( k−1 ) v + ∑ u∈N ( v ) h ( k−1 ) u ) , ( 5 ) where is a learnable parameter initialized with zero . Equation ( 5 ) can be easily reformulated into the matrix form [ 23 ] by : H ( k ) = σ ( ( ( k ) · I +A ) H ( k−1 ) W ( k ) ) , ( 6 ) where H ( k ) = [ h ( k ) 1 , · · · , h ( k ) N ] ∈ RD×N is the stack of node feature vectors , I is the identity matrix , A is the adjacency matrix between the node features , W is the network weights of the MLP , and σ is the nonlinearity function . The READOUT function takes the updated node features h ( k ) v to compute the representation of the whole graph : h ( k ) G = READOUT ( { h ( k ) v | v ∈ G } ) . ( 7 ) In general , the READOUT function is defined simply as computing the sum or average of the input node features . This is equivalent to multiplication with the length N pooling vectors φ > sum = [ 1 , ... , 1 ] or φ > mean = [ 1/N , ... , 1/N ] for the matrix form : h ( k ) G = H ( k ) φmean . ( 8 )

Summary: The paper proposes an attention based GNN method, which it names STAGIN. The new method handles dynamic FC matrices for gender and task classification using fMRI data and learns to focus on important time points and brain regions based on attention. The method achieves better accuracy and AUC for classification compared to the relevant competing methods that STAGIN is compared against. The temporal attention matches with the task-relevant time points in the input data.