Graph Embedding Using Infomax for ASD Classification and Brain Functional Difference Detection

Abstract

Significant progress has been made using fMRI to characterize the brain changes that occur in ASD, a complex neuro-developmental disorder. However, due to the high dimensionality and low signal-to-noise ratio of fMRI, embedding informative and robust brain regional fMRI representations for both graph-level classification and region-level functional difference detection tasks between ASD and healthy control (HC) groups is difficult. Here, we model the whole brain fMRI as a graph, which preserves geometrical and temporal information and use a Graph Neural Network (GNN) to learn from the graph-structured fMRI data. We investigate the potential of including mutual information (MI) loss (Infomax), which is an unsupervised term encouraging large MI of each nodal representation and its corresponding graph-level summarized representation to learn a better graph embedding. Specifically, this work developed a pipeline including a GNN encoder, a classifier and a discriminator, which forces the encoded nodal representations to both benefit classification and reveal the common nodal patterns in a graph. We simultaneously optimize graph-level classification loss and Infomax. We demonstrated that Infomax graph embedding improves classification performance as a regularization term. Furthermore, we found separable nodal representations of ASD and HC groups in prefrontal cortex, cingulate cortex, visual regions, and other social, emotional and execution related brain regions. In contrast with GNN with classification loss only, the proposed pipeline can facilitate training more robust ASD classification models. Moreover, the separable nodal representations can detect the functional differences between the two groups and contribute to revealing new ASD biomarkers.

1. INTRODUCTION

Autism spectrum disorder (ASD) affects the structure and function of the brain. Functional magnetic resonance imaging (fMRI) produces 4D spatial-temporal data describing functional activation but with very low signal-noise ratio (SNR). It can be used to characterize neural pathways and brain changes that occur in ASD. However, due to high dimension and low SNR, it is difficult to analyze fMRI. Here, we address the problem of embedding good fMRI representations for identifying ASD and detecting brain functional differences between ASD and healthy control (HC). To utilize the spatial-temporal information of fMRI, we represent the whole brain fMRI as a graph, where each brain region (ROI) is a node, the underlying connection can be calculated by fMRI correlation matrix and node features can be predetermined, hence preserving both geometrical and temporal information. The Graph Neural Network (GNN), a deep learning architecture to analyze graph structured data, has been used in ASD classification.¹ In addition to improving ASD classification, one core objective of our work is to discover useful representations to detect brain regional differences between ASD vs HC. The simple idea explored here is to train a representation-learning function related to the end-goal task, which maximizes the mutual information (MI) between nodal representation and graph-level representation and minimizes the loss of the end-goal task. MI is notoriously difficult to compute, particularly in continuous and high dimensional settings. Fortunately, the recently proposed MINE² enables effective computation of MI between high dimensional input/output pairs of a deep neural network, by training a statistics network as a classifier of samples coming from the joint distribution of two random variables and the product of their marginals. During training of a GNN, we simultaneously optimize the classification loss and Infomax loss,³ which maximizes the MI between local/global representation. In this way, we tune the suitability of learned representations for classification and detecting group-level regional functional differences. Results show the improvement of the classification task and reveal the functional differences between ASD and HC from the separable embedded brain regions encoded by the GNN.

Our contributions are summarized as follows:

1. We formulate ASD brain abnormal functional discovery as a node embedding problem using GNN.

2. To the best of our knowledge, we investigate deep graph infomax (DGI) for fMRI image analysis for the first time. The learnt representations benefit both summarizing nodes of interest, and downstream graph-level classification task.

3. We discover group-level brain ROI difference between ASD v.s. HC group and provide potential biomarkers for ASD diagnosis and treatment.

Paper structure:

In Section 2, we summarize transitional methods to analyze fMRI as a graph and GNNs techniques. In Section 3, we introduce the methods used for our study. Specifically, in Section 3.1, we propose how to covert fMRI images to graphs; in Section 3.2, we introduce the details of network architecture; and in Section 3.3, we develop the loss function. The experiments, results, and evaluation methods are presented in Section 4. We conclude the paper in Section 5.

2. RELATED WORK

