GeoDiff: A Geometric Diffusion Model for Molecular Conformation Generation

1 INTRODUCTION . Graph representation learning has achieved huge success for molecule modeling in various tasks ranging from property prediction ( Gilmer et al. , 2017 ; Duvenaud et al. , 2015 ) to molecule generation ( Jin et al. , 2018 ; Shi et al. , 2020 ) , where typically a molecule is represented as an atom-bond graph . Despite its effectiveness in various applications , a more intrinsic and informative representation for molecules is the 3D geometry , also known as conformation , where atoms are represented as their Cartesian coordinates . The 3D structures determine the biological and physical properties of molecules and hence play a key role in many applications such as computational drug and material design ( Thomas et al. , 2018 ; Gebauer et al. , 2021 ; Jing et al. , 2021 ; Batzner et al. , 2021 ) . Unfortunately , how to predict stable molecular conformation remains a challenging problem . Traditional methods based on molecular dynamics ( MD ) or Markov chain Monte Carlo ( MCMC ) are very computationally expensive , especially for large molecules ( Hawkins , 2017 ) . Recently , significant progress has been made with machine learning approaches , especially with deep generative models . For example , Simm & Hernandez-Lobato ( 2020 ) ; Xu et al . ( 2021b ) studied predicting atomic distances with variational autoencoders ( VAEs ) ( Kingma & Welling , 2013 ) and flow-based models ( Dinh et al. , 2017 ) respectively . Shi et al . ( 2021 ) proposed to use denoising score matching ( Song & Ermon , 2019 ; 2020 ) to estimate the gradient fields over atomic distances , through which the gradient fields over atomic coordinates can be calculated . Ganea et al . ( 2021 ) studied generating conformations by predicting both bond lengths and angles . As molecular conformations are roto-translational invariant , these approaches circumvent directly modeling atomic coordinates by leveraging intermediate geometric variables such as atomic distances , bond and torsion angles , which are roto-translational invariant . As a result , they are able to achieve very compelling performance . However , as all these approaches seek to indirectly model the intermediate geometric variables , they have inherent limitations in either training or inference process ( see Sec . 2 for a detailed description ) . Therefore , an ideal solution would still be directly modeling the atomic coordinates and at the same time taking the roto-translational invariance property into account . In this paper , we propose such a solution called GEODIFF , a principled probabilistic framework based on denoising diffusion models ( Sohl-Dickstein et al. , 2015 ) . Our approach is inspired by the diffusion process in nonequilibrium thermodynamics ( De Groot & Mazur , 2013 ) . We view atoms as particles in a thermodynamic system , which gradually diffuse from the original states to a noisy distribution in contact with a heat bath . At each time step , stochastic noises are added to the atomic positions . Our high-level idea is learning to reverse the diffusion process , which recovers the target geometric distribution from the noisy distribution . In particular , inspired by recent progress of denoising diffusion models on image generation ( Ho et al. , 2020 ; Song et al. , 2020 ) , we view the noisy geometries at different timesteps as latent variables , and formulate both the forward diffusion and reverse denoising process as Markov chains . Our goal is to learn the transition kernels such that the reverse process can recover realistic conformations from the chaotic positions sampled from a noise distribution . However , extending existing methods to geometric generation is highly non-trivial : a direct application of diffusion models on the conformation generation task lead to poor generation quality . As mentioned above , molecular conformations are roto-translational invariant , i.e. , the estimated ( conditional ) likelihood should be unaffected by translational and rotational transformations ( Köhler et al. , 2020 ) . To this end , we first theoretically show that a Markov process starting from an roto-translational invariant prior distribution and evolving with roto-translational equivariant Markov kernels can induce an roto-translational invariant density function . We further provide practical parameterization to define a roto-translational invariant prior distribution and a Markov kernel imposing the equivariance constraints . In addition , we derive a weighted variational lower bound of the conditional likelihood of molecular conformations , which also enjoys the rototranslational invariance and can be efficiently optimized . A unique strength of GEODIFF is that it directly acts on the atomic coordinates and entirely bypasses the usage of intermediate elements for both training and inference . This general formulation enjoys several crucial advantages . First , the model can be naturally trained end-to-end without involving any sophisticated techniques like bilevel programming ( Xu et al. , 2021b ) , which benefits from small optimization variances . Besides , instead of solving geometries from bond lengths or angles , the one-stage sampling fashion avoids accumulating any intermediate error , and therefore leads to more accurate predicted structures . Moreover , GEODIFF enjoys a high model capacity to approximate the complex distribution of conformations . Thus , the model can better estimate the highly multi-modal distribution and generate structures with high quality and diversity . We conduct comprehensive experiments on multiple benchmarks , including conformation generation and property prediction tasks . Numerical results show that GEODIFF consistently outperforms existing state-of-the-art machine learning approaches , and by a large margin on the more challenging large molecules . The significantly superior performance demonstrate the high capacity to model the complex distribution of molecular conformations and generate both diverse and accurate molecules . 2 RELATED WORK . Recently , various deep generative models have been proposed for conformation generation . Among them , CVGAE ( Mansimov et al. , 2019 ) first proposed a VAE model to directly generate 3D atomic coordinates , which fails to preserve the roto-translation equivariance property of conformations and suffers from poor performance . To address this problem , the majority of subsequent models are based on intermediate geometric elements such as atomic distances and torsion angles . A favorable property of these elements is the roto-translational invariance , ( e.g . atomic distances does not change when rotating the molecule ) , which has been shown to be an important inductive bias for molecular geometry modeling ( Köhler et al. , 2020 ) . However , such a decomposition suffers from several drawbacks for either training or sampling . For example , GRAPHDG ( Simm & HernandezLobato , 2020 ) and CGCF ( Xu et al. , 2021a ) proposed to predict the interatomic distance matrix by VAE and Flow respectively , and then solve the geometry through the Distance Geometry ( DG ) technique ( Liberti et al. , 2014 ) , which searches reasonable coordinates that matches with the predicted distances . CONFVAE further improves this pipeline by designing an end-to-end framework via bilevel optimization ( Xu et al. , 2021b ) . However , all these approaches suffer from the accumulated error problem , meaning that the noise in the predicted distances will misguide the coordinate searching process and lead to inaccurate or even erroneous structures . To overcome this problem , CONFGF ( Shi et al. , 2021 ) proposed to learn the gradient of the log-likelihood w.r.t coordinates . However , in practice the model is still aided by intermediate geometric elements , in that it first estimates the gradient w.r.t interatomic distances via denoising score matching ( DSM ) ( Song & Ermon , 2019 ; 2020 ) , and then derives the gradient of coordinates using the chain rule . The problem is , by learning the distance gradient via DSM , the model is fed with perturbed distance matrices , which may violate the triangular inequality or even contain negative values . As a consequence , the model is actually learned over invalid distance matrices but tested with valid ones calculated from coordinates , making it suffer from serious out-of-distribution ( Hendrycks & Gimpel , 2016 ) problem . Most recently , Ganea et al . ( 2021 ) proposed a highly systematic ( rule-based ) pipeline named GEOMOL , which learns to predict a set of geometric quantities ( i.e . length and angles ) and then reconstruct the local and global structures of the conformation in a sophisticated procedure . Despite its effectiveness , this method can be only adopted to molecules with a specific pattern , which limits it from wide applications . In our experiments , the pattern requirement filtered out nearly one third of molecular graphs , making it incomparable in our setting . Besides , there has also been efforts to use reinforcement learning for conformation search Gogineni et al . ( 2020 ) . Nevertheless , this method relies on rigid rotor approximation and can only model the torsion angles , and thus fundamentally differs from other approaches . 3 PRELIMINARIES . 3.1 NOTATIONS AND PROBLEM DEFINITION . Notations . In this paper each molecule with n atoms is represented as an undirected graph G = 〈V , E〉 , where V = { vi } ni=1 is the set of vertices representing atoms and E = { eij | ( i , j ) ⊆ |V| × |V| } is the set of edges representing inter-atomic bonds . Each node vi ∈ V describes the atomic attributes , e.g. , the element type . Each edge eij ∈ E describes the corresponding connection between vi and vj , and is labeled with its chemical type . In addition , we also assign the unconnected edges with a virtual type . For the geometry , each atom in V is embedded by a coordinate vector c ∈ R3 into the 3-dimensional space , and the full set of positions ( i.e. , the conformation ) can be represented as a matrix C = [ c1 , c2 , · · · , cn ] ∈ Rn×3 . Problem Definition . The task of molecular conformation generation is a conditional generative problem , where we are interested in generating stable conformations for a provided graph G. Given multiple graphs G , and for each G given its conformations C as i.i.d samples from an underlying Boltzmann distribution ( Noé et al. , 2019 ) , our goal is learning a generative model pθ ( C|G ) , which is easy to draw samples from , to approximate the Boltzmann function . 3.2 EQUIVARIANCE . Equivariance is ubiquitous in machine learning for atomic systems , e.g. , the vectors of atomic dipoles or forces should rotate accordingly w.r.t . the conformation coordinates ( Thomas et al. , 2018 ; Weiler et al. , 2018 ; Fuchs et al. , 2020 ; Miller et al. , 2020 ; Simm et al. , 2021 ; Batzner et al. , 2021 ) . It has shown effectiveness to integrate such inductive bias into model parameterization for modeling 3D geometry , which is critical for the generalization capacity ( Köhler et al. , 2020 ; Satorras et al. , 2021a ) . Formally , a function F : X → Y is equivariant w.r.t a group G if : F ◦ Tg ( x ) = Sg ◦ F ( x ) , ( 1 ) where Tg and Sg are transformations for an element g ∈ G , acting on the vector spaces X and Y , respectively . In this work , we consider the SE ( 3 ) group , i.e. , the group of rotation , translation in 3D space . This requires the estimated likelihood unaffected with translational and rotational transformations , and we will elaborate on how our method satisfy this property in Sec . 4 .

Summary: This paper tackles the problem of conditional generation of molecular conformations (i.e. 3D cartesian atom positions) given a molecular graph. The authors formulate the generation process via diffusion probabilistic models; Conformations are generated by learning a reverse diffusion process from isotropic gaussian noise to molecular conformations. They use a SE(3) invariant formulation of the diffusion process based on Kohler et al 2020, and they operate directly on atomic positions (i.e. a point cloud) instead of interatomic distances or an intermediate bond geometry representation. The authors show state of the art results evaluated by COV/MAT metrics on GEOM-Drugs and GEOM-QM9 datasets.