PaperlayerView exact revision on arXiv

Source-linked AI summary

Strategies for Pre-training Graph Neural Networks

Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, Jure Leskovec

arXiv:1905.12265v3cs.LGstat.ML

TL;DR

Graph pre-training must support prediction under distribution shift when downstream labels are scarce, yet effective strategies for graph data remain unclear. This paper systematically studies GNN pre-training and combines node- and graph-level objectives with expressive architectures. Across downstream tasks, the strategy avoids negative transfer and improves generalization, reaching up to 9.4% higher average ROC-AUC than non-pre-trained GNNs.

  • Problem

    Graph pre-training remains difficult because downstream labels are scarce and graph test samples can be structurally out-of-distribution.

  • Method

    The paper pre-trains expressive GNNs at both node and graph levels using domain-specific information and self-supervised methods.

  • Results

    Across downstream tasks, the strategy avoids negative transfer and yields up to 9.4% higher average ROC-AUC than non-pre-trained GNNs.

  • Takeaways & Limitations

    Combining local and global pre-training with an expressive GNN provides consistently better out-of-distribution generalization than non-pre-trained models.

  • Takeaways & Limitations

    Structural-similarity pre-training is left for future work because graph-distance ground truth is difficult and graph pairs grow quadratically.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Pre-training on graphs addresses scarce task-specific labels and distribution shifts, but naive strategies can cause negative transfer. The paper develops a strategy combining node- and graph-level pre-training and reports stronger downstream generalization.

  • Motivation: Few studies had generalized the successful transfer-learning paradigm from vision and language to graph data.Transfer learning repurposes models trained on related tasks for new tasks.
  • Motivation: Graph pre-training targets scarce labels and structurally different test graphs, challenges common in chemistry and biology.Scientific graph labels can require resource- and time-intensive experiments.
  • Study design: The study introduces chemistry and biology pre-training datasets containing 2 million and 395K graphs, respectively, for systematic evaluation.It also argues that existing downstream benchmarks are too small for statistically reliable evaluation.
  • Challenges: Naive graph-level supervised pre-training produced negative transfer on 2 of 8 molecular datasets and 13 of 40 protein prediction tasks.The strategy otherwise provided only marginal gains across many downstream tasks.
  • Proposed strategy: The proposed strategy combines node-level and graph-level pre-training so GNNs learn domain-specific local and global representations.It uses node information to produce meaningful, composable node embeddings that support robust graph representations.
  • Results: The strategy with GIN achieves state-of-the-art benchmark results, avoids negative transfer, and improves generalization across tested downstream tasks.The reported gains reach 9.4% higher average ROC-AUC than non-pre-trained GNNs.

2 PRELIMINARIES OF GRAPH NEURAL NETWORKS

The preliminaries formalize graph prediction and describe GNN representations built through neighborhood aggregation. A permutation-invariant readout pools final node features into a graph-level embedding.

  • Supervised learning of graphs: A graph is represented as G = (V, E), with node attributes X_v and edge attributes e_uv, and graph labels predicted from h_G.The graph-level task learns h_G so that a function g can predict the label y_G.
  • Graph Neural Networks: GNNs learn an embedding h_v for every node and h_G for the entire graph using connectivity, node features, and edge features.Node representations are iteratively updated by aggregating neighboring nodes and edges.
  • Graph Neural Networks: After k aggregation iterations, each node representation captures structural information from its local neighborhood.The supplied passage introduces this role before defining the associated symbols.
  • Graph representation learning: The READOUT function pools node features from the final iteration K to obtain the graph representation h_G.READOUT is permutation-invariant and can use averaging or more sophisticated graph-level pooling.

3 STRATEGIES FOR PRE-TRAINING GRAPH NEURAL NETWORKS

The strategy pre-trains GNNs at both node and graph levels, combining self-supervised local objectives with graph-level supervision to learn transferable representations. It addresses limitations of graph-level-only pre-training, including negative transfer and costly task selection.

  • Core strategy: Pre-training GNNs at both individual-node and entire-graph levels encourages domain-specific representations at local and global scales.This contrasts with strategies that pre-train only on graph-level or only on node-level predictions.
  • Node-level pre-training: Context Prediction uses a node’s K-hop neighborhood to predict its surrounding context graph, bringing nodes with similar structural contexts closer in embedding space.The context lies between r1 and r2 hops from the center node, and an auxiliary GNN encodes it into a fixed-length vector for negative-sampling-based learning.
  • Node-level pre-training: Attribute Masking randomly hides node or edge attributes and trains the GNN to predict them from neighboring graph structure.In molecular graphs, masked attributes can include atom types; the method targets regularities in richly annotated scientific graphs.
  • Graph-level pre-training: Graph-level multi-task supervised pre-training jointly predicts diverse graph labels to inject domain-specific information into graph representations.Linear classifiers are applied on top of graph representations for the binary property-prediction tasks.
  • Motivation and limitations: Naïve graph-level multi-task pre-training can hurt transfer because unrelated tasks interfere, while selecting relevant tasks requires substantial domain expertise and separate pre-training.The proposed strategy first regularizes node representations, addressing the weakness of graph-level supervision alone.
  • Alternative graph-level objectives: Graph-level structural-similarity prediction is left for future work because obtaining ground-truth graph distances is difficult and large datasets require considering a quadratic number of graph pairs.The paper identifies this as a natural but out-of-scope alternative.