2.1. Analyze fMRI As A Graph

Usually, brain can be parcellated by atlas resulting a particular number of ROIs, saying R. Then mean fMRI time series of each ROI is extracted then calculate its correlation to the mean fMRI time series of other ROIs, resulting in a R × R correlation matrices for each subject. With the connectiviy matrices at hand, we hope to reduce the dimensionality of the signal to a few hundreds components, which can contain some cognitive information. The functional network of the brain can be modelled as a graph in which each node is a brain region and the edges represent the direction and strength of the connection between those regions. There are several ways to measure the edge weights. These are generally divided into two groups. The first group of methods focus on using finding causality between different brain regions and thus deriving directed edges. These methods include Granger causality⁴ and dynamic casual modeling.⁵ The second group of methods is used to obtain undirected edges by using simpler statistical association measures to measure the intra-node connection strength. These are discussed in more detail in the “connectivity measures” section. Modelling the brain as a graph allows us to use graph theoretical metrics to study the functional network. Above all, to extract the whole-brain functional connectivity network of each subject, each ROI is seen as a network node and a measure of connectivity is used to connect these nodes.⁶ This connectivity measure a_ij must be able to quantify the relationship between the time series of ROI i and j. Correlation and mutual information metrics have been extensively used for this purpose.⁷

2.2. Graph Neural Network

Different from the traditional graph learning methods mentioned above, GNN is the state-of-the-art deep learning methods for graph-structured data analysis. GNNs combine node features, connection patterns, and graph struc-ture by using a neural network to embed node information and passing through edges in the graph. GNN methods aim to generalize the traditional convolutional neural networks (CNN) used in image classification to graph structures. For example, Graph Convolutional Neural Network (GCN)⁸ proposes spectral graph convolutions and provide approximated fast computation.⁹ proposes another variant for general graph classification from spatial space. Due to its convincing performance and high interpretability, GNN has been a widely applied graph analysis method recently, especially on fMRI analysis.^1,10,11

3. METHODOLOGY

3.1. Data Definition and Notations

On a labeled graph set $\mathcal{C}=\{(G_1,y_1),(G_2,y_2),\dots \}$, the general graph classification problem is to learn a classifier that maps G_i to its label y_i. In practice, the G_i is usually given as a triple G = (V, A, X) where V = {v₁,… v_N} is the set of N nodes, A = {a_ij}_N×N is the set of edges with a_ij denoting the edge weight, and $\mathbf{X}=\{{\overrightarrow{x}}_1,\dots ,{\overrightarrow{x}}_N{\}}_{N\times D}$ is the set of node features.

Suppose each brain is parcellated into N ROIs based on its T1 structural MRI. We define an undirected graph on the brain regions G = (V, A). Nodes V are associated with $\mathbf{X}\in {\mathbb{R}}^{N\times D}$ and $\mathbf{A}=[a_{ij}]\in {\mathbb{R}}^{N\times N}$, and D is the attribute dimensions of nodes feature. For node attributes, we concatenate handcrafted features: degree of connectivity, General Linear Model (GLM) coefficients, mean, and standard deviation of task-fMRI, and ROI center coordinates. A is calculated by the correlation of the mean fMRI time series in each ROI.

3.2. Graph Embedding

In addition to improving ASD classification, one core objective of our work is to discover useful representations to detect brain regional differences between ASD vs HC. The overview and each component of our architecture to achieve our goal are discussed in his section.

3.2.1. Overview of the pipeline

We present the pipeline in Figure 1. It contains three components: node encoder, graph classifier and data discriminator. Graph convolutional kernel (Section 3.2.2) will encode the input graph to a feature map $\mathbf{H}=({\overrightarrow{h}}_1,{\overrightarrow{h}}_2,\dots ,{\overrightarrow{h}}_N{)}^T\in {\mathbb{R}}^{N\times F}$, that reflects useful structure locally. Next, we summarize the node representation into a global feature $\overrightarrow{s}$ by pooling and reading out (Section 3.2.3). Given a G, we will generate a negative graph G′, whose embedded node representation is H′. The corresponding positive pair $({\overrightarrow{h}}_i,\overrightarrow{s})$ and negative pair $({\overrightarrow{h}}_i^{{}_{’}},\overrightarrow{s})$ will be encouraged to be separated by a discriminator $\mathcal{D}$ (Section 3.2.4).

