Deep Representations for Time-varying Brain Datasets

1 INTRODUCTION . Neuroimaging techniques such as fMRI ( functional magnetic resonance imaging ) and DWI ( diffusion-weighted imaging ) provide a window into complex brain processes . Yet , modeling and understanding these signals has always been a challenge . Network neuroscience ( Bassett & Sporns , 2017 ) views the brain as a multiscale networked system and models these signals in their graph representations : nodes represent brain ROIs ( regions of interest ) , and edges represent either structural or functional connections between pairs of regions . With larger imaging datasets and developments in Graph Neural Networks ( Scarselli et al. , 2009 ) , recent works leverage variants of the graph deep learning , modeling brain signals with data-driven models and getting rid of Gaussian assumptions typically existed in linear models ( Zhang et al. , 2019 ; Li et al. , 2019 ) . These methods are making progress on identifying physiological characteristics and brain disorders : In Kim & Ye ( 2020 ) , authors combine grad-CAM ( Selvaraju et al. , 2017 ) and GIN ( Xu et al. , 2018 ) to highlight brain regions that are responsible for gender classification with resting-state fMRI data . Li et al . ( 2020 ) utilizes the regularized pooling with GNN to identify fMRI biomarkers . Noman et al . ( 2021 ) embeds both topological structures and node signals of fMRI networks into low-dimensional latent representations for a better identification of depression . However , the first two works use time-averaged fMRI , losing rich dynamics in the temporal domain . The third combines nodes ’ temporal and feature dimensions instead of handling them separately , leading to a suboptimal representation ( as discussed in section 3.2 ) . To overcome these issues , we propose ReBraiD ( Deep Representations for Time-varying Brain Datasets ) , a graph neural network model that jointly models dynamic functional signals and structural connectivities , leading to a more comprehensive deep representation of brain dynamics . To simultaneously encode signals along spatial and temporal dimensions , some notable works in traffic prediction and activity recognition domains such as Graph WaveNet ( Wu et al. , 2019b ) alternate TCN ( temporal convolution network ) ( Lea et al. , 2016 ) and GCN ( graph convolutional network ) ( Kipf & Welling , 2017 ) . Others ( Song et al. , 2020 ; Liu et al. , 2020 ) use localized spatial-temporal graph to embed both domains ’ information in this extended graph . There are also works incorporating gated recurrent networks for the temporal domain such as ( Seo et al. , 2018 ; Ruiz et al. , 2020 ) . We choose the first option for ReBraiD , as it is more memory and time efficient , and can support much longer inputs . We also explore the best option when alternating spatial and temporal layers for encoding brain activities with extensive ablation studies . Upon this structure , we propose novel sample-level adaptive adjacency matrix learning and multi-resolution inner cluster smoothing , both of which learn and refine latent dynamic structures . We also make the model more efficient while being effective with the choice of the temporal layer . Equally important as finding a good representation of brain dynamics is interpreting them . We utilize integrated gradients ( Sundararajan et al. , 2017 ) to identify how brain ROIs participate in various processes . This can lead to better behavioral understanding , biomarker discoveries , and characterization of individuals or groups with their brain imagings . We also make the novel contribution of identifying temporally important frames with graph attribution techniques ; this can enable more fine-grained temporal analysis around keyframes when combined with other imaging modalities such as EEG ( electroencephalogram ) . In addition , our subject-level and group-level attribution studies unveil heterogeneities among ROIs , tasks , and individuals . 2 METHOD . 2.1 PRELIMINARIES . We utilize two brain imaging modalities mapped onto a same coordinate : SC ( structural connectivity ) from DWI scans , and time-varying fMRI scans . We represent them as a set of L graphs Gi = ( Ai , Xi ) with i ∈ [ 1 , L ] , in which Ai ∈ RN×N represents normalized adjacency matrix with an added self-loop : Ai = D̃ − 1 2 SCi ˜SCiD̃ − 1 2 SCi , ˜SCi = SCi + IN and D̃SCi = ∑ w ( ˜SCi ) vw is the diagonal node degree matrix . Graph signal matrix obtained from fMRI scans of the ith sample is represented as Xi ∈ RN×T . Here N is the number of nodes , and each node represents a brain region ; T is the input signal length on each node . Our objective focuses on classifying brain signals Gi into one of C task classes through learning latent graph structures . 2.2 MODEL . ReBraiD takes ( A , X ) as inputs , and outputs task class predictions . The overall model structure is shown in fig . 1 . For the ith sample Xi ∈ RN×1×T , the initial 1 × 1 convolution layer increases 1Axis order follows PyTorch conventions . Dimension at the second index is the expanded feature dimension . its hidden feature dimension to dh1 , outputting ( N , dh1 , T ) . The encoder then encodes temporal and spatial information alternately , reducing information to feature axis and generating a hidden representation of size ( N , dh2 , 1 ) . The encoder is followed by two linear layers to perform node embedding pooling and two MLP layers for classification . Cross entropy is used as the loss function : L ( i ) CE = − ∑ C yi log ŷi , where yi ∈ RC is the one-hot vector of ground truth task label and ŷi ∈ RC is model ’ s predicted distribution . We now explain different components of the model . ( I ) Learning sample-level latent graph structures . Structural scans serve as our graph adjacency matrices . However , they are static not only across temporal frames but also across different tasks . In contrast , functional connectivities ( FC ) are highly dynamic , as shown in appendix A.1.1 . To better capture dynamic graph structures , we assign an adaptive adjacency matrix to each input graph signal . Unlike other works such as Wu et al . ( 2019b ) that use a universal latent graph structure , our model do not assume all samples share the same latent graph . Instead , in ReBraiD , each sample has a unique latent structure reflecting its signal status . This implies that the latent adjacency matrix can not be treated as a learnable parameter as a part of the model . To solve this , we minimize the assumption down to a shared projection Θadp that projects each input sequence into an embedding space and use this embedding to generate the latent graph structure . Θadp can be learned in an endto-end manner . The generated adaptive adjacency matrix for the ith sample can be written as , where Softmax is applied column-wise : Ai adp = Softmax ( ReLU ( ( XiΘadp ) ( XiΘadp ) > ) ) , Θadp ∈ RT×hadp ( 1 ) ( II ) Gated TCN ( Temporal Convolutional Network ) . To encode temporal information , we use the gating mechanism as in Oord et al . ( 2016 ) in our temporal layers : H ( l+1 ) = tanh ( TCNemb ( H ( l ) ) ) σ ( TCNgate ( H ( l ) ) ) , where H ( l ) ∈ RN×d×t is one sample ’ s activation matrix of the lth layer , denotes the Hadamard product , and σ is the Sigmoid function . Different from TCNs generally used in sequence to sequence models that consist of dilated Conv1d and causal padding along the temporal dimension ( van den Oord et al . ( 2016 ) ) , we simply apply Conv1d with kernel 2 and stride 2 as our TCNemb and TCNgate to embed temporal information . The reason is twofold : first , for a sequence to sequence model with a length-T output , yτ should only depend on xt≤τ to avoid information leakage , and causal convolution can ensure this . In contrast , our model ’ s task is classification , and the goal of our encoder along the temporal dimension is to embed signal information into the feature axis while reducing temporal dimension to 1 . The receptive field of this single temporal point ( with multiple feature channels ) is meant to be the entire input sequence . Essentially , our TCN is the same as the last output node of a kernel 2 causal TCN whose dilation increases by 2 at each layer ( fig . 2 ) . Second , from a practical perspective , directly using strided non-causal TCN works the same as using dilated causal TCNs and taking the last node , while simplifies the model structure and reduces training time to less than a quarter . ( III ) Graph Network layer . In our model , every l temporal layers ( appendix A.2.3 studies the best l to choose ) are followed by a spatial layer to encode signals with the graph structure . Building temporal and spatial layers alternatively helps spatial modules to learn embeddings at different temporal scales , and this generates better results than putting spatial layers after all temporal ones . To encode spatial information , Kipf & Welling ( 2017 ) uses first-order approximation of spectral filters to form the layer-wise propagation rule of a GCN layer : H ( l+1 ) = GCN ( H ( l ) ) = f ( AH ( l ) W ( l ) ) . It can be understood as spatially aggregating information among neighboring nodes to form new node embeddings . In the original setting without temporal signals , H ( l ) ∈ RN×d is the activation matrix of lth layer , A ∈ RN×N denotes the normalized adjacency matrix with self-connections as discussed in section 2.1 , W ( l ) ∈ Rd×d ′ is learnable model parameters , and f is a nonlinear activation function of choice . Parameters d and d′ are the number of feature channels . We view a GCN layer as a local smoothing operation followed by an MLP , and simplify stacking K layers to AKH as in Wu et al . ( 2019a ) . In ReBraiD , every graph network layer aggregates information from each node ’ s K-hop neighborhoods based on both brain structural connectivity and the latent adaptive adjacency matrix : namely we have both AiKH ( l ) WK and Ai adpKH ( l ) WK adp for input H ( l ) . We also gather different levels ( from 0 to K ) of neighbor information with concatenation . In other words , one graph convolution layer here corresponds to a small module that is equivalent to K simple GCN layers with residual connections . We can write our layer as : H ( l+1 ) = GNN ( l ) ( H ( l ) ) = MLP [ ConcatKk=1 ( H ( l ) , ReLU ( Ai kH ( l ) ) , ReLU ( Ai adp kH ( l ) ) ) ] ( 2 ) With the additional temporal dimension , H ( l ) ∈ RN×d×t in eq . ( 2 ) , and Ai ∈ RN×N applies on H ( l ) ’ s first two dimensions while multiplying . Outputs of different GNN ( l ) layers are parametrized and then skip connected with a summation . Since the temporal lengths of these outputs are different because of TCNs , max-pooling is used before each summation to make the lengths identical . ( IV ) Multi-resolution inner cluster smoothing . While GNN layers can effectively passing information between neighboring nodes , long-range relationships among brain regions that neither appear in SC nor learned by latent Aadp can be better captured using soft assignments similar to DIFFPOOL ( Ying et al . ( 2018 ) ) . To generate the soft assignment tensor assigning N nodes into c clusters ( c chosen manually ) , we use GNN ( l ) pool that obeys the same propagation rule as in eq . ( 2 ) , followed by Softmax along c. This assignment is applied to the output of GNN ( l ) emb which carries out the spatial embedding for the lth layer input H ( l ) : S ( l ) = Softmax ( GNN ( l ) pool ( H ( ` ) ) , 1 ) ∈ RN×c×t Z ( l ) = GNN ( l ) emb ( H ( l ) ) ∈ RN×d×t H̃ ( l ) = S ( l ) > Z ( l ) ∈ Rc×d×t ( 3 ) The extra temporal dimension allows nodes to be assigned to heterogeneous clusters at different frames . We find that using coarsened A ( l+1 ) i = S ( l ) > A ( l ) i S ( l ) ∈ Rc×c as the graph adjacency matrix leads to worse performance compared to using SC-generated Ai and learned Ai adp ( comparison in section 3.1 ) . In addition , if the number of nodes is changed , residual connections coming from the beginning of temporal-spatial blocks can not be used and this impacts overall performance . To continue use Ai and Ai adp as graph adjacency matrices and allow residual connections , we reverseassign H̃ ( l ) with assignment tensor obtained from applying Softmax on S ( l ) > along N , so that the number of nodes is kept unchanged : S̃ ( l ) = Softmax ( S ( l ) > , 1 ) ∈ Rc×N×t H ( ` +1 ) = S̃ ( l ) > H̃ ( l ) ∈ RN×d×t ( 4 ) In fact , eqs . ( 3 ) and ( 4 ) perform signal smoothing on nodes within each soft-assigned cluster ( appendix A.1.2 shows a toy example ) . With the bottleneck c < N , the model is forced to pick up latent community structures . This inner-cluster smoothing is carried out at different spatial resolutions : as the spatial receptive field increases with more graph layers , we decrease cluster number c for the assignment operation .

Summary: This paper proposes a Graph Neural Network model to estimate latent dynamics in the human brain using functional Magnetic Resonance Imaging (fMRI) and Diffusion Weighted Imaging (DWI) while performing a classification task. The model consists of four parts 1) estimating each sample's adjacency matrix by learning a shared projection matrix 2) gated temporal convolutional neural network for learning information from time-series data 3) graph neural network to get embeddings of temporal and spatial information combined by 4) multi-resolution inner cluster smoothing. The authors use integrated gradients for estimating important features for the classification task helping with model interpretation. The authors show that the proposed method can perform better in the prediction task than GCN, GAT V2, etc. On the real dataset, it is shown that the method can capture temporal and spatial heterogeneity and provide regions important for the presented classification task.