[Paper Review] Learning Neural Generative Dynamics for Molecular Conformation Generation

Minkai Xu, Shitong Luo+3 more|arXiv (Cornell University)|Feb 20, 2021

Machine Learning in Materials Science41 references28 citations

TL;DR

The paper introduces a probabilistic framework that combines conditional graph continuous flow (CGCF) for distance geometry with an energy-based tilting model (ETM) to generate diverse and valid molecular conformations conditioned on a molecular graph, achieving state-of-the-art results on GEOM-QM9, GEOM-Drugs, and ISO17 benchmarks.

ABSTRACT

We study how to generate molecule conformations (i.e., 3D structures) from a molecular graph. Traditional methods, such as molecular dynamics, sample conformations via computationally expensive simulations. Recently, machine learning methods have shown great potential by training on a large collection of conformation data. Challenges arise from the limited model capacity for capturing complex distributions of conformations and the difficulty in modeling long-range dependencies between atoms. Inspired by the recent progress in deep generative models, in this paper, we propose a novel probabilistic framework to generate valid and diverse conformations given a molecular graph. We propose a method combining the advantages of both flow-based and energy-based models, enjoying: (1) a high model capacity to estimate the multimodal conformation distribution; (2) explicitly capturing the complex long-range dependencies between atoms in the observation space. Extensive experiments demonstrate the superior performance of the proposed method on several benchmarks, including conformation generation and distance modeling tasks, with a significant improvement over existing generative models for molecular conformation sampling.

Motivation & Objective

• Motivate and address the challenge of generating multi-modal, long-range dependent molecular conformations from graphs.
• Model the distance geometry p(d|G) with a continuous normalizing flow conditioned on the molecular graph.
• Refine generated conformations via an energy-based tilting term E(R,G) to capture long-range interactions.
• Enable efficient sampling and training by combining flow-based density estimation with energy-based refinement.

Proposed method

• Factorize p(R|G) as p(R|d,G) p(d|G) to separate coordinate generation from distance geometry.
• Model p(d|G) with a Conditional Graph Continuous Flow (CGCF) implemented via continuous dynamics (CNF) and Message Passing Neural Networks.
• Represent p(R|d,G) through a probabilistic form that allows sampling via a simple energy-based distribution over Cartesian coordinates.
• Introduce an Energy-based Tilting Model (ETM) Eφ(R,G) to model long-range interactions and tilt p(R|G) via pθ,φ(R|G) ∝ pθ(R|G) exp(-Eφ(R,G)).
• Train ETM using Noise-Contrastive Estimation (NCE) leveraging CGCF as the noise distribution.
• Perform two-stage sampling: (i) generate distances with CGCF and obtain initial R via EDG-like postprocessing, (ii) refine R with Langevin dynamics using the ETM.

Experimental results

Research questions

• RQ1 Can a flow-based model over distance geometry capture the multimodal distribution of inter-atomic distances conditioned on the molecular graph?
• RQ2 Does incorporating an energy-based tilting term improve long-range dependency modeling and the realism of generated conformations?
• RQ3 How does the two-stage CGCF + ETM sampling strategy compare to prior baselines on conformation generation and distance distribution tasks?
• RQ4 What is the impact of ETM on diversity and convergence to equilibrium-like conformations across standard benchmarks?

Key findings

• CGCF achieves higher diversity and accuracy than state-of-the-art baselines on conformation generation benchmarks.
• Incorporating ETM further improves the quality of generated conformations by explicitly modeling long-range interactions.
• The combined CGCF + ETM framework outperforms baselines on both conformation generation and distance distribution tasks across GEOM-QM9, GEOM-Drugs, and ISO17 datasets.
• The method maintains high diversity (low mode collapse) while achieving strong alignment with ground-truth distributions.