Figure 1:

The flowchart of our proposed ASD classification and graph embedding architecture. The top row of the flowchart is a Graph Neural Network architecture to classify ASD and HC. The bottom row is a graph infomax pipeline to encourage better graph embedding. Here, (a) and (b) are positive samples; (c) and (d) are negative samples. (a)(c) (or (b)(d)) is a paired graph. The inputs of discriminator D is the summary vector generated from positive samples, paired with node embedded representation (H or H′). $({\overrightarrow{h}}_i,\overrightarrow{s})$ pair will have True (T) output from D, whereas $({\overrightarrow{h}}_i^{{}_{’}},\overrightarrow{s})$ will have False (F) output. The encoder, classifier and discriminator are trained simultaneously.

3.2.2. Encoder: Graph Convolutional Layer

Generally speaking, GNNs inductively learn a node rep-resentation by recursively aggregating and transforming the featurevectors of its neighboring nodes. Our encoder $\mathcal{E}$ node embedding network is a L-layer supervised GraphSAGE architecture,⁹ which learns the embedding function $\mathcal{E}:\mathbf{X}\in {\mathbb{R}}^{N\times D}\to \mathbf{H}\in {\mathbb{R}}^{N\times F}$ mapping input nodes X to output H. The embedding function is based on the mean-pooling (MP) propagation rule $MP(\mathbf{X},\mathbf{A})={\widehat{\mathbf{D}}}^{-1}\widehat{\mathbf{A}}\mathbf{X}\Theta$ as used in Hamilton et al.,⁹ where $\widehat{\mathbf{A}}=\mathbf{A}+{\mathbf{I}}_N$ is the adjacency matrix with inserted self-loops and $\widehat{\mathbf{D}}$ is its corresponding degree diagonal matrix with ${\widehat{\mathbf{D}}}_{\mathbf{ii}}={\sum }_j{\widehat{\mathbf{A}}}_{ij}$. Our encoder can be written as:

$${\mathbf{H}}_1=\sigma (M{P}_1(\mathbf{X},\mathbf{A})){\mathbf{H}}_L=\mathcal{E}(\mathbf{X},\mathbf{A})=\sigma (M{P}_L({\mathbf{H}}_{L-1}+\cdots +{\mathbf{H}}_1+\mathbf{XWA}))$$ (1)

where W is a learnable projection matrix and σ is sigmoid function.

3.2.3. Classifier: Pooling and Readout Layer

To aggregate the information of each node for the graph level classification, we use Dense hierarchical pooling (DHP¹²) to cluster nodes together. After each DHP, the number of nodes in the graph decreases. At the last level L, the pooling layer is performed by a hltering matrix $F\in {\mathbb{R}}^{N\times Q}$.

$${\mathbf{H}}_p={\mathsf{F}}^T{\mathbf{H}}_L{\mathbf{A}}_p={\mathsf{F}}^T\mathbf{A}\mathsf{F}\overrightarrow{r}=\mathcal{R}({\mathbf{H}}_p,{\mathbf{A}}_p)=\frac{1}{Q}\sum _i^Q{\overrightarrow{h}}_i$$ (2)

produces pooled nodes ${\mathbf{H}}_p\in {\mathbb{R}}^{Q\times F}$ and adjacency matrix ${\mathbf{A}}_p\in {\mathbb{R}}^{Q\times Q}$, which generate readout vector $\overrightarrow{r}$ The final number of nodes Q is predefined. F was learned by another GraphSAGE convolutional layer optimized by a regularization loss L_reg = ∥A_p,FF^T∥_F, where ∥ · ∥_F denotes the Frobenius norm. Readout vector $\overrightarrow{r}$ will be submitted to a MLP for obtaining final classification outputs p, the probability of being an ASD subject.

3.2.4. Discriminator: Encouraging Good Representation

