PaperlayerView exact revision on arXiv

Source-linked AI summary

3D Infomax improves GNNs for Molecular Property Prediction

Hannes Stärk, Dominique Beaini, Gabriele Corso, Prudencio Tossou, Christian Dallago, Stephan Günnemann, Pietro Liò

arXiv:2110.04126v4cs.LGcs.AIq-bio.BM

TL;DR

Molecular property prediction benefits from 3D structure, but explicitly computing molecular geometry is too expensive at large scale. 3D Infomax pre-trains a GNN to encode latent 3D information from 2D molecular graphs, improving predictions and transferring across different molecular spaces.

  • Problem

    Explicit molecular geometry improves property prediction, but computing it is computationally intractable for many large-scale applications.

  • Method

    3D Infomax maximizes mutual information between 2D-graph GNN embeddings and learned 3D molecular representations, then transfers the GNN for fine-tuning without 3D inputs.

  • Results

    3D Infomax consistently improves quantum property predictions by about 22% and transfers across highly different molecular types.

  • Takeaways & Limitations

    The method supplies latent 3D and quantum information while retaining the inference speed of a standard GNN operating on 2D molecular graphs.

  • Takeaways & Limitations

    For conformer-specific QM9 properties, implicit 3D information can remain less accurate than explicit ground-truth conformers.

Abstract

from arXiv · show

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.

1. Introduction

3D Infomax pre-trains GNNs to encode geometric information from 2D molecular graphs, addressing the computational cost of explicitly obtaining molecular geometry. The learned representations improve property prediction and transfer across molecular spaces.

  • Motivation: GNNs capture covalent connectivity but reason poorly about interatomic forces dependent on relative spatial positions.Prior work shows that incorporating 3D positions improves molecular property prediction.
  • Motivation: Explicitly computing molecular geometry with simulations or recent conformation-generation methods is too costly for many large-scale applications.This computational barrier motivates learning implicit 3D information from 2D graphs.
  • 3D Infomax: 3D Infomax maximizes mutual information between a 2D-graph GNN embedding and a learned representation of the molecule’s 3D graph.After pre-training, the GNN can generate implicit 3D information for molecules whose geometry is unavailable during fine-tuning.
  • Evidence: 3D Infomax improves property predictions by large margins and generalizes across substantially different molecular spaces without observed negative transfer.The evaluations use multiple 3D pre-training datasets and quantum mechanical downstream properties.
  • Evidence: The method uses multiple conformers of the same molecule to further improve downstream property predictions.The paper also evaluates how effectively this multi-conformer information can be leveraged.

2. Background

Molecular graphs encode atoms and covalent bonds, while molecular conformers represent alternative spatial arrangements that can have different chemical properties. Explicitly generating accurate conformers remains too slow for applications processing very large molecular collections.

  • 2D Molecular Graphs: A 2D molecular graph G = (V, E) represents atoms as nodes and covalent bonds as edges, without 3D coordinates.Node features can include atomic numbers, while edge features can include bond types.
  • 3D Molecular Conformers: A molecule’s single 2D graph can correspond to multiple low-energy spatial arrangements called conformers.Different conformers may exhibit different chemical properties, so capturing likely conformations matters for 3D modeling.
  • 3D Molecular Conformers: Known conformers are represented as point clouds whose points specify the locations of all atoms.For c conformers, the representation is the set {R_j} over the conformers.
  • Conformer Generation: Conformer-generation methods trade speed against accuracy, with CREST requiring about 6 hours per drug-like molecule per CPU-core.This makes precise explicit structure generation infeasible for virtual screening datasets containing millions or billions of molecules.
  • Symmetries of Molecules: Molecular conformations are invariant to joint translations and rotations, known as SE(3) symmetry.Chiral molecules are not invariant to reflections because their properties depend on chirality.
  • Graph Neural Networks: MPNN-style GNNs iteratively update node representations from neighboring nodes and edges, then combine them with permutation-invariant readouts.These models learn graph representations through message passing over molecular connectivity.

3. Related Work

