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.

    Source: §1 · p. 1

  • Method

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

    Source: §1 · p. 2, §3.1.1 · p. 4

  • Results

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

    Source: §1 · p. 3, §6 · p. 9

  • Takeaways & Limitations

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

    Source: §1 · p. 2, §6 · p. 9

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

    Source: §3.2.2 · p. 6

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.Source: §1 · p. 1
  • 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.Source: §1 · p. 1

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.Source: §1 · p. 2

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.Source: §1 · p. 2

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.Source: §1 · p. 2, Figure 1 · p. 2

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.Source: §1 · pp. 2–3

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.Source: §2 · p. 3

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.Source: §2 · p. 3
  • 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.Source: §2 · p. 3

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.Source: §2 · p. 3

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.Source: §3 · p. 3

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.Source: §3.1.1 · p. 4
  • 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.Source: §3.1.2 · p. 5

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.Source: §3.2 · p. 5, §3.2.1 · p. 5

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.Source: §3.2.1 · pp. 5–6

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.Source: §3.2.2 · p. 6

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.Source: §4 · p. 6
  • 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.Source: §4 · p. 6

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.Source: §5.1 · p. 7
  • 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.Source: §5.1 · p. 7
  • 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.Source: §5.2 · p. 7

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.Source: §5.3 · pp. 8–9
  • 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%.Source: §5.3 · p. 9
  • 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.Source: §5.3 · p. 9

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.Source: §6 · p. 9
  • The new strategy achieves consistently better out-of-distribution generalization than non-pre-trained models across multiple datasets, downstream tasks, and GNN architectures.Source: §6 · p. 9
  • The work addresses negative transfer in graph transfer learning and identifies improving architectures, training approaches, interpretability, and applications to other domains as future directions.Source: §6 · pp. 9–10

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.

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.

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.Source: E DETAILS OF DATASET SPLITTING · p. 19
  • 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.Source: E DETAILS OF DATASET SPLITTING · p. 19

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.

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.Source: G FURTHER DETAILS OF THE EXPERIMENTAL… · p. 20

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.Source: G FURTHER DETAILS OF THE EXPERIMENTAL… · p. 20

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.Source: G FURTHER DETAILS OF THE EXPERIMENTAL… · p. 20

Computation time

GNN architectures

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

Additional scatter plot comparisons

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.Source: Table 4 · p. 21

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.Source: Figure 6 · p. 22

Related papers

Loading 1905.12265v3…