Year: 2021

Abstract:
Finding an appropriate representation of dynamic activities in the brain is crucial for many downstream applications. Due to its highly dynamic nature, temporally averaged fMRI (functional magnetic resonance imaging) cannot capture the whole picture of underlying brain activities, and previous works lack the ability to learn and interpret the latent dynamics in brain architectures. In this paper, we build an efficient graph neural network model that incorporates both region-mapped fMRI sequences and structural connectivities obtained from DWI (diffusion-weighted imaging) as inputs. Through novel sample-level adaptive adjacency matrix learning and multi-resolution inner cluster smoothing, we find good representations of the latent brain dynamics. We also attribute inputs with integrated gradients, which enables us to infer (1) highly involved brain connections and subnetworks for each task (2) keyframes of imaging sequences along the temporal axis, and (3) subnetworks that discriminate between individual subjects. This ability to identify critical subnetworks that characterize brain states across heterogeneous tasks and individuals is of great importance to neuroscience research. Extensive experiments and ablation studies demonstrate our proposed method’s superiority and efficiency in spatial-temporal graph signal modeling with insightful interpretations of brain dynamics.

heading: 1 INTRODUCTION
text: 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.

heading: 2 METHOD

heading: 2.1 PRELIMINARIES
text: 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.

heading: 2.2 MODEL
text: 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 cannot 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.

heading: 2.3 ATTRIBUTION WITH IG (INTEGRATED GRADIENTS).
text: As one approach to model interpretability, attribution assigns credits to each part of the input, assessing how important they are to the final predictions. Wiltschko et al. (2020) gives an extensive comparison between different graph attribution approaches, in which IG (Sundararajan et al. (2017)) is top-performing and can be applied to trained models without extra alterations of the model structure. IG also has other desirable properties such as implementation invariance that other gradient methods are lacking. It is also more rigorous and accurate than obtaining explanations from attention weights or pooling matrices that span multiple feature channels. Intuitively, IG calculates how real inputs contribute differently compared to a selected baseline; it does so by aggregating model gradients at linearly interpolated inputs between the real and baseline inputs. For each sample, we
calculate attributions at each point of both input A ∈ RN×N and X ∈ RN×T :
ATTRAvw = ( Avw −A′vw ) × M∑ m=1 ∂F (AIntrpl, X) ∂AIntrplvw × 1 M , AIntrpl = A ′ + m M × ( A−A′ ) ATTRXvt = ( Xvt −X ′vt ) ×
M∑ m=1 ∂F (A,XIntrpl) ∂XIntrplvt × 1 M , XIntrpl = X ′ + m M × ( X −X ′ ) (5)
F (A,X) here represents our signal classification model,M is the step number when doing Riemann approximation of the path integral, and A′, X ′ are baselines of A,X (see section 3.3 for more details). Note that eq. (5) calculates the attribution of one point on one sample. The process is repeated for every point of the input, so attributions have identical dimensions as inputs. To obtain brain region importance of a task, we aggregate attributions across multiple samples of that task.

heading: 3 EXPERIMENTS
text: We use fMRI signals from CRASH dataset (Lauharatanahirun et al. (2020)) for our experiments. The model classifies input fMRI into 6 tasks: resting state, visual working memory task (VWM), dynamic attention task (DYN), math task (MOD), dot probe task (DOT), and psychomotor vigilance task (PVT) (appendix A.2.1 has detailed task descriptions). We preprocess 4D voxel-level fMRI images into graph signals G = (A,X) by averaging voxel activities into regional signals with the 200-ROI cortical parcellation (voxel to region mapping) specified by Schaefer et al. (2018). We also standardize signals for each region and discard scan sessions with obvious abnormal spikes that may be caused by head movement, etc.. DWI scans are mapped into the same MNI152 coordinate and processed into adjacency matrices with the same parcellation as fMRI uses. Valid data contains 1940 scan sessions from 56 subjects, session length varies from 265 frames to 828 frames and TR (Repetition Time) is 0.91s. These 1940 scan sessions are separated into training, validation, and test sets with a ratio of 0.7-0.15-0.15. Hyperparameters including dropout rate, learning rate, and weight decay are chosen with grid search based on validation loss, and all results reported in this section are obtained from the test set. For each scan session, we use a sliding window to generate input sequences (in the following experiments T ∈ {8, 16, 32, 64, 128, 256}) and feed them to the model. To encode temporal and spatial information alternatively, we find stacking two TCN layers per one GNN layer leads to better performance most times (fig. 3b, see more on appendix A.2.3 (II)). Models are written in PyTorch, trained with Google Colab GPU runtimes, and 60 epochs are run for each experiment setting. Codes and data will be released upon acceptance.

