[Paper Review] GeoMol: Torsional Geometric Generation of Molecular 3D Conformer Ensembles

Octavian-Eugen Ganea, Lagnajit Pattanaik+5 more|arXiv (Cornell University)|Jun 8, 2021

Machine Learning in Materials Science51 references47 citations

TL;DR

GeoMol is an end-to-end SE(3)-invariant model that generates diverse low-energy 3D conformer ensembles from molecular graphs by predicting local structures and torsion angles, with an optimal-transport based loss for diversity.

ABSTRACT

Prediction of a molecule's 3D conformer ensemble from the molecular graph holds a key role in areas of cheminformatics and drug discovery. Existing generative models have several drawbacks including lack of modeling important molecular geometry elements (e.g. torsion angles), separate optimization stages prone to error accumulation, and the need for structure fine-tuning based on approximate classical force-fields or computationally expensive methods such as metadynamics with approximate quantum mechanics calculations at each geometry. We propose GeoMol--an end-to-end, non-autoregressive and SE(3)-invariant machine learning approach to generate distributions of low-energy molecular 3D conformers. Leveraging the power of message passing neural networks (MPNNs) to capture local and global graph information, we predict local atomic 3D structures and torsion angles, avoiding unnecessary over-parameterization of the geometric degrees of freedom (e.g. one angle per non-terminal bond). Such local predictions suffice both for the training loss computation, as well as for the full deterministic conformer assembly (at test time). We devise a non-adversarial optimal transport based loss function to promote diverse conformer generation. GeoMol predominantly outperforms popular open-source, commercial, or state-of-the-art machine learning (ML) models, while achieving significant speed-ups. We expect such differentiable 3D structure generators to significantly impact molecular modeling and related applications.

Motivation & Objective

• Predict low-energy 3D conformer ensembles of small molecules from molecular graphs.
• Model essential geometry elements (local 3D structures and torsion angles) without over-parameterizing degrees of freedom.
• Enable end-to-end, non-autoregressive generation with SE(3) invariance.
• Promote diverse conformer ensembles through an optimal transport loss without adversarial training.
• Achieve fast conformer generation without relying on expensive force-field fine-tuning.

Proposed method

• Represent molecules as graphs and compute atom embeddings with MPNNs.
• Predict local 3D structures (LS) for non-terminal atoms, including bond distances and angles, in a permutation-equivariant way.
• Assemble neighboring LSs by predicting a single torsion angle per bond in a canonical SE(3)-invariant formulation.
• Deterministically assemble full conformers at test time with ring corrections and chiral center handling via oriented volumes.
• Train with a non-adversarial optimal transport based loss that matches generated to ground-truth conformers for diversity.
• Use an OT-based ensemble loss to handle variable numbers of ground-truth conformers and promote coverage of low-energy states.

Experimental results

Research questions

• RQ1 Can a non-autoregressive, SE(3)-invariant ML model generate representative low-energy 3D conformer ensembles from molecular graphs?
• RQ2 How can local geometric predictions (torsions and local structures) be combined to deterministically assemble full 3D conformers?
• RQ3 Can an optimal transport loss promote diverse conformer generation without adversarial training?
• RQ4 What is the impact of explicit chirality handling and ring corrections on conformer accuracy and diversity?

Key findings

• GeoMol often outperforms popular ML and open-source/commercial baselines on GEOM-QM9 and GEOM-DRUGS benchmarks.
• GeoMol achieves competitive quality without expensive force-field fine-tuning.
• GeoMol generates conformers in seconds for drug-like molecules, substantially faster than baselines like ETKDG/RDKit.
• An OT-based loss enables diverse conformer ensembles by matching generated to ground-truth conformers without adversarial training.
• Explicit chirality handling via oriented volumes improves enantiomer accuracy.
• Single torsion angle per bond and local-structure predictions reduce over-parameterization while preserving reconstruction fidelity.