Related work established GNN-based molecular prediction, 3D-aware architectures, and self-supervised learning for molecular graphs. 3D Infomax extends this direction by using 3D structures for contrastive pre-training and evaluating transfer across chemical spaces and quantum tasks.

  • Molecular Property Prediction: GNNs are widely used for quantum chemistry, drug discovery, and general molecular property prediction.Accessible molecular datasets have supported progress in this established field.
  • 3D Molecular Models: Adding geometric information such as bond lengths, angles, and torsion angles improves on 2D-only molecular prediction.SchNet, DimeNet, SMP, and GemNet progressively capture richer aspects of molecular geometry.
  • Self-Supervised Learning: Self-supervised learning obtains representation-learning signals from unlabeled data by comparing embeddings of similar and dissimilar inputs.Molecular SSL is motivated partly by the small size of chemistry datasets caused by experimental and computational costs.
  • Self-Supervised Learning: Existing molecular SSL methods provide limited improvements and often fail to generalize.Earlier methods primarily leveraged 2D molecular information.
  • 3D Pre-training: 3D Infomax directly incorporates multiple conformers in a contrastive loss and pre-trains across multiple chemical spaces.Its evaluation includes quantum mechanical tasks and examines transfer between different molecular spaces.

4. 3D Infomax

3D Infomax pre-trains a 2D GNN to encode information about molecular geometry from 2D graphs by aligning its representations with those produced from 3D conformers. Contrastive learning, extended to multiple conformers, transfers implicit 3D knowledge to downstream prediction.

  • 4. 3D Infomax: After pre-training, the 2D network’s weights are transferred to property-prediction tasks where only 2D molecular graphs are available.The resulting GNN can use learned implicit 3D information during fine-tuning.
  • 4. 3D Infomax: 3D Infomax maximizes mutual information between a 2D graph representation and a learned representation of associated 3D conformers.The 2D network receives molecular graphs, while a separate 3D network encodes conformer coordinates.
  • 4.1. Contrastive Framework: The 2D network is pre-trained with contrastive learning so same-molecule 2D and 3D representations are similar and different-molecule representations are dissimilar.The loss uses positive pairs from the same molecule and negative pairs from other molecules in the batch.
  • 4.2. Using Multiple Conformers: With multiple conformers, the loss treats every conformer of the same molecule as a positive pair and aggregates their similarities.The method uses the c highest-probability conformers, repeating the lowest-energy conformer when fewer than c are available.
  • 4.3. 3D Network: The 3D network produces an SE(3)-invariant representation from molecular coordinates without access to 2D atom or bond features.Its architecture encodes pairwise Euclidean distances and processes them on a fully connected molecular graph.

5. Experiments

The experiments compare 3D Infomax with direct 3D-prediction, conformer-generation, and conventional GraphCL pre-training while fine-tuning without 3D inputs. Results show substantial improvements, cross-dataset transfer, benefits from several conformers, and diminishing returns from larger pre-training sets.

  • 5. Experiments: 3D Infomax is compared with distance prediction, conformer generation, and GraphCL as alternative GNN pre-training approaches.The evaluation fine-tunes models on property-prediction datasets while ignoring available 3D information.
  • 5.3. Quantum Mechanical Properties: 22% average MAE reduction is obtained on QM9 after 3D pre-training on one half of QM9.3D Infomax improves over random initialization and GraphCL across the reported QM9 properties.
  • 5.3. Quantum Mechanical Properties: 19% and 18% improvements are obtained when pre-training on GEOM-Drugs and QMugs, respectively, before fine-tuning on QM9.These gains occur despite the pre-training datasets containing substantially larger drug-like molecules than QM9.
  • 5.3. Quantum Mechanical Properties: For some properties, 3D Infomax remains less accurate than 3D GNNs using explicit ground-truth conformers.The authors note that conformer-specific properties may impose a maximum accuracy when only the molecule, not its target conformer, is known.
  • 5.3. Quantum Mechanical Properties: 3D Infomax does not show the negative transfer observed sometimes with conventional pre-training and generalizes across substantially different molecular spaces.One evaluation uses pre-training molecules containing only C, H, N, O, and F while target molecules contain eleven additional elements.
  • 5.4. Number of Conformers and pre-training Molecules: 12 hours of pre-training on 620k QMugs molecules with three conformers compares with 71 hours for GraphCL on 280k GEOM-Drugs molecules.Performance improves as the pre-training dataset grows, but the returns diminish.
  • 5.4. Number of Conformers and pre-training Molecules: Using more than one conformer improves performance, but adding higher-energy conformers beyond about three can worsen downstream results.Three conformers cover 70% of GEOM-Drugs’ cumulative Boltzmann weight on average; QM9 homo MAE improves from 82.57 with one conformer to 70.77 with three.

