ChemRL-GEM: Geometry Enhanced Molecular Representation Learning for Property Prediction

Xiaomin Fang; Lihang Liu; Jieqiong Lei; Donglong He; Shanzhuo Zhang; Jingbo Zhou; Fan Wang; Hua Wu; Haifeng Wang

Source-linked AI summary

ChemRL-GEM: Geometry Enhanced Molecular Representation Learning for Property Prediction

Xiaomin Fang, Lihang Liu, Jieqiong Lei, Donglong He, Shanzhuo Zhang, Jingbo Zhou, Fan Wang, Hua Wu, Haifeng Wang

arXiv:2106.06130v4Jun 11, 2021cs.LGphysics.chem-phq-bio.MN

TL;DR

Molecular property prediction needs representations that capture 3D geometry in addition to topology, especially when labeled molecules are scarce. ChemRL-GEM combines a dual-graph geometry-aware GNN with geometry-level self-supervised learning, and outperforms competitive baselines across molecular benchmarks. The paper reports state-of-the-art results on 11 of 12 datasets and an 8.8% average relative improvement on regression tasks.

Problem
Existing GNN representations and pre-training strategies mainly encode molecular topology, neglecting 3D geometry that helps determine physical, chemical, and biological properties.
Method
ChemRL-GEM uses GeoGNN with atom-bond and bond-angle graphs, together with geometry-level self-supervised tasks predicting local and global molecular spatial structure.
Results
ChemRL-GEM achieves state-of-the-art results on 11/12 datasets, with an 8.8% average relative improvement over previous SOTA results on regression tasks.
Takeaways & Limitations
The experiments demonstrate that incorporating molecular geometry improves ChemRL-GEM's performance over competitive baselines across multiple molecular property prediction benchmarks.
Takeaways & Limitations
N-Gram on ToxCast was not completed because it was too time-consuming and was planned for later addition.

Abstract

from arXiv · show

Effective molecular representation learning is of great importance to facilitate molecular property prediction, which is a fundamental task for the drug and material industry. Recent advances in graph neural networks (GNNs) have shown great promise in applying GNNs for molecular representation learning. Moreover, a few recent studies have also demonstrated successful applications of self-supervised learning methods to pre-train the GNNs to overcome the problem of insufficient labeled molecules. However, existing GNNs and pre-training strategies usually treat molecules as topological graph data without fully utilizing the molecular geometry information. Whereas, the three-dimensional (3D) spatial structure of a molecule, a.k.a molecular geometry, is one of the most critical factors for determining molecular physical, chemical, and biological properties. To this end, we propose a novel Geometry Enhanced Molecular representation learning method (GEM) for Chemical Representation Learning (ChemRL). At first, we design a geometry-based GNN architecture that simultaneously models atoms, bonds, and bond angles in a molecule. To be specific, we devised double graphs for a molecule: The first one encodes the atom-bond relations; The second one encodes bond-angle relations. Moreover, on top of the devised GNN architecture, we propose several novel geometry-level self-supervised learning strategies to learn spatial knowledge by utilizing the local and global molecular 3D structures. We compare ChemRL-GEM with various state-of-the-art (SOTA) baselines on different molecular benchmarks and exhibit that ChemRL-GEM can significantly outperform all baselines in both regression and classification tasks. For example, the experimental results show an overall improvement of 8.8% on average compared to SOTA baselines on the regression tasks, demonstrating the superiority of the proposed method.

1 Introduction

Molecular property prediction is important for drug and materials discovery, but topology-only GNNs and self-supervised tasks underuse molecular geometry. ChemRL-GEM addresses this with geometry-aware modeling and geometry-level pre-training.

Molecular property prediction supports evaluating, selecting, and generating molecules for drug and materials discovery.
GNNs model atoms and bonds as topological graphs, while self-supervised pre-training helps address scarce labeled molecular data.
Topology-only molecular representations neglect 3D geometry and may fail to distinguish molecules with identical topology but different geometries.
Existing self-supervised tasks focus on nodes, edges, or contexts and do not directly capture molecular geometry knowledge such as bond-angle differences.
ChemRL-GEM introduces GeoGNN with separate atom-bond and bond-angle graphs, plus geometry-level self-supervised tasks for molecular 3D spatial knowledge.

2 Preliminaries

Molecules can be represented as graphs, and GNNs learn node and graph representations through iterative neighborhood message passing. Molecular self-supervised learning complements this framework with local and global graph-level tasks.

A molecule is represented as G = (V, E), where atoms are nodes and chemical bonds are edges connecting neighboring atoms.
GNNs take node features, edge features, and graph structure as inputs to learn representations of nodes and the entire molecular graph.
At each iteration, a node representation is updated by aggregating messages from neighboring nodes and applying an update function.
A permutation-invariant READOUT function pools final node representations into a graph representation used for downstream predictions.
Molecular self-supervised learning uses node- or edge-level tasks for local knowledge and graph-level tasks for global information.

