Pre-training Molecular Graph Representation with 3D Geometry

TL;DR

GraphMVP introduces a novel pre-training framework that leverages 3D molecular geometry to enhance 2D graph representations via two self-supervised tasks: a contrastive objective aligning 2D-3D views and a generative objective (VRR) that reconstructs across representation space. It formalizes a multi-task objective and provides theoretical justification based on mutual information and the privileged-information view of 3D geometry. Empirical results on diverse molecular tasks show consistent improvements over prior SSL methods and state-of-the-art performance. The work suggests that incorporating 3D priors during pre-training can significantly boost data-efficient molecular property prediction and related tasks.

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.

Figures (9)

• Figure 1: Overview of the pre-training stage in GraphMVP. The black dashed circles denote subgraph masking, and we mask the same region in the 2D and 3D graphs. Multiple views of the molecules (herein: Halicin) are mapped to the representation space via 2D and 3D GNN models, where we conduct GraphMVP for SSL pre-training, using both contrastive and generative pretext tasks.
• Figure 2: We select the molecules whose properties can be easily resolved via 3D but not 2D. The randomly initialised 2D GNN achieves accuracy of $38.9\pm0.8$ and $77.9\pm1.1$, respectively. The GraphMVP pre-trained ones obtain scores of $42.3\pm1.3$ and $81.5\pm0.4$, outperforming all the precedents in \ref{['sec:main_results']}. We plot cases where random initialization fails but GraphMVP is correct.
• Figure 3: Venn diagram of mutual information. Inspired by wikipedia.
• Figure 4: Contrastive SSL in GraphMVP. The black dashed circles represent subgraph masking.
• Figure 5: VRR SSL in GraphMVP. The black dashed circles represent subgraph masking.