Predicting Molecular Conformation via Dynamic Graph Score Matching

1 Introduction . Graph-based representations of molecules has become prevalent in a variety of tasks such as property prediction [ 13 , 31 ] and molecule generation [ 17 , 33 , 45 ] . However , a more natural and intrinsic representation of a molecule is its 3D geometry or conformation , which represents a molecule as a set of 3D coordinates . The 3D representation of molecules is central to many tasks , such as molecular properties prediction and virtual screening . Nevertheless , determining the conformation of a molecule remains a challenging task — both computational approaches , e.g. , molecular dynamics ( MD ) [ 9 ] , and experimental approaches , e.g. , crystallography , are expensive and time-consuming . Recently , machine learning approaches have demonstrated promising performance for molecular conformation generation . Pioneering methods such as GRAPHDG [ 36 ] and CGCF [ 43 ] first predict interatomic distances between bonded atoms and then solve 3D coordinates from the predicted distances via a post-processing algorithm . Very recently , Shi et al . proposed the CONFGF [ 34 ] , which employs the score matching technique [ 38 ] to learn pseudo-forces between bonded atoms and iteratively applies the forces to a randomly initialized 3D structure until convergence . CONFGF gets * Equal contribution . Work was done during Shitong ’ s internship at Mila . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) , virtual . rid of the two-stage fashion in prior works and significantly improves the performance . Nevertheless , these approaches have a common major limitation — they mainly focus on modeling the local interactions between bonded atoms defined by the input molecular graphs but fail to capture long-range interactions between non-bonded atoms1 , as they only model distances ( or gradients ) between bonded atoms . While in molecular mechanics , the potential energy of a molecule that alters conformations can be modeled as a sum of four parts [ 22 ] : E = Ebond + Eangle + Etorsion + Enon-bonded , ( 1 ) where Ebond , Eangle , and Etorsion model local interactions between bonded atoms , which are modeled in previous methods [ 34 , 36 , 43 ] . Long-range interactions between non-bonded atoms , denoted as Enon-bonded , are also non-trivial , which shape the molecular geometry via non-negligible electrostatic forces or van der Waals forces , etc . For multi-molecular complexes , non-bonded interactions dominate complexes ’ geometry . An ideal solution to conformation generation should therefore capture both the local and long-range interactions . In Figure 1 , we present three typical molecular systems where long-range interactions play a key role in determining their conformations . To tackle the aforementioned challenge of modeling long-range interactions , in this paper , we propose the Dynamic Graph Score Matching ( DGSM ) for molecular conformation generation , following the principle of CONFGF [ 34 ] that learns the gradients of the logarithm density of atomic coordinates . Instead of relying on the static input molecular graph as existing work , the basic idea is to dynamically construct graph structures between atoms based on their spatial proximity during both training and inference . This allows the model to ( 1 ) dynamically learn molecular graph representations with evolved graph structures that take long-range interactions into consideration , and ( 2 ) dynamically determine a set of interatomic distances that contribute to gradients of the current atomic coordinates . Specifically , the edges in the dynamic graph consist of two parts . The first part of edges are determined by covalent bonds , which capture local interactions between atoms ( Ebond , Eangle and Etorsion ) . The second part of edges are determined dynamically by spatial proximity between atoms at each training or sampling step , i.e. , two atoms are connected as long as they are proximal , no matter whether they are bonded . Such a strategy is able to effectively capture non-local interactions ( Enon-bonded ) since the magnitude of long-range interactions is inversely correlated with distances between atoms [ 30 ] . It remains meanwhile scalable as we avoid connecting all the atom-pairs , which has quadratic complexity . In addition , modeling non-bonded interactions enable the model to sample conformations for multi-molecular complexes , which represent a broader range of problems . We conduct extensive experiments and compare DGSM against previous state-of-the-art methods on both standard conformation generation and property prediction tasks . Numerical results show that DGSM outperforms previous methods by a clear margin , confirming the benefit of modeling long-range interactions . Besides , to further demonstrate the advantage of DGSM , we bring attention to two more challenging tasks — protein sidechain conformation prediction and multi-molecular complex structure prediction . These two new tasks represent two classes of practical challenges : predicting structures for macro-molecules and multi-molecular complexes . 2 Related Work . Prior works on conformation generation mainly rely on molecular dynamics ( MD ) [ 9 ] , where new conformations are sequentially generated based on an initial conformation and a physical model for interatomic potentials [ 25 , 27 ] . Although capable of accurately sampling equilibrium conformations , these methods are computationally intensive , especially for large molecular systems [ 2 , 35 ] , e.g. , 1Note that the receptive field of graph neural networks is much smaller than the diameter of molecular graphs in large systems , e.g. , protein . proteins . Another category of approaches leverage distance geometry [ 8 ] and fix distances between atoms to idealized values heuristically [ 4 ] , which are much faster but less accurate . Recently , a variety of deep generative models have been proposed for molecular conformation generation , which strike a good balance between computational efficiency and accuracy . Among these methods , Mansimov et al . [ 24 ] first propose a variational autoencoder to directly generate 3D atomic coordinates . Albeit simple , this method fails to model the roto-translation equivariance of molecular conformations , leading to unsatisfactory performance . To preserve roto-translation equivariance , Simm and Hernandez-Lobato [ 36 ] and Xu et al . [ 43 ] first model the molecular distance geometry and then reconstruct atomic coordinates from generated distances by solving an optimization problem . The state-of-the-art method CONFGF [ 34 ] estimates pseudo-forces acting on atoms and generates conformations via Langevin MCMC [ 42 ] , which bypasses the two-stage fashion in previous works and enhances the performance significantly . Two concurrent works [ 12 , 44 ] exist which generate conformations in end-to-end fashion via geometry elements assembly and bilevel programming respectively . Recently there has also been attempt to use reinforcement learning for conformation search [ 14 ] . Such a method is incapable of modeling bond lengths explicitly , and is fundamentally different from other approaches . To summarize , all of the previous methods focus mainly on modeling the local interactions based on the static input molecular graphs ( or augmented graphs by adding auxiliary edges between atoms that are two- and three-hops away ) and overlook the long-range non-bonded interactions between atoms . In contrast , our DGSM explicitly models both the local and long-range interactions via dynamic graph score matching and effectively addresses the above issue . 3 Preliminaries . 3.1 Notations and Problem Formulation . Notations . Let G = 〈V , E〉 be a molecular graph , where V = { v1 , v2 , · · · , v|V| } is the set of nodes representing atoms , and E = { eij | ( i , j ) ⊆ V × V } is the set of edges representing inter-atomic bonds in the molecule . Each node vi ∈ V is labeled with atomic attributes , e.g. , the element type Zi and the atomic coordinate ri ∈ R3 . Each edge eij ∈ E is labeled with a bond type . The conformation of the molecular graph G can be represented as a matrixR ∈ R|V|×3 . The distances between all pairs of atoms can be represented as a matrixD ∈ R|V|×|V| , whereDij : = dij = ‖ri − rj‖2 denotes the Euclidean distance between the positions of vi and vj . Following previous work [ 34 , 36 , 43 ] , we expand the original molecular graph by adding auxiliary edges between atoms that are second and third neighbors in G to reduce the degrees of freedom in 3D coordinates . Problem Formulation . Given a molecular graph G = 〈V , E〉 , the task of molecular conformation generation is the conditional generation of conformationsR = [ r1 ; r2 , · · · ; r|V| ] ∈ R|V|×3 based on G , while being able to capture long-range interactions between non-bonded atoms . Note that G may has multiple connected components , e.g. , protein-ligand complexes and multi-molecular complexes . 3.2 Score-Based Generative Modeling . Score-based generative modeling [ 15 , 38–40 ] is a class of generative models that has recently proven effective in a variety of tasks , ranging from image generation [ 15 , 40 ] , audio synthesis [ 7 , 21 ] to shape generation [ 6 , 23 ] . For any continuously differentiable probability density p ( x ) , we define its score function s ( x ) as∇x log p ( x ) , i.e. , the direction where the logarithm data density grows most rapidly . Score-based generative modeling perturbs the data with different levels of Gaussian noise and jointly estimates the score function of p ( x ) using neural networks . Samples are then generated by sampling from a sequence of decreasing noise levels with Langevin dynamics [ 42 ] . Formally , given a data distribution pdata ( x ) , let { σi } Li=1 be a sequence of noise levels that satisfies σ1 > σ2 > · · · > σL and σi/σi−1 = γ . Consider a series of noise distributions pσi ( x̃ | x ) : = N ( x̃ ; x , σ2i I ) , and denote the corresponding perturbed data distribution as pσi ( x̃ ) : = ∫ pσi ( x̃ | x ) pdata ( x ) dx . Song and Ermon [ 38 ] propose to jointly approximate the score function of each noise level , denoted by sθ ( x , σi ) , with the following objective : θ∗ = argminθ 1 2L L∑ i=1 σ2i Epdata ( x ) Epσi ( x̃|x ) [ ∥∥∥sθ ( x̃ , σi ) −∇x̃ log pσi ( x̃ | x ) ∥∥∥ 2 2 ] . ( 2 ) Assuming sufficient data and model capacity , the optimal noise conditional score network sθ∗ ( x , σi ) matches∇x log pσi ( x ) almost everywhere [ 38 ] . After training score networks , Song and Ermon [ 38 ] run annealed Langevin dynamics for each pσi ( x ) sequentially , where samples from each noise level serve as initializations for Langevin dynamics of the next noise level . Given σL small enough , the final samples from pσL ( x ) approximate to samples from pdata under minor conditions . 4 Model . Our approach treats the conformation generation as sequentially moving atoms towards high-density regions guided by pseudo-forces , i.e. , gradients of atoms . Following Shi et al . [ 34 ] , we leverage the denoising score matching [ 38 , 41 ] to approximate the gradients of the logarithm density of atomic coordinates , denoted as ∇R log p ( R | G ) . To model atomic gradients that are sensitive to both the local and long-range interactions ( Eq . 1 ) and inspired by the fact that long-range interactions decrease rapidly as distances increase , we propose to dynamically construct graph structures with non-bonded edges between atom pairs within a distance based on the current spatial proximity . In this way , we enable the model to effectively capture long-range non-bonded interactions while avoid connecting all atoms , which is computationally expensive . To ensure that the distribution of graph structures during training matches with that during generation , we devise a dynamic graph score matching algorithm , where graph structures are also dynamically determined during training depending on added perturbations . The whole framework is illustrated in Figure 2 and Figure 3 . Below we describe the framework of score estimation for Cartesian coordinates in Section 4.1 , the dynamic graph score matching algorithm in Section 4.2 , and the generation procedure in Section 4.3 .

Summary: This paper addresses the task of molecular conformer generation, i.e. generate the 3D conformer ensemble conditioned on the molecular graph. The introduced method, Dynamic Graph Score Matching (DGSM), is using score matching to generate 3D structures by explicitly capturing both local and long-range interactions. This is done using dynamic graphs constructed via atom coordinate perturbations using different noise levels as required by the denoising score matching procedure.  The score is estimated using message passing neural networks which operate on these dynamic graphs. At inference time, Langevin dynamics is used to sample conformations from the generated distribution. Empirically, the authors show clear improvements over previous machine learning (ML) models and over one popular open-source non-ML method (RDKit) on molecular conformer generation on two different tasks, as well as improvements on experiments on property prediction based on molecular conformations, protein sidechain conformation generation, and multi-molecular complex conformation.  In short, I think this approach is quite compelling and would be very impactful for future work related to graph conditional 3D generative models of various data in biology, chemistry, etc. The paper is very well written and the method is clean, addressing nicely using dynamic graphs the very important problem of modeling long-range interactions in 3D generation.