[Paper Review] Strategies for Pre-training Graph Neural Networks

Weihua Hu, Bowen Liu+5 more|arXiv (Cornell University)|May 29, 2019

Advanced Graph Neural Networks62 references186 citations

TL;DR

The paper introduces a dual-level pre-training strategy for GNNs—node-level self-supervised tasks and graph-level supervised tasks—combined with an expressive GNN (GIN) to improve out-of-distribution generalization and achieve state-of-the-art results, avoiding negative transfer.

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.

Motivation & Objective

• Address the challenge of learning on graph data with scarce task-specific labels and distribution shifts.
• Evaluate whether naïve pre-training strategies help or harm downstream graph tasks.
• Develop a combined node-level and graph-level pre-training approach to improve transferability.
• Demonstrate effectiveness across molecular property prediction and protein function prediction.
• Provide datasets and benchmarks to enable large-scale pre-training studies on graphs.

Proposed method

• Develop node-level pre-training via Context Prediction and Attribute Masking to capture local domain-specific knowledge.
• Develop graph-level pre-training via supervised multi-task pre-training to inject global graph-level signals.
• Jointly optimize main GNN and auxiliary context GNN with negative sampling for Context Prediction.
• Use an expressive Graph Isomorphism Network (GIN) as the backbone to maximize benefits from pre-training.
• Sequence: node-level self-supervision first, then graph-level supervised pre-training, followed by end-to-end fine-tuning.
• Show that naïve graph-only or node-only pre-training can cause negative transfer; the combined approach mitigates this.

Experimental results

Research questions

• RQ1 Does pre-training GNNs improve performance on downstream graph classification and function prediction tasks?
• RQ2 Do node-level and graph-level pre-training provide complementary benefits for transfer learning on graphs?
• RQ3 Is naive pre-training (graph-level only or node-level only) beneficial or detrimental across domains?
• RQ4 Which GNN architectures most benefit from the proposed pre-training strategy?

Key findings

• The combined node- and graph-level pre-training strategy yields up to 9.4% absolute ROC-AUC gains over non-pre-trained models.
• GIN with pre-training achieves state-of-the-art results on molecular property prediction and protein function prediction.
• Graph-level pre-training alone can cause negative transfer on a substantial portion of tasks; combining with node-level pre-training avoids this.
• Naïve extensive graph-level multi-task pre-training provides limited gains and may hurt downstream performance on several tasks.
• Pre-trained models converge orders of magnitude faster during fine-tuning than non-pre-trained models.