PaperlayerView exact revision on arXiv

Source-linked AI summary

Self-Supervised Graph Transformer on Large-Scale Molecular Data

Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, Junzhou Huang

arXiv:2007.02835v2q-bio.BMcs.LG

TL;DR

Molecular property learning is challenged by scarce labels and poor generalization to new molecules. GROVER pretrains a message-passing Transformer with self-supervised tasks on unlabelled molecular graphs, then fine-tunes it for downstream prediction. The authors report more than 6% average improvement over current state-of-the-art methods across 11 benchmarks.

  • Problem

    Molecular property learning faces insufficient labeled data and poor generalization to newly synthesized molecules.

    Source: §1 · p. 2

  • Method

    GROVER combines message-passing networks with a Transformer-style encoder and node-, edge-, and graph-level self-supervised tasks on unlabelled molecules.

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

  • Results

    More than 6% average improvement over current state-of-the-art methods was reported across 11 molecular property prediction benchmarks.

    Source: §6 · p. 9

  • Takeaways & Limitations

    The authors conclude that well-designed self-supervision and expressive pretrained models can substantially boost performance on molecular property prediction benchmarks.

    Source: §6 · p. 9

    The authors call for more self-supervised tasks and broader downstream task coverage, noting that different task categories may favor different pretraining strategies.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Molecular graphs preserve structural information, but molecular property prediction faces scarce labels and poor generalization to new molecules. GROVER combines self-supervised learning with a Transformer-style message-passing architecture and reports improved results on 11 benchmarks.

1 Introduction

  • Molecular tasks have few labels because obtaining molecular-property labels requires time-consuming, resource-costly wet-lab experiments.The introduction notes that public molecular benchmarks consequently contain far-from-adequate labels.Source: §1 · p. 2
  • Models trained on these benchmarks are prone to overfitting and can struggle with out-of-distribution molecules in the enormous chemical space.The authors identify both insufficient labeled data and poor generalization as impediments to real-world use.Source: §1 · p. 2
  • Earlier SMILES-based pretraining approaches do not explicitly encode molecular topology, while prior graph pretraining has limitations in its masking and graph-level tasks.The authors criticize atom-type masking for representation ambiguity and supervised graph-property pretraining for dependence on labels and risk of negative transfer.Source: §1 · p. 2
  • More than 6% average relative improvement over current state-of-the-art methods was reported across 11 molecular property prediction benchmarks.The introduction also reports 22.4% relative improvement over [29] and 7.4% over [18] on classification tasks.Source: §1 · p. 2
  • GROVER uses node- and edge-level contextual prediction plus graph-level semantic-motif prediction to learn from unlabelled molecular data.Motifs are obtained by low-cost pattern matching, so graph-level training does not require property labels.Source: §1 · p. 2
  • GROVER integrates message-passing networks into a Transformer-style architecture to encode molecular structure and semantics.The model was pretrained with 100 million parameters on 10 million unlabelled molecules.Source: §1 · p. 2

2 Related Work

Molecular representation work spans fingerprints, sequence models, and graph neural networks, while graph self-supervision offers another route to learned representations. The section contrasts GROVER’s approach with prior molecular pretraining methods.

Molecular Representation Learning

  • Molecular representations include fixed-length ECFP fingerprints, learned neural fingerprints [10], SMILES-based RNNs [21], and graph neural networks [23] [45] [46].Other graph approaches include attention-based aggregation [44], message passing for atom interactions [13], bond-interaction extensions [63], and hierarchical GNNs [30].Source: §2 · pp. 2–3

Self-supervised Learning on Graphs

  • Graph self-supervision has used graph proximity [14] [37], random-walk objectives [15], and mutual-information maximization [36] [50] [55].Molecular graph pretraining includes Liu et al.’s N-gram approach [29] and Hu et al.’s three-task approach [18].Source: §2 · p. 3
  • Hu et al. [18] separate correlated context and node/edge prediction tasks, while their supervised graph-property task depends on scarce labels and has shown negative transfer downstream.The authors contrast this with their method’s preservation of domain knowledge between local structure and node attributes, and its reported downstream effectiveness without negative transfer.Source: §2 · p. 3