Following the intuition in Deep Graph Infomax,³ the good representation may not benefit from encoding counter information. In order to obtain a representation more suitable for classification, we maximize the average MI between the high-level representation and local aggregated embedding of each node, which favours encoding aspects of the data that are shared across the nodes and reduces noisy encoding.¹³ The graph-level summary vector can be $\overrightarrow{s}=\sigma (\overrightarrow{r})$ as the input of discriminator, here σ is the logistic sigmoid nonlinearity. A discriminator $\mathcal{D}({\overrightarrow{h}}_i,\overrightarrow{s}):{\mathbb{R}}^F\times {\mathbb{R}}^F\to \mathbb{R}$ is used as a proxy for maximizing the MI representing the probability scores assigned to the local-global pairs. We randomly sample an instance from the opposite class as the negative sample (X′, A′). The discriminator scores summary-node representation pairs by applying a simple bi-linear scoring function³

$$\mathcal{D}({\overrightarrow{h}}_i,\overrightarrow{s})=\sigma ({\overrightarrow{h}}_i^T\mathbf{M}\overrightarrow{s})$$ (3)

where M is a learnable scoring matrix and σ is the logistic sigmoid nonlinearity, used to convert scores into probabilities of (${\overrightarrow{h}}_i,\overrightarrow{s}$) being positive.

3.3. Loss function

In order to learn useful, predictive representations, the Infomax loss function L₂ encourages nodes of the same graph to have similar representations, while enforcing that the representations of disparate nodes are highly distinct. In order to insure the performance of downstream classification, we use binary cross-entropy as the classification loss L₁. Therefore, the loss function of our model is written as:

$$\mathcal{L}=\underset{L_1}{\underbrace{-\frac{1}{N}\sum _{i=1}^N(y_i\log (p_i)+(1-y_i)\log (1-p_i))}}+\underset{L_2}{\underbrace{\frac{1}{2N}(\sum _{i=1}^N{\mathbb{E}}_{(\mathbf{X},\mathbf{A})}[\log \mathcal{D}({\overrightarrow{h}}_i,\overrightarrow{s})]+{\mathbb{E}}_{({\mathbf{X}}^{\prime },{\mathbf{A}}^{\prime })}[1-\log \mathcal{D}({\overrightarrow{h}}_i^{\prime },\overrightarrow{s})])}}$$ (4)

Therefore, we can achieve graph classification and different node representations detection goals by end-to-end training manner.

4. EXPERIMENT AND RESULTS

4.1. Data Acquisition and Preprocessing

For the fMRI scans, subjects viewed point light animations of coherent and scrambled biological motion in a block design.¹⁴ We tested our method on a group of 75 ASD children and 43 age and IQ-matched healthy controls collected at Yale Child Study Center.¹ The experimental paradigm features coherent and scrambled point-light animations created from motion capture data. The coherent biological motion depicts an adult male actor performing movements relevant to early childhood experiences.¹⁵ The random motion combine the trajectories of 16 randomly selected points from the coherent displays. Six biological motion clips and six scrambled motion clips were presented in an alternating-block design (24s per block). The fMRI data was preprocessed following the pipeline in Yang et al.¹⁶ The fMRI data was preprocessed using Freesurfer as follows: 1) motion correction using MCFLIRT, 2) interleaved slice timing correction, 3) BET brain extraction, 4) spatial smoothing (FWHM=5mm), and 5) high-pass temporal filtering. The functional and anatomical data were registered to the MNI152 standard brain atlas.¹⁷ The graph data was augmented as described in our previous work,¹ resulting in 750 ASD graphs and 860 HC graphs. We split the data into 5 folds based on subjects. Four folds were used as training data and the left out fold was used for testing. Each node attribute ${\overrightarrow{x}}_i\in {\mathbb{R}}^{10}$, which include node degree, mean fMRI, std fMRI, ROI center coordinates and GLM parameters. Specifically, the GLM parameters of the “biopoint task” are: β₁, coefficient of biological motion matrix; β₃, coefficient of scrambled motion matrix; β₂ and β₄, coefficients of the previous two matrices’ derivatives.

