[Paper Review] 3D Infomax improves GNNs for Molecular Property Prediction

H. Stärk, Dominique Beaini+5 more|arXiv (Cornell University)|Oct 8, 2021

Machine Learning in Materials Science53 references85 citations

TL;DR

The paper introduces 3D Infomax, a self-supervised pre-training method that teaches 2D GNNs to implicitly reason about 3D molecular geometry, improving downstream quantum property predictions and enabling cross-dataset generalization.

ABSTRACT

Molecular property prediction is one of the fastest-growing applications of deep learning with critical real-world impacts. Including 3D molecular structure as input to learned models improves their performance for many molecular tasks. However, this information is infeasible to compute at the scale required by several real-world applications. We propose pre-training a model to reason about the geometry of molecules given only their 2D molecular graphs. Using methods from self-supervised learning, we maximize the mutual information between 3D summary vectors and the representations of a Graph Neural Network (GNN) such that they contain latent 3D information. During fine-tuning on molecules with unknown geometry, the GNN still generates implicit 3D information and can use it to improve downstream tasks. We show that 3D pre-training provides significant improvements for a wide range of properties, such as a 22% average MAE reduction on eight quantum mechanical properties. Moreover, the learned representations can be effectively transferred between datasets in different molecular spaces.

Motivation & Objective

• Motivate molecular property prediction on large scales where explicit 3D structures are unavailable.
• Develop a pre-training scheme that teaches 2D GNNs to encode latent 3D information from 2D graphs.
• Demonstrate that the learned 3D-aware representations improve downstream quantum property predictions.
• Show that 3D Infomax representations generalize across distinct molecular spaces and datasets.
• Explore the benefit of using multiple conformers during pre-training to further boost performance.

Proposed method

• Pre-train a 2D GNN f^a on 2D molecular graphs to produce z^a by maximizing mutual information with a 3D representation z^b derived from conformers.
• Use a contrastive NT-Xent style loss to align z^a and z^b for same molecules and repel non-matching pairs.
• Incorporate multiple conformers (c) by treating all conformers of a molecule as positive examples when computing the contrastive loss (L^multi3D).
• The 3D network processes 3D coordinates as a point cloud and encodes pairwise distances into an SE(3)-invariant z^b via a distance-based MPNN with sinusoidal feature mappings γ(d_uv).
• Fine-tune the 2D network on downstream tasks using only 2D inputs, leveraging the implicitly learned 3D information for improved predictions.

Experimental results

Research questions

• RQ1 Can a 2D GNN learn to encode latent 3D information from 2D molecular graphs through contrastive pre-training with 3D conformers?
• RQ2 Do 3D Infomax pre-trained models improve quantum property predictions (e.g., QM9) compared to 2D-only pre-training baselines?
• RQ3 Are the learned representations transferable across molecular spaces with different sizes and compositions?
• RQ4 Does leveraging multiple conformers during pre-training yield additional gains in downstream properties?

Key findings

• 3D Infomax achieves large improvements in quantum property predictions, including a ~22% average MAE reduction on QM9 properties.
• Pre-trained representations generalize across distinct molecular spaces, showing gains when transferring from small to drug-like molecules and vice versa.
• Using multiple conformers during pre-training provides additional performance gains, with diminishing returns beyond three conformers.
• 3D Infomax outperforms conventional SSL baselines like GraphCL and distance-prediction pre-training across evaluated datasets.
• The method does not exhibit negative transfer, unlike some prior pre-training approaches.