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 describes and tests GeoDiff, an end-to-end denoising diffusion model for generating molecular conformations from a molecular graph. The idea is to learn how to denoise a random conformation so as to generate a realistic structure of a small molecule. To do so, a molecular structure is first transformed into a random configuration from an uncorrelated multivariate normal distribution via a diffusion process (forward process). By training a backwards model that inverts the diffusion process, molecular conformations can then be generated from white noise. The reverse process is modeled as a sequence of Markov kernels where each kernel is a Gaussian whose means are neural networks that are trained based on the forward  trajectories. Molecular conformations are directly parameterized in Cartesian coordinates. Therefore, an important property that should be imposed on the reverse process is *equivariance*. Equivariance refers to the following behavior of a mapping under the action of a group: Transforming the input and then mapping it to some output is the same as applying the mapping first and then transforming the output. Here, the relevant group, the special Euclidean group $SE(3)$, is formed by the rigid transformations in $\mathbb R^3$. Equivariance is implemented by drawing the initial conformation from a distribution that is invariant under rigid transformations (isotropic Gaussian that removes the center of mass) and by using equivariant Markov kernels for the reverse process. The equivariance of the Markov kernels is implemented with equivariant convolutional layers (graph field network (GFN)). Two invariant objective functions for training are derived ("alignment approach" and "chain-rule approach"). GeoDiff is tested on two standard benchmarks and shows a superior performance compared to other molecular configuration generators. Also property predictions made by GeoDiff are better than with concurrent methods.