heading: 3.1 MODEL COMPONENTS
text: Graph adjacency matrices. For each input sample Gi, we test different options to provide graph adjacency matrices to the GNN layer. They include (i) our proposed method: using both adaptive adjacency matrix Ai adp and SC-induced Ai; (ii) only using Ai; (iii) only using Ai adp; (iv) replacing Ai adp in setting i with Ai FC derived from functional connectivity; (v) only using random graph
adjacency matrices with the same level of sparsity as real As. We set hadp to be 5 in eq. (1), which works better for our data than larger hadp choices. K is set to 2 for eq. (2), meaning each GNN aggregates information from 2-hop neighbors based on the provided adjacency matrices. We evaluate our model with weighted F1 as the metric because of the imbalance among tasks. The results under different settings are reported in fig. 3a (and table 3 in appendix for numerical values).
From the results of setting (ii) plotted in fig. 3a, we see that removing the adaptive adjacency matrix impacts the performance differently at different input lengths: the gap peaks for signals of length 64 - 128, and becomes smaller for either shorter or longer sequences. This could suggest the existence of more distinct latent states of brain signals of this length that cannot be captured by structural connectivities. On the other hand, removing SC (setting (iii)) seems to have a more constant impact on the model performance, with shorter inputs more likely to see a slightly larger drop. In general, only using Aadp leads to smaller performance drop than only using SC, indicating the effectiveness of Aadp in capturing useful latent graph structures. More detailed studies in appendix A.2.3 shows Aadp learns distinct representations not captured by A.
As mentioned in section 2, our motivation behind creating sample-level adaptive adjacency matrices is FC’s highly dynamic nature. Therefore, for setting (iv), we test directly using adjacency matrices Ai FC obtained from FC instead of the learned Ai adp. In particular, Ai FC = D̃ − 1 2
FCi ˜FCiD̃
− 1 2 FCi ∈ R200×200, where (FCi)vw = corr((Xi)v, (Xi)w), ˜FCi = FCi + IN and D̃FCi = ∑ w( ˜FCi)vw. Fig. 3a shows Ai FC constantly underperforms Ai adp, except for being really close with length-8 inputs. Larger performance gaps are observed for longer inputs, where Corr((Xi)v, (Xi)w) struggles to capture the changing dynamics in the inputs. This demonstrates that our input-based latent Ai adp has better representation power than input-based FC. We also notice batch correlation coefficients calculation for Ai FC results in a slower training speed than computing Ai adp.
An interesting result comes from setting (v), where we use randomly generated Erdős-Rényi graphs with the edge creation probability the same as averaged edge existence probability of As. Its performance is at a similar level or even better than settings (ii) and (iii). Our hypothesis is the model can learn the latent graph structure out of randomness, and we will verify this hypothesis in section 3.3.
Multi-resolution inner cluster smoothing. To verify the capability of inner-cluster smoothing operation in capturing latent graph dynamics, we test the following settings: (vi) using our proposed model and inputs, except removing paralleled GNNpool and inner-cluster smoothing module; (vii) previous setting (v) but remove GNNpool and inner-cluster smoothing module; (viii) keep GNNpool, but using coarsened graph instead of performing smoothing and increasing the node number back (essentially performing DIFFPOOL with an added temporal dimension). In this last setting, we hierarchically pool graph nodes until node number reaches 1, and we keep the total number of GNN layers the same as our other settings. Values of soft-assigned cluster number c are chosen to be halved per smoothing module, starting from half of the graph nodes number, namely: 100, 50, 25, etc. for our experiments. Different choices of c affect model converging rate, but only have minor impacts on the final performance (see appendix A.2.3 (III)). Results are reported in fig. 3a (and table 3 in appendix). Apart from these three settings, we also test adding pooling regularization terms (described in appendix A.1.3) into the loss function but this does not lead to much of a difference.
The above results demonstrate that both setting (vi) and (vii) outperforms (viii) by a large margin, indicating the importance of keeping the original node number when representing brain signals. In addition, all three settings underperform our proposed method, and they are mostly worse than changing graph adjacency matrices as in settings (ii)-(v): this shows inner-cluster smoothing module has a large impact in learning latent graph dynamics. We also find using adaptive adjacency matrices and inner cluster smoothing can stabilize training, making the model less prone to over-fitting and achieving close-to-best performance over a larger range of hyperparameters (see fig. 11).