3 Preliminaries of Transformer-style Models and Graph Neural Networks

The preliminaries describe supervised graph prediction, Transformer attention, and GNN message passing. Together, these concepts establish how node and graph representations are formed.

Supervised learning tasks of graphs

  • A molecular graph has atoms as nodes and bonds as edges; supervised graph learning includes node-level and graph-level classification or regression.Graph-level tasks use labeled graphs to predict the target of a new graph.Source: §3 · p. 3

Attention mechanism and the Transformer-style architectures

  • Transformer attention uses queries, keys, and values; scaled dot-product attention computes query-key scores and softmax weights over values.Stacking attention inputs as matrices supports optimized matrix multiplication.Source: §3 · p. 3

Graph Neural Networks (GNNs)

  • GNN message passing iteratively updates each node by aggregating information from neighboring nodes and edges across iterations and hops.After message passing, a readout operation produces a graph-level representation.Source: §3 · p. 4

4 The GROVER Pre-training Framework

GROVER combines graph-tailored Transformer attention and dynamic message passing with self-supervised tasks to learn molecular representations. Its node, edge, and graph embeddings can then be fine-tuned for downstream prediction.

4.1 Details of Model Architecture

  • GTransformer derives attention queries, keys, and values from graph nodes using tailored GNNs, combining local structural information with global relations between nodes.A single long-range residual connection also carries initial node or edge features to the final layers and is described as alleviating vanishing gradients and over-smoothing.Source: §4.1 · pp. 4–5, Figure 2 · p. 5
  • Randomized message-passing hops create random receptive fields and yield better generalization than vanilla message passing without randomization.The authors report this comparison from extensive experimental verification.Source: §4.1 · p. 5

4.2 Self-supervised Task Construction for Pre-training

  • Pre-training uses contextual-property prediction for nodes and edges and graph-level motif prediction, rather than supervised labels.Contextual targets are local-subgraph statistics; motif targets indicate whether recurrent motifs, including functional groups, occur in a molecule.Source: §4.2 · pp. 5–6

4.3 Fine-tuning for Downstream Tasks

  • For downstream prediction, GROVER's node and edge embeddings are pooled into a graph embedding, then an MLP predicts properties while the encoder and task-specific parameters are fine-tuned.The described workflow applies to graph-level molecular property prediction.Source: §4.3 · p. 6

5 Experiments

GROVER is evaluated on molecular property benchmarks and through experiments testing pre-training and architectural choices. Its strongest reported result is consistent improvement over state-of-the-art methods, while ablations examine pre-training and model variants.

5 Experiments

  • Experiments use 11 MoleculeNet benchmarks with scaffold splits and compare GROVER against 10 popular baselines and several state-of-the-art approaches.The data include 11 million unlabelled molecules sampled from ZINC15 and ChEMBL for pre-training.Source: §5 · p. 6

5.1 Results on Downstream Tasks

  • 6.1% overall relative improvement over state-of-the-art methods across all datasets, with GROVERlarge outperforming them on all 11 datasets.The relative improvements are 2.2% on classification and 10.8% on regression tasks.Source: Table 1 · p. 7, §5.1 · p. 8

5.2.1 How Useful is the Self-supervised Pre-training?

  • Pre-training raises average classification AUC by 3.8% over GROVER without pre-training, with larger gains on datasets with fewer samples.The comparison uses the same hyper-parameter setting; SIDER, ClinTox, and BACE are named as examples of smaller datasets.Source: §5.2.1 · p. 8, Table 2 · p. 9

5.2.2 How Powerful is GTransformer Backbone?

  • Figure 4 compares training and validation losses across different backbones.The supplied passage does not report which backbone performs best.Source: Figure 4 · p. 8

5.2.3 Effect of the Proposed dyMPN and GTransformer.

  • GROVER without GTransformer performs worst on training and validation, while dyMPN slightly worsens training loss but improves validation loss.The authors attribute dyMPN’s better generalization to randomizing the receptive field at each message-passing step.Source: §5.2.3 · p. 8, Figure 5 · p. 9

6 Conclusion and Future Works

The paper reports that fine-tuning GROVER improves results over state-of-the-art methods on 11 molecular property benchmarks. It identifies additional self-supervised tasks, downstream tasks, and larger models as future directions for GNN pre-training.