6. Conclusion

3D Infomax pre-trains GNNs to extract latent 3D and quantum information from 2D molecular graphs, improving property prediction while retaining standard 2D-GNN inference speed. The learned representations transfer across substantially different molecule types, and multiple conformers provide additional downstream information.

  • 3D Infomax teaches a GNN to produce latent 3D and quantum information from 2D molecular graphs.
  • ∼22% improvements were found for quantum properties, overshadowing gains from conventional self-supervised learning methods.
  • Representations transfer from molecules averaging 18 atoms to drug-like molecules averaging 44.4 atoms.
  • Fine-tuning uses the latent information while retaining the inference speed of a standard GNN operating on 2D molecular graphs.
  • Using multiple molecular conformers during pre-training provides additional information that further improves downstream property predictions.

B.2. Data Details

The study pre-trains on three 3D molecular datasets spanning small, drug-like, and conformer-rich molecules, then evaluates transfer across quantum, physical-chemical, biophysical, and physiological tasks.

  • Three 3D datasets are used for pre-training, varying in molecule size and dataset scale.
  • QM93 contains 134k small organic molecules with one low-energy conformer and 12 quantum mechanical regression targets.
  • GEOM-Drugs contains 304k biologically and pharmacologically relevant molecules with multiple 3D conformers and energy targets.
  • QMugs contains 665k drug-like molecules with three diverse conformers each and conformer-specific quantum mechanical properties.
  • Fine-tuning covers quantum mechanics, physical chemistry, biophysics, and physiology using datasets with graph-only annotations.
  • The omitted OGB datasets muv and pcba were excluded because of their larger size.

C.1. Non-Quantum Properties

For non-quantum properties, 3D Infomax helps selectively: it improves four OGB datasets, with stronger gains for intrinsic molecular properties than abstract clinical outcomes, and does not reduce performance.

  • 3D Infomax improved performance for 4 out of 10 OGB datasets.
  • 3D Infomax never decreases performance on the evaluated OGB datasets, although GraphCL and ConfGen are sometimes superior.
  • The most significant improvements occur for solubility and lipophilicity tasks such as ESOL and Lipo.
  • Abstract tasks such as clinical test outcome prediction benefit less than tasks tied more directly to molecular mechanics and intrinsic properties.
  • ConfGen often produces significant improvements for intrinsic-property tasks, providing additional evidence for the value of 3D pre-training.
  • For BACE binding prediction, improvements are only modest despite the expected value of 3D information.

C.2. Different 3D Networks and Ablation

Net3D outperforms the alternative 3D networks in the reported QM9 comparison, and its distance encoding is a valuable component. The architecture’s quadratic molecule-size complexity remains a practical trade-off.

  • The γ function maps pairwise distances to a higher-dimensional space and is a valuable component of Net3D.
  • Net3D achieves the best MAE, while removing the γ distance encoding defines the Net3D w/o γ ablation.
  • Net3D yields the best downstream performance among the compared 3D networks.
  • SMP has the worst downstream performance because its 3D input does not uniquely define all relative molecular positions.
  • SMP distinguishes chiral molecules and uses less memory, whereas Net3D and EGNN have quadratic complexity in molecule size.

C.3. Different Methods for Multiple Conformers

