[Paper Review] Learning Gradient Fields for Molecular Conformation Generation

Chence Shi, Shitong Luo+2 more|arXiv (Cornell University)|May 9, 2021

Machine Learning in Materials Science50 references40 citations

TL;DR

ConfGF directly learns gradient fields of the log density over atomic coordinates by connecting to interatomic distances, enabling one-stage conformation generation via Langevin dynamics and achieving state-of-the-art results on GEOM benchmarks.

ABSTRACT

We study a fundamental problem in computational chemistry known as molecular conformation generation, trying to predict stable 3D structures from 2D molecular graphs. Existing machine learning approaches usually first predict distances between atoms and then generate a 3D structure satisfying the distances, where noise in predicted distances may induce extra errors during 3D coordinate generation. Inspired by the traditional force field methods for molecular dynamics simulation, in this paper, we propose a novel approach called ConfGF by directly estimating the gradient fields of the log density of atomic coordinates. The estimated gradient fields allow directly generating stable conformations via Langevin dynamics. However, the problem is very challenging as the gradient fields are roto-translation equivariant. We notice that estimating the gradient fields of atomic coordinates can be translated to estimating the gradient fields of interatomic distances, and hence develop a novel algorithm based on recent score-based generative models to effectively estimate these gradients. Experimental results across multiple tasks show that ConfGF outperforms previous state-of-the-art baselines by a significant margin.

Motivation & Objective

• Motivate a one-stage approach to molecular conformation generation that avoids distance-to-geometry post-processing.
• Propose learning roto-translation equivariant gradient fields of the log density w.r.t. atomic coordinates via distances as intermediate variables.
• Develop a scalable score-based framework that trains on equilibrium conformations to generate diverse, realistic geometries.
• Ensure the method respects physical symmetries and outperforms distance-based two-stage methods on standard benchmarks.

Proposed method

• Parameterize log p_theta(R|G) as f_G(d) with d = interatomic distances and GNNs estimating an energy-like quantity from distances.
• Train a noise-conditional score network s_theta(d, sigma) to model grad_d log q_sigma(d|G) for perturbed distances.
• Compute atomic-coordinate scores via the chain rule: s_theta(R, sigma)_i = sum_{j in N(i)} (1/d_ij) * s_theta(d, sigma)_ij * (r_i - r_j).
• Use roto-translation invariant distance scores to ensure the resulting coordinate scores are roto-translation equivariant.
• Train with denoising score matching across multiple noise levels; use lambda(sigma_i) = sigma_i^2 as in score-based models.
• Generate conformations by annealed Langevin dynamics across noise levels, updating coordinates with the learned score and Gaussian noise.

Experimental results

Research questions

• RQ1 Can molecular conformations be generated in a single stage by directly modeling gradient fields over coordinates?
• RQ2 How can one estimate gradient fields while preserving roto-translation equivariance via interatomic distances as an intermediate representation?
• RQ3 Do distance-based score networks yield higher-quality and more diverse conformations than two-stage distance-to-geometry methods?
• RQ4 How does ConfGF perform in conformation generation, distance distributions, and ensemble property prediction compared to state-of-the-art baselines?

Key findings

• ConfGF achieves state-of-the-art performance on conformation generation metrics (COV and MAT) across GEOM-QM9 and GEOM-Drugs benchmarks.
• ConfGF outperforms baselines on distributions over distances (mean and median MMD) across Single, Pair, and All categories.
• ConfGF yields substantially better MAE for ensemble properties (average energy, minimum energy, HOMO-LUMO gaps) than neural baselines and outperforms RDKit on several metrics.
• Two-stage distance-to-geometry methods (GraphDG, CGCF) incur additional errors not present in the one-stage ConfGF approach.
• RDKit remains competitive on some metrics but generally underperforms compared with neural models when not using post-processing force fields.