Strategies for Pre-training Graph Neural Networks

1 INTRODUCTION . Transfer learning refers to the setting where a model , initially trained on some tasks , is re-purposed on different but related tasks . Deep transfer learning has been immensely successful in computer vision ( Donahue et al. , 2014 ; Girshick et al. , 2014 ; Zeiler & Fergus , 2014 ) and natural language processing ( Devlin et al. , 2019 ; Peters et al. , 2018 ; Mikolov et al. , 2013 ) . Despite being an effective approach to transfer learning , few studies have generalized pre-training to graph data . Pre-training has the potential to provide an attractive solution to the following two fundamental challenges with learning on graph datasets ( Pan & Yang , 2009 ; Hendrycks et al. , 2019 ) : First , task-specific labeled data can be extremely scarce . This problem is exacerbated in important graph datasets from scientific domains , such as chemistry and biology , where data labeling ( e.g. , biological experiments in a wet laboratory ) is resource- and time-intensive ( Zitnik et al. , 2018 ) . Second , graph data from real-world applications often contain out-of-distribution samples , meaning that graphs in the training set are structurally very different from graphs in the test set . Out-of-distribution prediction is common in real-world graph datasets , for example , when one wants to predict chemical properties of a brand-new , just synthesized molecule , which is different from all molecules synthesized so far , and thereby different from all molecules in the training set . However , pre-training on graph datasets remains a hard challenge . Several key studies ( Xu et al. , 2017 ; Ching et al. , 2018 ; Wang et al. , 2019 ) have shown that successful transfer learning is not only a ∗Equal contribution . Project website , data and code : http : //snap.stanford.edu/gnn-pretrain matter of increasing the number of labeled pre-training datasets that are from the same domain as the downstream task . Instead , it requires substantial domain expertise to carefully select examples and target labels that are correlated with the downstream task of interest . Otherwise , the transfer of knowledge from related pre-training tasks to a new downstream task can harm generalization , which is known as negative transfer ( Rosenstein et al. , 2005 ) and significantly limits the applicability and reliability of pre-trained models . Present work . Here , we focus on pre-training as an approach to transfer learning in Graph Neural Networks ( GNNs ) ( Kipf & Welling , 2017 ; Hamilton et al. , 2017a ; Ying et al. , 2018b ; Xu et al. , 2019 ; 2018 ) for graph-level property prediction . Our work presents two key contributions . ( 1 ) We conduct the first systematic large-scale investigation of strategies for pre-training GNNs . For that , we build two large new pre-training datasets , which we share with the community : a chemistry dataset with 2 million graphs and a biology dataset with 395K graphs . We also show that large domain-specific datasets are crucial to investigate pre-training and that existing downstream benchmark datasets are too small to evaluate models in a statistically reliable way . ( 2 ) We develop an effective pretraining strategy for GNNs and demonstrate its effectiveness and its ability for out-of-distribution generalization on hard transfer-learning problems . In our systematic study , we show that pre-training GNNs does not always help . Naïve pre-training strategies can lead to negative transfer on many downstream tasks . Strikingly , a seemingly strong pre-training strategy ( i.e. , graph-level multi-task supervised pre-training using a state-of-the-art graph neural network architecture for graph-level prediction tasks ) only gives marginal performance gains . Furthermore , this strategy even leads to negative transfer on many downstream tasks ( 2 out of 8 molecular datasets and 13 out of 40 protein prediction tasks ) . We develop an effective strategy for pre-training GNNs . The key idea is to use easily accessible node-level information and encourage GNNs to capture domain-specific knowledge about nodes and edges , in addition to graph-level knowledge . This helps the GNN to learn useful representations at both global and local levels ( Figure 1 ( a.iii ) ) , and is crucial to be able to generate graph-level representations ( which are obtained by pooling node representations ) that are robust and transferable to diverse downstream tasks ( Figure 1 ) . Our strategy is in contrast to naïve strategies that either leverage only at graph-level properties ( Figure 1 ( a.ii ) ) or node-level properties ( Figure 1 ( a.i ) ) . Empirically , our pre-training strategy used together with the most expressive GNN architecture , GIN ( Xu et al. , 2019 ) , yields state-of-the-art results on benchmark datasets and avoids negative transfer across downstream tasks we tested . It significantly improves generalization performance across downstream tasks , yielding up to 9.4 % higher average ROC-AUC than non-pre-trained GNNs , and up to 5.2 % higher average ROC-AUC compared to GNNs with the extensive graph-level multitask supervised pre-training . Furthermore , we find that the most expressive architecture , GIN , benefits more from pre-training compared to those with less expressive power ( e.g. , GCN ( Kipf & Welling , 2017 ) , GraphSAGE ( Hamilton et al. , 2017b ) and GAT ( Velickovic et al. , 2018 ) ) , and that pre-training GNNs leads to orders-of-magnitude faster training and convergence in the fine-tuning stage . 2 PRELIMINARIES OF GRAPH NEURAL NETWORKS . We first formalize supervised learning of graphs and provide an overview of GNNs ( Gilmer et al. , 2017 ) . Then , we briefly review methods for unsupervised graph representation learning . Supervised learning of graphs . Let G = ( V , E ) denote a graph with node attributes Xv for v ∈ V and edge attributes euv for ( u , v ) ∈ E. Given a set of graphs { G1 , . . . , GN } and their labels { y1 , . . . , yN } , the task of graph supervised learning is to learn a representation vector hG that helps predict the label of an entire graph G , yG = g ( hG ) . For example , in molecular property prediction , G is a molecular graph , where nodes represent atoms and edges represent chemical bonds , and the label to be predicted can be toxicity or enzyme binding . Graph Neural Networks ( GNNs ) . GNNs use the graph connectivity as well as node and edge features to learn a representation vector ( i.e. , embedding ) hv for every node v ∈ G and a vector hG for the entire graph G. Modern GNNs use a neighborhood aggregation approach , where representation of node v is iteratively updated by aggregating representations of v ’ s neighboring nodes and edges ( Gilmer et al. , 2017 ) . After k iterations of aggregation , v ’ s representation captures the structural information within its k-hop network neighborhood . Formally , the k-th layer of a GNN is : h ( k ) v = COMBINE ( k ) ( h ( k−1 ) v , AGGREGATE ( k ) ( { ( h ( k−1 ) v , h ( k−1 ) u , euv ) : u ∈ N ( v ) } ) ) , ( 2.1 ) where h ( k ) v is the representation of node v at the k-th iteration/layer , euv is the feature vector of edge between u and v , and N ( v ) is a set neighbors of v. We initialize h ( 0 ) v = Xv . Graph representation learning . To obtain the entire graph ’ s representation hG , the READOUT function pools node features from the final iteration K , hG = READOUT ( { h ( K ) v ∣∣ v ∈ G } ) . ( 2.2 ) READOUT is a permutation-invariant function , such as averaging or a more sophisticated graph-level pooling function ( Ying et al. , 2018b ; Zhang et al. , 2018 ) . 3 STRATEGIES FOR PRE-TRAINING GRAPH NEURAL NETWORKS . At the technical core of our pre-training strategy is the notion to pre-train a GNN both at the level of individual nodes as well as entire graphs . This notion encourages the GNN to capture domain-specific semantics at both levels , as illustrated in Figure 1 ( a.iii ) . This is in contrast to straightforward but limited pre-training strategies that either only use pre-training to predict properties of entire graphs ( Figure 1 ( a.ii ) ) or only use pre-training to predict properties of individual nodes ( Figure 1 ( a.i ) ) . In the following , we first describe our node-level pre-training approach ( Section 3.1 ) and then graph-level pre-training approach ( Section 3.2 ) . Finally , we describe the full pre-training strategy in Section 3.3 . 3.1 NODE-LEVEL PRE-TRAINING . For node-level pre-training of GNNs , our approach is to use easily-accessible unlabeled data to capture domain-specific knowledge/regularities in the graph . Here we propose two self-supervised methods , Context Prediction and Attribute Masking . 3.1.1 CONTEXT PREDICTION : EXPLOITING DISTRIBUTION OF GRAPH STRUCTURE . In Context Prediction , we use subgraphs to predict their surrounding graph structures . Our goal is to pre-train a GNN so that it maps nodes appearing in similar structural contexts to nearby embeddings ( Rubenstein & Goodenough , 1965 ; Mikolov et al. , 2013 ) . Neighborhood and context graphs . For every node v , we define v ’ s neighborhood and context graphs as follows . K-hop neighborhood of v contains all nodes and edges that are at most K-hops away from v in the graph . This is motivated by the fact that a K-layer GNN aggregates information across the K-th order neighborhood of v , and thus node embedding h ( K ) v depends on nodes that are at most K-hops away from v. We define context graph of node v as graph structure that surrounds v ’ s neighborhood . The context graph is described by two hyperparameters , r1 and r2 , and it represents a subgraph that is between r1-hops and r2-hops away from v ( i.e. , it is a ring of width r2 − r1 ) . Examples of neighborhood and context graphs are shown in Figure 2 ( a ) . We require r1 < K so that some nodes are shared between the neighborhood and the context graph , and we refer to those nodes as context anchor nodes . These anchor nodes provide information about how the neighborhood and context graphs are connected with each other . Encoding context into a fixed vector using an auxiliary GNN . Directly predicting the context graph is intractable due to the combinatorial nature of graphs . This is different from natural language processing , where words come from a fixed and finite vocabulary . To enable context prediction , we encode context graphs as fixed-length vectors . To this end , we use an auxiliary GNN , which we refer to as the context GNN . As depicted in Figure 2 ( a ) , we first apply the context GNN ( denoted as GNN′ in Figure 2 ( a ) ) to obtain node embeddings in the context graph . We then average embeddings of context anchor nodes to obtain a fixed-length context embedding . For node v in graph G , we denote its corresponding context embedding as cGv . Learning via negative sampling . We then use negative sampling ( Mikolov et al. , 2013 ; Ying et al. , 2018a ) to jointly learn the main GNN and the context GNN . The main GNN encodes neighborhoods to obtain node embeddings . The context GNN encodes context graphs to obtain context embeddings . In particular , the learning objective of Context Prediction is a binary classification of whether a particular neighborhood and a particular context graph belong to the same node : σ ( h ( K ) > v c G′ v′ ) ≈ 1 { v and v′ are the same nodes } , ( 3.1 ) where σ ( · ) is the sigmoid function , and 1 ( · ) is the indicator function . We either let v′ = v and G′ = G ( i.e. , a positive neighborhood-context pair ) , or we randomly sample v′ from a randomly chosen graph G′ ( i.e. , a negative neighborhood-context pair ) . We use a negative sampling ratio of 1 ( one negative pair per one positive pair ) , and use the negative log likelihood as the loss function . After pre-training , the main GNN is retained as our pre-trained model

Summary: The authors introduce strategies for pre-training graph neural networks. Pre-training is done at the node level as well as at the graph level. They evaluate their approaches on two domains, biology and chemistry on a number of downstream tasks. They find that not all pre-training strategies work well and can in fact lead to negative transfer. However, they find that pre-training in general helps over non pre-training.