Unified 2D and 3D Pre-Training of Molecular Representations

TL;DR

This paper introduces a unified pre-training approach that combines 2D graph and 3D conformation information for molecular representations, leading to improved performance on property prediction and conformation generation tasks.

Contribution

It proposes a novel joint 2D and 3D pre-training method that encodes and fuses both types of molecular information using graph neural networks.

Findings

Achieves state-of-the-art results on 10 out of 11 molecular property prediction tasks.

Improves average performance on 2D-only tasks by 8.3%.

Significantly enhances 3D conformation generation results.

Abstract

Molecular representation learning has attracted much attention recently. A molecule can be viewed as a 2D graph with nodes/atoms connected by edges/bonds, and can also be represented by a 3D conformation with 3-dimensional coordinates of all atoms. We note that most previous work handles 2D and 3D information separately, while jointly leveraging these two sources may foster a more informative representation. In this work, we explore this appealing idea and propose a new representation learning method based on a unified 2D and 3D pre-training. Atom coordinates and interatomic distances are encoded and then fused with atomic representations through graph neural networks. The model is pre-trained on three tasks: reconstruction of masked atoms and coordinates, 3D conformation generation conditioned on 2D graph, and 2D graph generation conditioned on 3D conformation. We evaluate our method on 11 downstream molecular property prediction tasks: 7 with 2D information only and 4 with both 2D and 3D information. Our method achieves state-of-the-art results on 10 tasks, and the average improvement on 2D-only tasks is 8.3%. Our method also achieves significant improvement on two 3D conformation generation tasks.