heading: 3.2 MODEL COMPARISONS
text: In this section, we compare our model with the vanilla GCN from Kipf & Welling (2017), Chebyshev Graph Convolutional Gated Recurrent Unit (GConvGRU) from Seo et al. (2018), GraphSAGE from Hamilton et al. (2017), GAT V2 from Brody et al. (2021) and Graph Transformer as in Shi et al. (2021). To use them in our fMRI classificaiton problem, we directly take corresponding layers from
PyTorch Geometric 2 and PyTorch Geometric Temporal 3 and construct the models similar to ours: four graph layers taking in both signals and adjacency matrices, followed by two linear layers along node axis and two linear layers for the final classification. We train baseline models with the same setting as our model: 256-frame inputs, Adam optimizer, cross entropy loss, and 60 epochs (all models well-converged). Grid search is used to optimize the rest of hyperparameters. We compare accuracy, weighted F1 score, and training time per epoch in table 1; we also plot our model and Graph Transformer’s confusion matrices in appendix fig. 12.
We observe that our proposed method significantly outperforms the baseline graph models by a margin of 20 to 40 percent and has much less training time. This demonstrates our proposed model’s effectiveness in capturing latent brain dynamics. For these baseline models, temporal content is used as features; the comparison shows separating them into different axis is more advantageous. This is further confirmed with models only having TCN layers: we test both removing GNN layers all together and replacing them with 1×1 CNN layers. Both outperform graph models that focuses on the spatial modeling aspect. Although temporal modeling is crucial, including the spatial information in its graph format as our proposed model can improve the performance much further.

