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: This paper proposes new pre-training strategies for GNN with both a node-level and a graph-level pretraining. For the node-level pretraining, the goal is to map nodes with similar surrounding structures to nearby context (similarly to word2vec). The main problem is that directly predicting the context is intractable because of combinatorial explosion. The main idea is then to use an additional GNN to encode the context and to learn simultaneously the main GNN and the context GNN via negative sampling. Another method used is attribute masking where some masked node and edge attributes need to be predicted by the GNN. For graph-level pretraining, some general graph properties need to be predicted by the graph.