[Paper Review] Pre-training Molecular Graph Representation with 3D Geometry

Shengchao Liu, Hanchen Wang+4 more|arXiv (Cornell University)|Oct 7, 2021

Computational Drug Discovery Methods85 references150 citations

TL;DR

GraphMVP pre-trains a 2D molecular graph encoder by leveraging 3D geometry through contrastive and generative self-supervised tasks, yielding improved downstream molecular property predictions.

ABSTRACT

Molecular graph representation learning is a fundamental problem in modern drug and material discovery. Molecular graphs are typically modeled by their 2D topological structures, but it has been recently discovered that 3D geometric information plays a more vital role in predicting molecular functionalities. However, the lack of 3D information in real-world scenarios has significantly impeded the learning of geometric graph representation. To cope with this challenge, we propose the Graph Multi-View Pre-training (GraphMVP) framework where self-supervised learning (SSL) is performed by leveraging the correspondence and consistency between 2D topological structures and 3D geometric views. GraphMVP effectively learns a 2D molecular graph encoder that is enhanced by richer and more discriminative 3D geometry. We further provide theoretical insights to justify the effectiveness of GraphMVP. Finally, comprehensive experiments show that GraphMVP can consistently outperform existing graph SSL methods.

Motivation & Objective

• Motivate improved molecular representations by incorporating 3D geometry information into 2D graph encodings.
• Propose a self-supervised pre-training framework (GraphMVP) that uses contrastive and generative tasks across 2D and 3D views.
• Demonstrate that 3D geometry acts as privileged information to enhance 2D graph embeddings during pre-training.
• Analyze the effects of masking, conformer count, and objective functions on performance.
• Provide theoretical insights and empirical evidence for the benefits of multi-view pre-training with 3D geometry.

Proposed method

• Represent molecules with two views: a 2D topological graph and a 3D geometric conformer graph.
• Use two self-supervised pretext tasks: a contrastive objective (InfoNCE or EB-NCE) aligning 2D and 3D views across molecules, and a generative objective (VRR) reconstructing representations between views.
• Introduce a variational representation reconstruction (VRR) surrogate loss to measure cross-view reconstruction quality in representation space.
• Combine the two SSL objectives into a unified GraphMVP objective, with variants GraphMVP-G and GraphMVP-C incorporating 2D SSL auxiliary tasks.
• Pre-train on a shared dataset (50k molecules with 2D/3D structures and conformers) and fine-tune on downstream 8 low-data tasks using a 2D GNN backbone (GIN) and 3D geometry (SchNet) for pretraining signals.

Experimental results

Research questions

• RQ1 Does incorporating 3D geometric information into pre-training improve 2D molecular graph representations for downstream tasks?
• RQ2 How do contrastive and generative SSL signals between 2D and 3D views complement each other in GraphMVP?
• RQ3 What is the impact of masking ratio and number of conformers on learning and performance?
• RQ4 Is 3D geometry acting as privileged information that accelerates learning during pre-training but is unavailable at test time?
• RQ5 Do extensions that add 2D SSL objectives (GraphMVP-G, GraphMVP-C) yield further gains?

Key findings

• GraphMVP outperforms random initialization and most SSL baselines on 8 molecular property tasks, with notable average improvements.
• Hybrid variants GraphMVP-G and GraphMVP-C consistently improve over GraphMVP, confirming that 3D geometry complements 2D topology.
• VRR-based generative SSL with stochastic representation reconstruction provides robust gains over deterministic reconstruction.
• Combining contrastive and generative SSL yields the best overall performance across tasks.
• Ablation shows masking and conformer count affect performance, with diminishing returns beyond moderate conformer counts and benefits from nonzero masking.