3 The ChemRL-GEM Framework

ChemRL-GEM combines a geometry-based GNN with geometry-level self-supervised tasks to represent molecular topology and 3D structure. GeoGNN uses atom-bond and bond-angle graphs, while pre-training targets local and global geometric information.

3.1 Geometry-based Graph Neural Network: GeoGNN models molecular geometry through two linked graphs: an atom-bond graph and a bond-angle graph.The atom-bond graph represents atoms and bonds; the bond-angle graph represents bonds and the angles connecting consecutive bonds.
3.1 Geometry-based Graph Neural Network: Bond representations learned from the bond-angle graph serve as edge features in the atom-bond graph, linking the two message-passing processes.Bond-angle information is aggregated first, and the resulting bond representations help update atom representations.
3.1 Geometry-based Graph Neural Network: The final atom representations are integrated into a molecular vector used to predict molecular properties.GeoGNN iteratively updates atom and bond representations before applying a readout function.
3.2 Geometry-level Self-supervised Learning Tasks: GeoGNN is pre-trained with three geometry-level tasks: bond-length prediction, bond-angle prediction, and atomic-distance-matrix prediction.Bond lengths and angles capture local spatial structure, whereas atomic distance matrices capture global spatial structure.
3.2.1 Local Spatial Structures: Local geometry tasks mask 15% of selected atoms and their one-hop contexts, then predict extracted bond lengths and bond angles.The masked context includes neighboring atoms, bonds, and bond angles; the resulting predictions form node-level self-supervised tasks.
3.2.2 Global Spatial Structures: The global geometry task predicts discretized atomic distances from distance matrices constructed using molecular 3D coordinates.It treats distance prediction as multi-class classification and emphasizes the distribution of global spatial structures.

4 Experiments

Experiments evaluate ChemRL-GEM across 12 MoleculeNet benchmarks using scaffold splits, simulated 3D geometry, and comparisons with pretrained and non-pretrained baselines. ChemRL-GEM achieves leading performance overall, with geometry-aware architecture and pre-training contributing to regression gains.

Experimental Setup: ChemRL-GEM is evaluated against multiple state-of-the-art baselines on 12 MoleculeNet datasets covering classification and regression tasks.The experiments use scaffold splitting, with training, validation, and test sets divided 8:1:1.
Experimental Setup: 20 million unlabeled Zinc15 molecules are used to pre-train GeoGNN with geometry-level and graph-level self-supervised tasks.Simulated 3D coordinates provide bond lengths, bond angles, and atomic distance matrices for geometry-level tasks.
Overall Performance: ChemRL-GEM achieves SOTA results on 11/12 datasets, with an 8.8% average relative improvement on regression tasks and 3.7% average ROC-AUC improvement on classification tasks over previous SOTA.The previous classification SOTA comparison is D-MPNN.
Ablation Studies: GeoGNN yields a 7.9% overall relative improvement over previous methods on regression datasets.The comparison includes GIN, GAT, GCN, D-MPNN, AttentiveFP, and GTransformer.
Ablation Studies: More accurate QM9 coordinates reduce GeoGNN’s average MAE from 0.00746 with coarse coordinates to 0.00652.The result supports stronger representation learning when more accurate 3D coordinates are available.
Ablation Studies: Geometry-level pre-training tasks outperform GeoGNN without pre-training on regression datasets, while adding graph-level tasks can reduce performance.The paper attributes the latter pattern to a weak connection between molecular fingerprints and regression tasks.

5 Related Work

Prior molecular representations include fingerprints, sequence-based encodings, and graph-based representations, while self-supervised learning addresses limited labeled data. These approaches generally leave limitations in representing global structure, unambiguous molecular form, or complete geometry.

Representation Types: Molecular representations are categorized as fingerprints, sequence-based representations, and graph-based representations.The categories organize the main representation strategies discussed in related work.
Molecular Fingerprints: Molecular fingerprints encode topological substructures as bits but lack complex global structures because they focus on local substructures.Examples include ECFP and MACCS.
Sequence-based Representations: Sequence-based methods use SMILES with recurrent or Transformer models, but the same molecule can have multiple SMILES representations.This creates representation ambiguity and complicates modeling some molecular topological structures.
Graph-based Representations: Graph-based methods model atoms as nodes and chemical bonds as edges, with some approaches adding atomic distances as edge features.These methods still lack complete molecular geometry modeling, including bond-angle information.
Self-supervised Learning: Self-supervised learning trains unlabeled samples in a supervised manner and has been applied to molecular GNNs to address insufficient labeled samples.The related methods learn molecular representations through self-supervised tasks.

