PRE-TRAIN GRAPH NEURAL NETWORKS FOR BRAIN NETWORK ANALYSIS

Abstract

Human brains, controlling behaviors and cognition, are at the center of complex neurobiological systems. Recent studies in neuroscience and neuroimaging analysis have reached a consensus that interactions among brain regions of interest (ROIs) are driving factors for neural development and disorders. Graph neural networks (GNNs) as a powerful tool for analyzing graph-structured data are naturally applied to the analysis of brain networks. However, training of deep learning models including GNNs often requires a significant amount of labeled data. Due to the complicated data acquisition process and restrictions on data sharing, brain network datasets are still small compared to other types of graphs (e.g., social networks, molecules, proteins). Moreover, real clinical tasks (e.g., mental disorder analysis) are often conducted on local datasets with even smaller scales and larger noises. To this end, we propose to leverage pre-training to capture the intrinsic brain network structures regardless of specific clinical outcomes, for obtaining GNN models that are easily adaptable to downstream tasks. Specifically, (1) we design brain-network-oriented unsupervised pre-training techniques to utilize large-scale brain imaging studies without highly relevant task labels; (2) we develop a data-driven parcellation atlas mapping pipeline to facilitate effective knowledge transfer across studies with different ROI systems. The proposed framework is validated with various GNN models, with extensive empirical results demonstrating consistent improvement in performance as well as robustness.

1. INTRODUCTION

In recent years, the analysis of brain networks has attracted considerable interest in neuroscience studies. Brain networks are essentially graphs, where anatomical regions of interest (ROIs) given a parcellation atlas are formed into nodes, and the connectivities among ROIs are formed into edges. Based on brain networks constructed from different modalities such as Diffusion Tensor Imaging (DTI) and functional Magnetic Resonance Imaging (fMRI), effective graph analysis plays a pivotal role in understanding the biological structures and functions of complex neural systems, which can be helpful in the early diagnosis of neurological disorders and facilitate neuroscience research Martensson et al. (2018); Yahata et al. (2016); Lindquist (2008); Smith (2012) . Deep learning has replenished the fields of computer science and beyond. Among various modern deep learning models, the emerging graph neural networks (GNNs) have demonstrated superior performance and even plausible interpretability on a variety of network datasets, including social networks, recommender systems, knowledge graphs, protein and gene networks, molecules and so forth Kipf & Welling ( 2017 One solution toward data scarcity is transfer learning which transfers models trained on the largescale brain network datasets onto small-scale local studies while retaining favorable performance. However, one limitation of transfer learning is its reliance on the availability of similar tasks as supervision during training in the source dataset. In reality, similar tasks in the smaller local studies may not always be available in the large-scale public studies. 2022). We explore GNN pre-training on brain networks without supervision and study its effectiveness in the prediction of specific clinical outcomes. However, unique challenges impede the direct application of existing GNN pre-training paradigms to brain networks. For example, brain networks within one study usually share the same node system, which is not properly exploited, whereas different studies often use different node systems, which hinders the transferability of pre-trained models. To fully unleash the power of GNNs for brain network data by collaborating across datasets, we propose to develop an unsupervised multi-dataset GNN pre-training framework for brain networks. Specifically, we adapt the popular data-efficient framework of MAML with carefully designed twolevel contrastive learning that works in concert with the functional neural system modules in brain networks. In addition, to address the problem of dataset misalignment, we propose a novel datadriven atlas mapping technique based on an auto-encoder that transforms the original features in multiple datasets into low-dimensional representations in a uniform embedding space. The transformed features are aligned via variance-based projection with locality-preserving, module-aware, and sparsity-oriented regularizations. In summary, our contributions are three-folded: • We identify the intrinsic data-scarcity issue in brain network learning and formulate our problem into an unsupervised pre-training objective. • We propose a novel brain-network-oriented two-level contrastive sampling strategy for multidataset GNN pre-training (Section 3.2). In addition, we implement a data-driven brain atlas mapping strategy with customized regularizations and variance-based sorting to facilitate cross-dataset model sharing (Section 3.3). • Extensive experiments are conducted to compare our proposed framework with various baselines (Section 4.1), and we further the analysis through in-depth studies on the contributions of each constituent component of our framework (Sections 4.2, 4.3, 4.4 ). 



); Hamilton et al. (2017); Schlichtkrull et al. (2018); Vashishth et al. (2020); Xu et al. (2019); Ying et al. (2018); He et al. (2020b); Zhang et al. (2020); Liu et al. (2022); Xiong et al. (2020), due to its powerful representations and efficient computations of complex graph structures towards specific downstream tasks. Such achievements on other types of networked data propel studies on GNNs for brain networks, especially models for graph-level classification/regression Ying et al. (2018); Xu et al. (2019); Errica et al. (2020) and important vertex/edge identification Ying et al. (2019); Luo et al. (2020); Vu & Thai (2020), towards tasks such as connectome-based disease prediction and multi-level neural pattern discovery. However, the training of powerful deep learning models including GNNs often requires significant amounts of labeled data Hu et al. (2020a); You et al. (2020); Zhu et al. (2021a). For brain network analysis, there are limited big imaging datasets from a few large-scale national neuroimaging studies such as the ABCD Casey et al. (2018), ADNI Hinrichs et al. (2009), and PPMI Aleksovski et al. (2018), but such datasets are still rather small compared with graph datasets in other domains (e.g., datasets with 41K to 452K graphs on OGBHu et al. (2020b)  and datasets with thousands to millions of graphs on NetRepo Rossi & Ahmed (2016)).

Pre-training has shown its effectiveness in the fields of computer vision He et al. (2020a); Chen et al. (2020b), natural language processing Devlin et al. (2019b); Radford et al. (2018), as well as graph mining Sun et al. (

He et al. (2020a) has inspired contrastive learning on graphs. For instance, GBT Bielak et al. (2022) designs a Barlow Twins Zbontar et al. (2021) loss function based on empirical cross-correlation of node embeddings learned from a pair of augmented graph views Zhao et al. (2021). Similarly, GraphCL You et al. (2020) contrasts among graph representations learned from different graphs. Generally, graph contrastive learning captures robust latent representations on well-attributed graphs with a sufficient number of graph samples. However, its accountability may be undermined in sample-scarce and attribute-lacking brain networks.