4 FURTHER RELATED WORK

Prior graph representation-learning methods primarily learn embeddings for individual nodes using local or reconstruction-based objectives. These approaches generally use separate embeddings across substructures, limiting dataset transfer, end-to-end fine-tuning, and coverage of large contexts.

  • Unsupervised node representation learning includes local random-walk objectives and graph reconstruction objectives such as edge-existence prediction.
  • Existing node-embedding methods use distinct, non-shared embeddings for different substructures, making them inherently transductive.The passage also states that they cannot transfer between datasets or be fine-tuned end-to-end.

5 EXPERIMENTS

The experiments evaluate pre-training strategies across molecular and protein graph prediction tasks using realistic out-of-distribution splits. Combining graph-level supervised and node-level self-supervised pre-training with expressive GNNs avoids negative transfer and delivers the strongest performance, while classic small benchmarks are statistically unreliable.

  • Experimental setup: The study uses 2 million unlabeled molecules, 456K molecules with 1310 biochemical assays, and 395K unlabeled protein ego-networks for pre-training.Downstream evaluation uses eight larger MoleculeNet binary classification datasets rather than MUTAG and PTC alone.
  • Experimental setup: Out-of-distribution generalization is evaluated with scaffold-based chemistry splits and corresponding structurally distinct graph splits rather than conventional random splits.The split design reflects cases where test graphs differ structurally from training graphs.
  • Experimental setup: The experiments compare graph-level multi-task supervised pre-training, node-level self-supervised pre-training, and their combination, primarily with GIN and also with less expressive GNN architectures.The self-supervised methods include Context Prediction and Attribute Masking; coarse-grained labels are used for pre-training.
  • Results: Naive graph-level pre-training gives negative transfer on 2 of 8 molecular datasets and 13 of 40 protein tasks, while node-level pre-training alone provides only limited improvement.The graph-level baseline is a strong supervised multi-task strategy, yet its gains remain marginal overall.
  • Results: 7.2% average chemistry ROC-AUC improvement over the non-pre-trained baseline and 9.4% average biology ROC-AUC improvement demonstrate substantial gains across the two domains.On HIV, the best pre-trained model reaches 79.9% ROC-AUC, exceeding reported graph-model scores of 76.3%, 77.6%, and 76.2%.
  • Results: Pre-trained models converge orders of magnitude faster; on MUV, reaching 85.3% validation ROC-AUC took 5 minutes versus 1 hour to reach 74.9% without pre-training.The authors characterize pre-training as a one-time effort reusable across downstream tasks.

6 CONCLUSIONS AND FUTURE WORK

The paper’s combined node-level and graph-level pre-training strategy improves out-of-distribution generalization across datasets, tasks, and GNN architectures while addressing negative transfer. Future work includes extending methods to additional scientific graph domains.

  • Combining node-level and graph-level pre-training with an expressive GNN yields meaningful local and global representations for downstream tasks.Node embeddings capture local neighborhood semantics and are pooled into graph-level representations.
  • The new strategy achieves consistently better out-of-distribution generalization than non-pre-trained models across multiple datasets, downstream tasks, and GNN architectures.
  • The work addresses negative transfer in graph transfer learning and identifies improving architectures, training approaches, interpretability, and applications to other domains as future directions.

A DETAILS OF GNN ARCHITECTURES

The experiments use GIN-based architectures with embedded node and edge features, iterative neighborhood updates, and mean-pooled graph representations. Hyperparameters are fixed to emphasize differences caused by pre-training strategies.

  • GIN architectures are used for molecular and protein prediction, with minor modifications for edge features and protein center-node information.
  • All experimental hyperparameters are fixed to compare relative improvements directly attributable to the pre-training strategies.
  • Molecular inputs use categorical node and edge features, including distinct categories for masked nodes, masked edges, and self-loop edges.
  • The final GNN layer omits ReLU so dot-product pre-training methods can produce negative similarities as well as positive ones.
  • Graph-level representations are formed by averaging final-layer node embeddings, and labels are predicted with a linear model on top of hG.
  • Protein inputs use uniform node features and binary edge-feature vectors, which are embedded before GNN processing.