6 Conclusions and Future Work

ChemRL-GEM combines geometry-aware message passing with geometry-level self-supervision and outperforms competing methods across 12 molecular benchmarks. Future work will extend the representation to additional geometric parameters and molecule-related prediction tasks.

ChemRL-GEM models atom-bond-angle relations with a geometry-based GNN and uses multiple geometry-level self-supervised methods to capture spatial knowledge.
ChemRL-GEM significantly outperforms multiple competitive methods on 12 molecular property prediction benchmarks.
Future work will consider torsional angles and apply molecular representations to drug-target and intermolecular interaction prediction.

A.1 GNN Architecture

GeoGNN uses parallel atom-bond and bond-angle processing to model molecular structure, with task-specific headers built on shared node or graph representations. Its inputs combine atom, bond, and bond-angle features, including continuous-feature expansions and discrete one-hot encodings.

Architecture: GeoGNN consists of two stacks of blocks, one for the atom-bond graph and one for the bond-angle graph.GIN provides the message-passing backbone, with layer normalization, graph size normalization, and residual connections used in the architecture.
Architecture: The atom-bond graph and bond-angle graph use aggregate functions to summarize their respective node and edge features.The combine function is a two-layer MLP with hidden size 32; eight GeoGNN blocks are used for each graph.
Task headers: Different self-supervised and downstream headers share GeoGNN as their common foundation, with geometry-level headers applied to node representations h(K).
Task headers: Downstream headers operate on graph representations hG through a three-layer MLP, with separate functions representing different downstream tasks.
Input features: GeoGNN inputs include atom, bond, and bond-angle features extracted with RDKit and the Merck molecular force field function.Bond lengths and bond angles are continuous features expanded with radial basis functions, while discrete features become one-hot vectors.

B Training and Test Processes for ChemRL-GEM

ChemRL-GEM first pre-trains GeoGNN on an unlabeled molecular dataset with self-supervised geometry objectives, then fine-tunes it on labeled downstream data. Validation selects the best epoch before final test evaluation.

Model: GeoGNN parameters are first pre-trained with self-supervised tasks and then fine-tuned for downstream objectives.
Optimization: The pre-training loss combines length, angle, and distance prediction objectives.
Training: Training iterates over the pre-training dataset to optimize the pre-training loss, then iterates over labeled downstream data to fine-tune the model.
Evaluation: Validation data selects the best epoch, and the selected model is evaluated on the downstream test dataset.

C.1 Dataset Description

The evaluation uses 12 MoleculeNet benchmarks divided evenly between regression and classification tasks, covering physical chemistry, quantum mechanics, biophysics, and physiology. The datasets include molecular properties ranging from solubility and hydration free energy to toxicity, permeability, binding, and quantum-mechanical measurements.

Benchmark design: The benchmark suite contains six molecular regression and six molecular classification datasets from MoleculeNet.All 12 datasets use scaffold splitting, which evaluates generalization on out-of-distribution molecular structures more challengingly than random splitting.
Benchmark design: The benchmarks span physical chemistry, quantum mechanics, biophysics, and physiology.Physical-chemistry and quantum-mechanics benchmarks are described as more related to molecular geometries than biophysics and physiology benchmarks.
Regression datasets: ESOL measures water solubility, FreeSolv records hydration free energy, and Lipophilicity contains electronic and related molecular properties.
Quantum-mechanical datasets: QM8 and QM96 provide quantum-mechanical measurements involving electronic spectra, excited-state energies, geometry, energy, and thermodynamic properties.
Classification datasets: BACE provides binary inhibitor-binding labels, BBBP measures blood-brain-barrier permeability, and ClinTox covers approved and failed drugs distinguished by toxicity outcomes.
Classification datasets: SIDER contains marketed-drug adverse drug reactions grouped into 27 system organ classes, while Tox21 and ToxCast provide multi-target or high-throughput toxicity labels.

C.2 Dataset Splitting Methods

The experiments use scaffold-based splitting to create more challenging and realistic molecular property prediction evaluations. The paper also standardizes scaffold calculation by accounting for molecular chirality and reports dataset statistics separately for classification and regression benchmarks.

Splitting strategies: Scaffold splitting and random scaffold splitting keep molecules sharing a scaffold within the same train, validation, or test set.These strategies are described as more challenging and realistic than random splitting.
Splitting strategies: Random scaffold splitting is adopted because its additional randomness supports evaluation across different methods.
Implementation consistency: Scaffold calculation includes molecular chirality to avoid inconsistent train, validation, and test partitions across baseline implementations.The authors found that some public implementations consider chirality while others do not, producing substantially different splits.