The section compares strategies for incorporating multiple molecular conformers into 3D pre-training. Multi-conformer information can substantially improve predictions, with the best strategy depending on pre-training data and computational constraints.

  • multi3D: Multi3D represents each molecule with multiple conformers and encourages one 2D embedding to match all corresponding 3D representations.The 3D network processes all selected conformers, which serve as additional positive samples in contrastive learning.
  • multi3D+2D: Multi3D+2D produces one 2D representation per conformer and compares these representations with all corresponding 3D representations.The approach uses a set-based similarity function rather than a single cosine similarity.
  • Results: Large improvements from multiple conformers were observed after QMugs pre-training, reducing homo MAE from 82.57 to 70.77 and Gibbs free-energy MAE from .1966 to .1831 for GEOM-Drugs.The improvements were larger than after GEOM-Drugs pre-training, where the lowest-energy conformer already captured much relevant information.
  • Results: Multi3D performed best after QMugs pre-training, while uniform conformer sampling was simpler and used substantially less memory.Multi3D processes conformers in parallel; more than five conformers with batch size 500 did not fit on a 48GB vRAM GPU.
  • Results: The hypothesis that larger molecules benefit more from multiple conformers was not supported, because improvements on QM9 were larger than on GEOM-Drugs.QM9 contains smaller molecules, whereas GEOM-Drugs contains larger drug-like molecules.
  • Different losses: The mutual-information loss and InfoNCE outperformed Donsker–Varadhan and Jensen–Shannon estimators across the evaluated settings.The mutual-information loss sometimes exceeded InfoNCE, although this may reflect additional temperature-parameter search.

C.5. SSL Methods

This section compares 3D Infomax with other self-supervised latent-space methods. 3D Infomax achieved the strongest result in the reported comparison.

  • Methods: The comparison included Barlow Twins, multimodal BYOL, and VICReg alongside 3D Infomax.Method-specific hyperparameters were selected before downstream tuning.
  • Results: 3D Infomax reduced QM9 homo MAE from 82.10 ± 0.33 to 68.96 ± 0.32, whereas BYOL reached 79.16 ± 0.58.The other compared methods produced no improvement or a smaller reduction.
  • Interpretation: The authors attribute 3D Infomax’s advantage to its direct optimization of a lower bound on mutual information.The alternative methods maximize similarity with mechanisms intended to prevent collapse.

C.6. Pre-training a 3D GNN

The section tests whether 3D Infomax pre-training can initialize a 3D GNN and examines performance with generated conformers. Pre-training improved the 3D GNN, but GeoMol conformers yielded poor downstream performance.

  • Pre-training a 3D GNN: Pre-training improved SMP’s performance when the 3D GNN was fine-tuned with accurate QM9 conformers.The pre-trained model lacked atom features but could use covalent-bond and other 2D edge information available during pre-training.
  • Future direction: Using pre-training for 3D GNNs is identified as a future direction for predicting quantum properties with accurate 3D information.The section does not report that this direction already surpasses state-of-the-art methods.
  • Generated conformers: SMP performed poorly with GeoMol-generated conformers, while RDKit ETKDG conformers consistently performed better.The authors hypothesize that unrealistic GeoMol outliers caused high MAE because SMP could not handle them well.
  • Batch size: A batch size of 500 was necessary for the good performance of 3D Infomax.The evaluation was motivated by contrastive learning’s reliance on many negative samples.

C.9. Combining Pre-Training Methods

The section evaluates combining 3D Infomax with conventional node-dropping pre-training on 2D or 3D inputs. Performance is reported across several biophysical-property OGB datasets using RMSE or ROC-AUC.

  • Method: The combined method applies GraphCL-style node dropping to the 2D graph and corresponding 3D information during 3D pre-training.The drop ratio is 0.2, and pairwise distances for removed atoms are also removed.
  • Evaluation: Table 18 compares combined 3D Infomax against random initialization and GraphCL across biophysical-property OGB datasets.The table reports RMSE, where lower is better, or ROC-AUC, where higher is better.
Loading 2110.04126v4…