6 Conclusion and Future Works

  • Fine-tuning GROVER improves performance by more than 6% on average over current state-of-the-art methods on 11 challenging molecular property benchmarks.The paper presents this as evidence of the potential of self-supervised pre-trained approaches in graph learning.Source: §6 · p. 9
  • Future work includes distance-preserving and 3D-informed pre-training tasks, more downstream task categories, and wider or deeper models.The authors note that very large GNNs may require addressing gradient vanishing and oversmoothing.Source: §6 · p. 9

Broader Impact

GROVER uses self-supervised pretraining on unlabelled molecules to extract information that downstream tasks can use. The authors describe potential benefits for both machine learning research and drug discovery.

Broader Impact

  • GROVER extracts useful implicit information from massive unlabelled molecules, which the authors say downstream tasks can benefit from.Source: Broader Impact · p. 9
  • GROVER demonstrates pre-training on Graph Neural Networks and may open a venue for exploring pre-trained GNNs in social networks and knowledge graphs.The authors frame these as potential broader applications.Source: Broader Impact · p. 9
  • For drug discovery, GROVER produces atom feature vectors and molecule fingerprints for downstream tasks and can be fine-tuned with specific tasks.The authors identify molecular property prediction and virtual screening as example applications.Source: Broader Impact · p. 9

A The Overall Architecture of GROVER Model

GROVER combines node-view and edge-view GTransformers to produce molecular embeddings from two information flows. Its edge contextual properties are built from local subgraph statistics, and its property-prediction model combines both views.

A The Overall Architecture of GROVER Model

A.1 Fine-tuning Model for Molecular Property Prediction

  • For molecular property prediction, a shared self-attentive READOUT combines the two node embeddings, while separate MLPs produce predictions with a disagreement loss.The loss combines supervised losses for both predictions with the L2 distance between them.Source: §A.1 · p. 15

A.2 Constructing Contextual Properties for Edges

  • Edge contextual properties count node-edge pairs in a k-hop local subgraph and alphabetically list those counts to form the property.Subgraphs with the same statistical property receive the same contextual property, so this also acts as clustering by those statistics.Source: §A.2 · pp. 15–16

B.1 Dataset Description

The section describes benchmark datasets and a scaffold-splitting protocol that tests on molecules with structural frameworks distinct from training molecules. It also identifies the atom- and bond-feature tables used in GROVER.

B.1 Dataset Description

  • The benchmarks span molecular classification and regression datasets, including toxicity, drug-response, solubility, and quantum-property tasks.The dataset descriptions include Tox21, ToxCast, SIDER, ClinTox, BACE, QM7, QM8, ESOL, Lipophilicity, FreeSolv, and BBBP.Source: §B.1 · p. 16
  • Scaffold splitting separates molecules with distinct two-dimensional structural frameworks, creating a more challenging but practical train, validation, and test setting.Source: §B.1 · p. 17

B.2 Feature Extraction Processes for Molecules

C Implementation and Pre-training Details

The appendix reports a regression-task comparison of GROVER with and without self-supervised pretraining, alongside implementation and fine-tuning settings. Fine-tuning uses a task-specific hyperparameter search.

C Implementation and Pre-training Details

  • Pretraining uses 500 epochs, four attention heads, and a learning rate that warms up for two epochs before decaying exponentially.The learning rate starts at 0.00015 and decreases to 0.00001; the reported hidden sizes are 800 and 1200.Source: C Implementation and Pre-training… · p. 17

D Fine-tuning Details

E.1 Effect of Self-supervised Pre-training on Regression Tasks

  • Table 7 compares pre-trained GROVER with GROVER without pre-training on regression tasks.The supplied table caption identifies the comparison but gives no result values.Source: Table 7 · p. 18, §E.1 · p. 18

E.2 GROVER Fine-tuning Tasks with Other Backbones

GROVER-GIN matches Hu et al. without pre-training and is nearly 2% higher after pre-training, supporting the proposed self-supervised loss.

E.2 GROVER Fine-tuning Tasks with Other Backbones

Related papers

Loading 2007.02835v2…