heading: 3.3 INTERPRETATIONS WITH IG
text: In this section, we study the contributions of different brain ROIs and subnetworks defined by their functionalities. For the subnetwork definition, we choose to use the 17 networks specified in Thomas Yeo et al. (2011)) which has a mapping from our previous 200-ROI parcellation. See table 4 in appendix for all subnetwork names. We compute IG of a model trained on length-256 input signals because the model has higher performance with longer inputs, leading to more accurate attributions. To select baseline inputs, we follow the general baseline selection principle for attribution methods: when the model takes in a baseline input, it should produce a near-zero prediction and Softmax(outputs) should give each class about the same probability in a classification model. Allzero A′ and X ′ can roughly achieve this for our model, so we choose them as our baseline inputs. For each task, the IG computation is done on 900 inputs to get an overall distribution.
Temporal importance. On the single input level, we can attribute which parts of the inputs in Gi are more important in predicting the target class by looking into (ATTRX)i. This attribution map not only shows which brain regions contribute more but also reveals the important signal frames. One critical drawback of fMRI imaging is its low temporal resolution, but if we can know which part is more important, we can turn to more temporally fine-grained signals such as EEG to see if there are any special activities during that time. To confirm that the attributions we get are valid and consistent, we perform a sanity check of IG results on two overlapped inputs with an offset τ : the first input is obtained from window [t0, t0 + T ] and the second is obtained from window [t0 + τ, t0 + τ +T ]. Offset aligned results are shown in fig. 4a, and we can see the attributions agree with each other quite well.
Spatial importance. We examine the connection importance between brain ROIs by looking at ATTRA (task-averaged ATTRA are plotted in fig. 14 in appendix). In particular, columns in ATTRA with higher average values are sender ROIs of high-contributing connections, which is what matters in the GNN operation. We also explore why using random graph adjacency matrices (setting (v) in section 3.1) can produce a similar result for length-256 inputs compared to using both SC-induced
2https://pytorch-geometric.readthedocs.io/ 3https://pytorch-geometric-temporal.readthedocs.io/
Ai and Ai adp (setting (i)). By examining ATTRA under both settings (fig. 5), we see that the column averages of ATTRA under these two settings are similar for almost all tasks, meaning the model can learn where to pay attention to even using random adjacency matrix inputs. We credit this ability partially to multi-resolution inner cluster smoothing, as the performance would drop notably without it (setting (vii)). However, using ground truth SC not only gives us higher performance for shorter inputs but also provides the opportunity to better interpret brain region connections. We can directly use task-averaged ATTRA as the weighted adjacency matrix to plot edges between brain ROIs, just as in fig. 6. Important brain regions obtained from ATTRA mostly comply with the previous literature (see appendix A.2.5 for details).
ATTRX can also give us insights on spatial importance when the attribution maps are averaged or summed up along the temporal dimension. But it does so from another perspective: instead of showing important structural connections that support information passing, it reveals regions or subnetworks that are sources of the important signals. In fig. 4b, we plot the distribution of taveraged and subnetwork-averaged (mapping 200 ROIs into 17 subnetworks) ATTRX during VWM task. We can see the clear dominance of VisCent, DorsAttnA, and ContA subnetworks (numbered as 1, 5, 11), indicating signals from these regions are useful for model to decide if the input is from VWM task. For the boxplots of other tasks and subnetwork rankings, please see fig. 16, table 5 and table 6 in appendix. More informative than the rankings is the distribution itself: even though VisCent, DorsAttnA, and ContA ranked top 3 for both resting state and VWM task for signal attributions, their relative importance and attribution distribution variances are totally different. In a sense, the distribution can act as a task fingerprint based on brain signal states.
We notice that signal-important ROIs are not necessarily the same as connection-important ROIs: top-ranked subnetworks for resting state are DefaultA and DefaultB by ATTRA, and VisCent and DorsAttnA by ATTRX ; although they do coincide with each other for tasks like VMN. This disparity is reflected in fig. 6 as edge and node differences.
Group, session, and region heterogeneity. Average variances of attributions are very different across tasks, especially those of ATTRX : VWM and DYN have much smaller attribution variances compared to other tasks. This can be caused by either task dynamics when certain tasks have more phase transitions and brain status changes, or/and group heterogeneity when individuals carry out
certain tasks more differently than the others. We investigate the cause by looking into 3 subjects; each of them has multiple scan sessions for every task.
We report the following findings: (1) Even only aggregating attributions over a single subject’s sessions, attribution variances of the other four tasks are still larger than VWM and DYN. The variance values are comparable to that of aggregating over many subjects. This means the large variances are not mainly due to group heterogeneity, rather some tasks having more states than others. (2) Apart from different task dynamics, there is still group heterogeneity. For tasks with more dynamics (high attribution variances), the group heterogeneity is also more obvious. We can see from fig. 7 that attributions for VMM are much more concentrated and universal across subjects than that of MOD. (3) Flexibility of different subnetworks varies: subnetworks that have small distribution IQR (Interquartile Range) of the same subject’s different sessions will also be more consistent across subjects. One example is subnetwork 18 during MOD task has both higher withinsubject IQR and larger across-subject differences compared to subnetwork 19. This indicates for a certain task, some subnetworks are more individual and flexible (may activate differently across t), while others are more universal and fixed. In summary, we can reveal both critical regions that a particular task must rely on, and regions that can characterize individual differences during tasks.