B EXPERIMENTS ON CLASSIC GRAPH CLASSIFICATION BENCHMARKS

Classic graph benchmarks are evaluated with 10-fold cross-validation using the paper’s molecular and protein graph representations. Results are difficult to compare reliably on very small datasets because validation folds contain few molecules.

  • MUTAG and PTC are molecule datasets for binary classification, while NCI1 is excluded because it lacks edge information needed to reconstruct input representations.
  • The classic benchmark experiments evaluate graph classification with 10-fold cross-validation accuracy.The protocol follows Xu et al. (2019), with dropout and batch size additionally tuned during fine-tuning.
  • On MUTAG and PTC, all accuracies fall within one standard deviation, limiting reliable comparison because validation folds contain only about 19 to 35 molecules.
  • Molecular inputs use minimal node and bond features that unambiguously describe two-dimensional molecular structure.
  • The molecular benchmark suite includes eight binary graph-classification datasets from MoleculeNet.
  • Protein subgraphs contain edge attributes representing seven protein-protein relationship types, with weighted evidence thresholded for dataset construction.
  • Fine-grained protein functions are Gene Ontology leaf annotations, while coarse-grained functions are their immediate parents.

E DETAILS OF DATASET SPLITTING

The study uses distribution-shifting dataset splits designed to test generalization beyond familiar molecular scaffolds and species-specific protein annotations.

  • Molecular graphs are clustered by scaffold before constructing validation and test sets with structurally different molecules.This scaffold split is intended to provide a more realistic prospective-evaluation estimate than random splitting.
  • The protein species split uses half of human protein subgraphs as a test set and half as a prior set containing only coarse-grained annotations.The setup represents predicting fine-grained functions for proteins in a species with only high-level knowledge available for part of the data.

F TIME COMPLEXITY OF PRE-TRAINING

Both Attribute Masking and Context Prediction have at most linear time complexity in the number of graph edges, matching GNN message-passing efficiency while adding almost no memory overhead.

  • Both pre-training methods run in at most linear time with respect to the number of graph edges.Attribute Masking samples nodes or edges, while Context Prediction samples a center node and extracts its K-hop neighborhood and context graph.
  • Context Prediction uses breadth-first search to extract neighborhood and context graphs within at most linear time in the number of edges.
  • The methods are as efficient as ordinary supervised GNN learning because their complexity matches message-passing computation.
  • On-the-fly data transformations result in almost no memory overhead during pre-training.

G FURTHER DETAILS OF THE EXPERIMENTAL SETUP

The experiments use standardized optimization, fine-tuning, and ROC-AUC evaluation protocols across molecular and protein datasets, while also comparing architectures and reporting pre-training costs.

  • Optimization: All models use Adam with learning rate 0.001, and pre-training runs for 100 epochs.Self-supervised pre-training uses batch size 256; supervised pre-training uses batch size 32 and 20% dropout.
  • Fine-tuning: Fine-tuning uses batch size 32 and 50% dropout, with joint fitting for datasets containing multiple prediction tasks.Molecular models train for 100 epochs, while protein-function models train for 50 epochs.
  • Evaluation: Performance is measured by test ROC-AUC at the best validation epoch, averaged across tasks where applicable and reported over 10 random seeds.Results include mean ROC-AUC and standard deviation.
  • Computation time: Self-supervised pre-training takes about 24 hours in chemistry and 3.8 hours in biology, compared with 11 and 2.5 hours for supervised pre-training.
  • GNN architectures: The most expressive GIN architectures benefit most from pre-training on the chemistry datasets.

I ADDITIONAL TRAINING AND VALIDATION CURVES

Additional plots examine training dynamics and compare the combined Context Prediction plus graph-level supervised strategy with non-pre-trained and graph-level supervised baselines across molecular and protein tasks.

  • Training and validation curves: Figure 5 plots training and validation curves for all molecular property prediction datasets.Solid lines represent training curves, while dashed lines represent validation curves.
  • Additional scatter plot comparisons: The combined strategy completely avoids negative transfer across all 40 downstream protein-function tasks.
  • Molecular prediction benchmarks: Table 4 reports test ROC-AUC percentages across eight molecular benchmarks and GNN architectures using Context Prediction plus graph-level supervised pre-training.Its rightmost column averages mean test performance across the eight datasets.
  • Additional scatter plot comparisons: Figure 6 compares ROC-AUC scores across 40 fine-grained protein-function tasks for the combined strategy and two baseline strategies.The baselines are non-pre-trained and graph-level supervised pre-trained models.
Loading 1905.12265v3…