Strategies for Pre-training Graph Neural Networks

TL;DR

The paper tackles the challenge of transfer learning on graphs under label scarcity and distribution shift. It introduces a dual-level pre-training framework that combines node-level self-supervision (Context Prediction and Attribute Masking) with graph-level multi-task supervision, paired with an expressive GNN (GIN). Large-scale experiments on chemistry (MoleculeNet) and biology (PPI) demonstrate robust out-of-distribution generalization and state-of-the-art performance, with the combined strategy avoiding negative transfer and delivering substantial ROC-AUC gains. The approach also yields orders-of-magnitude faster fine-tuning and provides a practical pathway for scalable graph representation learning across domains.

Abstract

Many applications of machine learning require a model to make accurate pre-dictions on test examples that are distributionally different from training ones, while task-specific labels are scarce during training. An effective approach to this challenge is to pre-train a model on related tasks where data is abundant, and then fine-tune it on a downstream task of interest. While pre-training has been effective in many language and vision domains, it remains an open question how to effectively use pre-training on graph datasets. In this paper, we develop a new strategy and self-supervised methods for pre-training Graph Neural Networks (GNNs). The key to the success of our strategy is to pre-train an expressive GNN at the level of individual nodes as well as entire graphs so that the GNN can learn useful local and global representations simultaneously. We systematically study pre-training on multiple graph classification datasets. We find that naive strategies, which pre-train GNNs at the level of either entire graphs or individual nodes, give limited improvement and can even lead to negative transfer on many downstream tasks. In contrast, our strategy avoids negative transfer and improves generalization significantly across downstream tasks, leading up to 9.4% absolute improvements in ROC-AUC over non-pre-trained models and achieving state-of-the-art performance for molecular property prediction and protein function prediction.

Figures (6)

• Figure 1: (a.i) When only node-level pre-training is used, nodes of different shapes (semantically different nodes) can be well separated, however, node embeddings are not composable, and thus resulting graph embeddings (denoted by their classes, $+$ and $-$) that are created by pooling node-level embeddings are not separable. (a.ii) With graph-level pre-training only, graph embeddings are well separated, however the embeddings of individual nodes do not necessarily capture their domain-specific semantics. (a.iii) High-quality node embeddings are such that nodes of different types are well separated, while at the same time, the embedding space is also composable. This allows for accurate and robust representations of entire graphs and enables robust transfer of pre-trained models to a variety of downstream tasks. (b) Categorization of pre-training methods for GNNs. Crucially, our methods, i.e., Context Prediction, Attribute Masking, and graph-level supervised pre-training (Supervised Attribute Prediction) enable both node-level and graph-level pre-training.
• Figure 2: Illustration of our node-level methods, Context Prediction and Attribute Masking for pre-training GNNs. (a) In Context Prediction, the subgraph is a $K$-hop neighborhood around a selected center node, where $K$ is the number of GNN layers and is set to 2 in the figure. The context is defined as the surrounding graph structure that is between $r_1$- and $r_2$-hop from the center node, where we use $r_1=1$ and $r_2 = 4$ in the figure. (b) In Attribute Masking, the input node/edge attributes (e.g., atom type in the molecular graph) are randomly masked, and the GNN is asked to predict them.
• Figure 3: Test ROC-AUC of protein function prediction using different pre-training strategies with GIN.(Left) Test ROC-AUC scores (%) obtained by different pre-training strategies, where the scores are averaged over the 40 fine-grained prediction tasks. (Middle and right): Scatter plot comparisons of ROC-AUC scores for a pair of pre-training strategies on the 40 individual downstream tasks. Each point represents a particular individual downstream task. (Middle): There are many individual downstream tasks where graph-level multi-task supervised pre-trained model performs worse than non-pre-trained model, indicating negative transfer. (Right): When the graph-level multi-task supervised pre-training and Attribute Masking are combined, negative transfer is avoided across downstream tasks. The performance also improves over pure graph-level supervised pre-training.
• Figure 4: Training and validation curves of different pre-training strategies on GINs. Solid and dashed lines indicate training and validation curves, respectively.
• Figure 5: Training and validation curves of different pre-training strategies. The solid and dashed lines indicate the training and validation curves, respectively.