heading: 4 CONCLUSIONS
text: In this paper, we propose ReBraiD, a high-performing and efficient graph neural network model that embeds both structural and dynamic functional brain signals for task classifications. To better capture latent dynamics, we propose input-dependent adjacency matrix learning and inner-cluster smoothing at multiple resolutions. Apart from quantitative results showing ReBraiD’s superiority in representing brain activities, we also leverage integrated gradients to attribute and interpret the importance of both spatial brain regions and temporal keyframes, as well as presenting heterogeneities among subnetworks, tasks, and individuals. These findings can potentially reveal new neural basis or biomarkers of tasks and brain disorders when combined with behavioral metrics, and enable more fine-grained temporal analysis around keyframes when combined with other imaging techniques.

heading: A APPENDIX

heading: A.1 MODELS

heading: A.1.1 DYNAMIC FUNCTIONAL CONNECTIVITIES.
text: Fig. 8 shows functional connectivities (FCs) among N brain regions, where each FC ∈ RN×N . The value at FCij is calculated as the Pearson correlation coefficient between signals of brain region i and region j. The figure shows 6 FCs calculated from 6 consecutive sliding windows within a same fMRI session, with signal window length being 30 and sliding stride being 30. From the figure, we can clearly tell that FCs are highly dynamic.
A.1.2 INNER-CLUSTER SMOOTHING TOY EXAMPLE.
Here we show a toy example demonstrating the inner-cluster smoothing module described in eqs. (3) and (4). Note that we will only show one time slice, and the same operation is done along every t: on a particular t, we have Z ∈ RN×d, S ∈ RN×c. We will use N = 3, c = 2 and node values a, b, c ∈ Rd for this toy example. In addition, this example is just to illustrate the concept behind the smoothing operation, and Softmax along the axis 1 is simplified as row normalization for a clearer presentation.
 1 01 ⇒ H̃ = S>Z = ( a+ 12b1
In this example, 1st and 2nd nodes are assigned to the first cluster, and 2nd and 3rd node are assigned to the second cluster. The final Hnew after our smoothing module will mingle the first two nodes’ values, and the last two nodes’ values (based on assignment weights) while keeping their original node number unchanged.

heading: A.1.3 REGULARIZATION TERMS FOR THE SOFT-ASSIGNMENT.
text: For each soft assignment matrix S ∈ RN×c×t in eq. (3), we test three regularization terms:
• Similar to DIFFPOOL, for ensuring a more clearly defined node assignment, namely each node is only assigned to few clusters (the closer to one the better), we can minimize the entropy of single node assignments: LE1 = 1 c ∑c i=1H(Si), where Si is along c.
• To ensure a representation separation among nodes, meaning the assignment shouldn’t assign all the nodes a same way, we maximize the entropy of node assignment patterns across all nodes: LE2 = − 1c ∑c i=1H( ∑n j=1 Sij), where j is along n and i is along c.
• To make assignment along temporal axis smoother, we penalize assignment variances within a small window [t̂, t̂+τ ]: LT = 1t−τ ∑t−τ t̂=0
σ(S[t̂,t̂+τ ]), where σ represents standard deviation.
Together with cross entropy classification loss LCE , the final loss function of the model will be: Lreg = α1LCE + α2LE1 + α3LE2 + α4LT , ∑ i αi = 1 (6)

heading: A.2 EXPERIMENTS

heading: A.2.1 TASK DESCRIPTIONS.
text: The following are task descriptions of CRASH (Cognitive Resilience and Sleep History) dataset:
Resting state: The subject simply lays in the scanner awake, with eyes open for 5 minutes.
Visual working memory task (VWM): The subject is presented with a pattern of colored squares on a computer screen for a very brief period (100ms). After 1000ms, they are presented with a single square and must determine if it is the same or different color as the previously presented square at that location. Responses are made with a button press (Luck & Vogel (1997)).
Dynamic Attention Task (DYN): Two streams of orientation gratings are presented to the left and right of fixation. Subjects monitor specified stream for a target (about 2 degree shift in orientation, clockwise or counter clockwise) that indicates whether the subject should continue to monitor the current stream (hold) or monitor the other stream (shift) and respond with a button press (Yantis et al. (2002)).
Dot Probe Task (Faces) (DOT): On each trial, two faces are presented, one neutral and the other happy or angry for 500ms. Then, either of two simple symbols is presented at the position of either of the faces. The subject must make a forced choice discrimination of the symbol. Reaction time differences as a function of the valance for the preceding facial expression are calculated. There is increased variability of the bias with PTSD and fatigue (Sipos et al. (2014)).
Math task (MOD): Subjects perform a modular math computation every 8 seconds and respond with a yes or no button press. The object of modular arithmetic is to judge the validity of problems such as 51=19(mod 4). One way to solve it is to subtract the middle number from the first number (i.e., 51–19) and then divide this difference is by the last number (32/4). If the dividend is a whole number, the answer is “true.” Otherwise the answer is false (Mattarella-Micke et al. (2011)).
Psychomotor vigilance task (PVT): The subject monitors the outline of a red circle on a computer screen for 10 minutes, and whenever a counter clockwise red sweep begins, they press a button as fast as possible. Subjects are provided with response time feedback. The experimenter records response latencies (Loh et al. (2004)).

heading: A.2.2 DATA DETAILS
text: After discarding scan sessions with abnormal spikes that may be caused by head movements, the valid session details for different tasks are listed in table 2.

heading: A.2.3 ABLATION STUDIES
text: Numerical values of fig. 3a are reported in table 3. Training time ranges from 51 seconds / epoch for length-8 inputs to 298 seconds / epoch for length-256 inputs. Although the model is trained for 60 epochs in all experiments, it converges to a relatively stable loss level within 20 epochs.
(I) Latent adaptive adjacency matrix Aadp. As we mentioned in section 3.1, latent Aadp can complement the task- and temporal-fixed A. We will now show that the learned Ai adp is sparse for each sample, has evident task-based patterns, and differs from what Ai can provide: fig. 10 shows visualizations of latent Aadp, which we can tell is quite sparse as in fig. 10a: each input only gets few important columns (information providing nodes in GNN), and they vary from one sample to another, indicating Aadp’s ability to adapt to changing inputs within a same task. However, when we look into inputs (not shuffled) generated by consecutive sliding windows from a same scan session as in fig. 10b, we can see the latent structures appear to be in a smooth transition. In addition, when averaged across many samples for each task, undeniable task-based patterns emerge, as in figs. 10c
to 10e. But the task-average patterns are also different from what we saw from AttrA in fig. 5, suggesting Aadp is embedding dynamics that are not captured by A. One interesting phenomenon in fig. 10e is that across all tasks, LimbicB OFC, Default B, and LimbicA TempPole are always among the most important subnetworks appeared in the latent Aadp. Further exploration is needed for explaining the case.
(II) Number of GNN layers. The total number of temporal layers depends on the input signal length since each strided TCN layer reduces the temporal length by a factor of 2: if the input length is 2i, there need to be i temporal layers. But is alternating every TCN with GNN the best strategy, or do we only need to follow one GNN after a few TCNs? We study this question with different input lengths.
Model weighted F1 are plotted in fig. 3b for all possible GNN to total TCN ratios (e.g. length-256 inputs requires 8 TCN layers. The possible ratios are 18 , 1 4 , 1 2 , 1 since we can insert one GNN per 8, 4, 2, 1 TCN layers). The figure shows alternating every layer rarely yields the highest performance and the best ratio lines around one GNN per two TCN layers for our dataset. We repeat the experiment for K = 1, 3 (in eq. (2)) to rule out the possibility that this result is related to how many neighbors one GNN layer can reach; we find they have roughly the same pattern as the K = 2 case. Our hypothesis is that lower GNN to TCN ratio does not capture enough spatial context, while higher ones might be overfitting. We leave exploring the relationship between this ratio and nodes number N to a future study.
The best GNN to TCN ratio also depends on whether model incorporates latent adjacency matrices or not: without Aadp, length-128 signals achieves its relative best (among all ratios) when having one GNN per two TCNs, but it only needs one GNN per three TCNs if using Aadp. This shows learning latent structures Aadp not only improves overall model accuracy but can also reduce model parameters, thus complexity, in achieving the relative best results.
(III) Effects of soft-assignment cluster numbers. During our experiments, we find as long as the smoothing module is used, the final performance will be close to each other, only the convergence rates are different. Fig. 11b shows how validation loss converges with different c settings and when there is no smoothing module used. From it, we can see halving numbers (100-50-25-12) is most helpful and we use it for our other experiments; decreasing numbers (160-120-80-40) or all larger numbers (all 100) works better than increasing numbers (12-25-50-100) or all smaller numbers (all 12). Using the inner-cluster smoothing module, all cluster number settings converges to around 0.23 at their smallest when trained for 60 epochs and have test weighted F1 from 89.47 (model with 12- 25-50-100) to 90.85 (model with 100-50-25-12). On the contrary, if no smoothing module is used, the model overfits easily and the validation loss can only reach about 0.4 before going up (with the best set of learning rate and weight decay parameters found with grid search). It is understandable that the model is prone to overfitting given the complexity of GNN and the relatively small dataset size. Our added inner-cluster smoothing module seem to effectively countering the effect and brings the loss down further and stabler.

heading: A.2.4 MODEL COMPARISONS.
text: We plot confusion matrices of our model, model from ablation study setting (viii), and the best performing baseline in fig. 12. Misclassification pairs clustered as the first three tasks (resting, VWM, DYN) and the latter three (DOT, MOD, PVT). Shown confusion matrices are from models trained on 256-frame inputs. We note that these misclassification pairs may be different for models trained on other input lengths (like 128-frame, etc.).
A.2.5 IG AttrA AND AttrX
See fig. 13 for the visualization of important ROIs based on AttrA, and fig. 14 for task-averaged AttrA under real and random SC settings. Many discriminatory regions obtained from AttrA complies with the literature: For resting state: the top attributed ROIs belong to the default mode network, which is regarded salient during the resting state (Raichle, 2015). For VWM: the dominant attributions are from visual regions and posterior parietal regions, which complies with Todd & Marois (2004). For DYN: attributions from our model suggest regions along cingulate gyrus (defaultA-
SalValAttnB-ContA-ContC-defaultC), as well as peripheral visual and somatomotor regions. Literature suggests anterior cingulate cortex (ACC) to be active (Kim et al., 2016) and posterior cingulate cortex (PCC) to be inactive (Leech & Sharp, 2014) during visual attention tasks. This means both regions provide discriminative information about classifying DYN states, which is what our attribution method is voting for. For DOT: important ROIs from our analysis are located in control networks, in particular both ACC and PCC, as well as in the peripheral visual system. In the literature, dorsal and rostral regions of the ACC are proved to be involved with dot-probe performance (Carlson et al., 2012; 2013).
For MOD: our important ROIs are mostly in temporal-parietal regions and default mode network (anatomically fronto-parietal), and literature suggests similar regions: parietal (Grabner et al., 2011) and prefrontal (Friedrich & Friederici, 2013). For PVT: our top attributed ROIs belong to control networks, attention networks, and somatomotor regions. Similar as stated in Drummond et al. (2005), where both attention and motor systems are considered important.
To view brain regions in the 17-networks setting instead of 200-ROI parcellations, table 5 has the subnetwork rankings based on column-average AttrA and table 6 has the subnetwork rankings based on temporal-averaged AttrX . For a visualization of the 17-network parcellation, please refer to fig. 15.
See fig. 16 for the complete attribution distributions for every task based on temporal-averaged AttrX . Corresponding brain regions of a certain number are listed in table 4. For boxplots showing 17 regions, we combine LH and RH for their common network.

heading: A.2.6 ATTRIBUTION REPRODUCIBILITY
text: In order to be used for downstream tasks, the extracted important regions should be reproducible across different initializations. Since the overall problem is non-convex, we test the reproducibility empirically with two models trained on two data splits, both initialized randomly and taking in length 256 inputs.

heading: A.2.7 SIMULATION STUDY
text: To validate interpretation results, we perform simulation studies with know ground truth. For generating graph signals, we first define graph structures. All graphs are generated with stochastic block model (SBM) using a same community structure (200 nodes, 10 communities), but each graph has its own adjacency matrix. This mimics brain structures in that samples share similar community structures but have different structural connectivities. Fig 18a shows a typical adjacency matrix of a synthetic graph. All adjacency matrices are binary. Time-series on each node are then generated with codes adapted from this repo 4. In particular, the value at each time step of each node is a small temporal Gaussian random noise plus signals from neighbors (a small Gaussian spatial noise is added to the adjacency matrix).
Simulation (I) We create two classes for this simulation: in class one, only the first three communities (nodes 1 - 60) generate small temporal noises and other nodes are only affected by neighbors; in class two, only the last three communities (nodes 141 - 200) generate small temporal noises and other nodes are only affected by neighbors. We visualize task aggregated AttrX and Aadp and in figs. 18b and 18c. The signal importance differences are relatively well reflected in AttrX . For the generated series, signals are more important in node 1 - 60 for class 1 and 141 - 200 for class 2: Aadp catches this up and help propagating signals in these regions better. We notice that AttrA are mostly random and no obvious patterns are shown. This also reflects the graph signal generation: when aggregating information from neighbors, all connected edges are weighted the same (binary), thus the connections not really affect generated signals. To see the opposite, we perform another study below.
Simulation (II) We again create two classes for the simulation: in class one, connections from nodes 61 - 100 are strengthened; in class two, connections from nodes 101 - 140 are strengthened. The weights of strengthened edges are increased from 1 to 5 during signal generation, but model still takes in binary adjacency matrices as inputs (processed as mentioned in section 2.1 before feeding to the model). We visualize task aggregated Aadp and AttrA in figs. 18d and 18e. We can see this time the connection differences are reflected in AttrA. Signals in node 61 - 100 for class 1 or 101 -140 for class 2 are less important because they can be sent out by stronger connections: this results in smaller values for corresponding columns in Aadp. Combined with previous simulation, regions that are strong signal senders and connections from them are weak or not reflected in graph adjacency matrices tend to have higher Aadp values. In other words, Aadp complements both signals and connections to encode latent dynamics, while attributions obtained from IG are better at interpreting the modalities separately.
4https://github.com/alelab-upenn/graph-neural-networks

Review summaries:
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.

This paper presents a graph neural network architecture for analyzing dynamic brain imaging data. The proposed approach uses both structural connectivity (via DWI) and functional connectivity (via fMRI) to setup the graph edge information, while the fMRI signals are used to define node features. The proposed network alternates between analyzing temporal information via gated temporal convolutional network layers and spatial information via GNN layers. A multi-resolution soft cluster smoothing is applied as the pooling operation in the GNN. Attribution of structural and functional imaging inputs is also explored using integrated gradients interpretation method. The proposed method is tested on a multi-task fMRI dataset and demonstrated better performance in both ablation studies and comparison to state-of-the-art GNN methods.

This paper proposes a method of graph neural network to jointly models dynamic functional signals and structural connectivities, to learn a good deep representation of brain dynamics. In this article, the author proposes some improved methods for graph neural networks, such as learning sample-level latent graph structures, temporal convolutional network and multi-resolution inner cluster smoothing methods to better learn insightful interpretations of brain dynamics. In the experiment, the author verified the effectiveness of the method in this paper on the fMRI signal data.