[Paper Review] Self-Supervised Graph Transformer on Large-Scale Molecular Data

Yu Rong, Yatao Bian+5 more|arXiv (Cornell University)|Jun 18, 2020

Computational Drug Discovery Methods62 references415 citations

TL;DR

GROVER pre-trains a graph transformer with self-supervised tasks on 10M unlabeled molecules, achieving large gains (average >6%) on 11 MoleculeNet benchmarks after fine-tuning.

ABSTRACT

How to obtain informative representations of molecules is a crucial prerequisite in AI-driven drug design and discovery. Recent researches abstract molecules as graphs and employ Graph Neural Networks (GNNs) for molecular representation learning. Nevertheless, two issues impede the usage of GNNs in real scenarios: (1) insufficient labeled molecules for supervised training; (2) poor generalization capability to new-synthesized molecules. To address them both, we propose a novel framework, GROVER, which stands for Graph Representation frOm self-superVised mEssage passing tRansformer. With carefully designed self-supervised tasks in node-, edge- and graph-level, GROVER can learn rich structural and semantic information of molecules from enormous unlabelled molecular data. Rather, to encode such complex information, GROVER integrates Message Passing Networks into the Transformer-style architecture to deliver a class of more expressive encoders of molecules. The flexibility of GROVER allows it to be trained efficiently on large-scale molecular dataset without requiring any supervision, thus being immunized to the two issues mentioned above. We pre-train GROVER with 100 million parameters on 10 million unlabelled molecules -- the biggest GNN and the largest training dataset in molecular representation learning. We then leverage the pre-trained GROVER for molecular property prediction followed by task-specific fine-tuning, where we observe a huge improvement (more than 6% on average) from current state-of-the-art methods on 11 challenging benchmarks. The insights we gained are that well-designed self-supervision losses and largely-expressive pre-trained models enjoy the significant potential on performance boosting.

Motivation & Objective

• Motivate the need for robust molecular representations under limited labeled data and large chemical space.
• Propose a self-supervised pre-training framework for molecular graphs to improve generalization.
• Design a Transformer-based encoder with graph-aware attention and dynamic message passing.
• Show that pre-training on large unlabeled data boosts downstream molecular property prediction.

Proposed method

• Introduce GROVER: a graph representation from self-supervised message passing transformer with node and edge GNN transformers.
• Use a bi-level information extraction: GNN-based queries/keys/values feeding a Transformer encoder over all nodes.
• Implement dynamic message passing (dyMPN) with randomized hop counts to improve generalization.
• Design node/edge contextual property prediction as node/edge-level self-supervision.
• Add graph-level motif prediction using RDKit-detected motifs as multi-label targets for graph representations.
• Pre-train on 11M unlabeled molecules (ZINC15, ChEMBL) with 100M parameter model across 250 GPUs.

Experimental results

Research questions

• RQ1 Can self-supervised pre-training on large unlabeled molecular graphs improve downstream property prediction after fine-tuning?
• RQ2 Do graph-aware Transformer encoders and dynamic message passing yield superior representations compared to traditional GNNs?
• RQ3 What is the impact of context-aware node/edge and motif-based graph-level pretext tasks on performance and generalization?
• RQ4 How does GROVER scale with model size and training data on MoleculeNet benchmarks?

Key findings

• GROVER models consistently achieve the best performance on all 11 datasets, with an average relative improvement of 6.1% (2.2% on classification, 10.8% on regression).
• GROVER_large surpasses state-of-the-art baselines on all datasets; GROVER_base surpasses on 8/11 datasets.
• Self-supervised pre-training provides an average AUC increase of 3.8% over no pre-training on classification tasks, especially helping small datasets.
• In ablations, GROVER with its GTransformer backbone outperforms GIN and MPNN backbones, confirming higher expressiveness; dyMPN improves generalization despite slight training loss impact.
• GROVER achieves notable gains on small-label regime (e.g., FreeSolv) with 23.9% relative improvement over SOTA.