4.2. Experiment and Results

We tested classifier performance on the Destrieux atlas¹⁸ (148 ROIs) using the proposed GNN with L₁ and $\mathcal{L}$ separately, to examine the advantage of including graph infomax loss L₂. In our GNN setting, D = 10 and pooling ratios r = 0.5. We used the Adam optimizer with initial learning 0.001, then decreased it by a factor of 2 every 20 epochs. We trained the network 100 epochs for all of the splits and measured the instance classification by F-score (Table 1). We changed the architectures by tuning either two graph convolutional layers with kernel size (F, F) or one graph convolutional layer with kernel size (F). F was tested at 8 and 16. The regularization parameters are adjusted correspondingly to get the best performance.

Table 1:

Performance of different loss functions and GNN architectures (mean± std)

Loss + (conv-layer)	L1(16,16)	L1(8,8)	L1(16)	$\mathcal{L}$(16,16)	$\mathcal{L}$(8,8)	$\mathcal{L}$(16)
F-score	0.57±0.11	0.70±0.06	0.63±0.01	0.68±0.08	0.69±0.05	0.66±0.03

For notation convenience, we use (·) model and $\mathcal{L}∕L_1$ model to represent the model of certain GNN architecture and corresponding training loss. Under model (8, 8), we could not find obvious advantage of using $\mathcal{L}$. However, if we increase the encoders’ complexity to (16, 16), the L₁ model became easily overfitted while $\mathcal{L}$ model kept similar performance. This may indicate L₂ can perform as regularization and restrain embedding from data noise. In (16) case, the L₁ model was underhtted, while the $\mathcal{L}$ model performed slightly better. It’s probably because L₂ encourages encoding common nodal signals over the graph hence ignoring data noise or just because $\mathcal{L}$ model had more trainable parameters.

After training, we extracted the nodal embedded vectors after the last Graph Convolutional Layer and used t-SNE¹⁹ to visualize the node presentations in 2D space shown in Fig. 2. Only with $\mathcal{L}$ did we find linearly separable nodal representations of ASD and HC for certain regions. We visually examined all the nodal representation embeddings by L₁ and verified they cannot be linearly separated into the two groups. We marked the regions whose Silhouette score²⁰ was greater than 0.1 (resulting in 31 regions using $\mathcal{L}$(8, 8), shown in Fig. 3a) as the brain ROIs with functional difference between ASD and HC. We compared the results with GLM z-stats analysis using FSL²¹ (shown in Fig. 3b ). Our proposed method marked obvious prefrontal cortex, while GLM method did not highlight those regions. Both our method and GLM analysis highlighted cingulate cortex. These regions were indicated as ASD biomarkers in Yang et al.¹⁶ and Kaiser et al.¹⁴ Also, we used Neurosynth²² to decode the functional keywords associated with separable regions found by our methods, as shown in Fig. 4. The decoded functional keywords of our detected regions showed evidence that these regions might have social-, mental-, visual-related and default mind functional differences between ASD and HC group. Potentially, our proposed method can be used as a tool to identify new brain biomarkers for better understanding the underlying roots of ASD.

Figure 2:

Embedded representations of 24 out of 148 brain regions visualized by t-SNE. HC is colored in green and ASD is colored in red. The regions representations circled by blue box are linear separable.

Figure 3:

Highlight functional differences between ASD and HC.

Figure 4:

The correlation between the functional keyword and the regions in Figure 3a decoded by Neurosynth.

5. CONCLUSION

We applied GNN to identify ASD and designed a loss function to encourage better node representation and detect separable brain regions of ASD and HC. By incorporating mutual information of local and global representations, the proposed loss function improved classification performance in certain cases. The added L₂ Infomax loss potentially regularizes the embedding of noisy fMRI and increases model robustness. By examining the embedded node representations, we found that ASD and HC had separable representations in regions related to default mode, social function, emotion regulation and visual function, etc. The finding is consistent with prior literature^{1, 14} and our approach could potentially discover new functional differences between ASD and HC. Overall, the proposed method provides an efficient and objective way of embedding ASD and HC brain graphs.