A Generative Model for Molecular Distance Geometry

1 INTRODUCTION . Over the last few years , many highly-effective deep learning methods generating small molecules with desired properties ( e.g. , novel drugs ) have emerged ( Gómez-Bombarelli et al. , 2018 ; Segler et al. , 2018 ; Dai et al. , 2018 ; Jin et al. , 2018 ; Bradshaw et al. , 2019a ; Liu et al. , 2018 ; You et al. , 2018 ; Bradshaw et al. , 2019b ) . These methods operate using graph representations of molecules in which nodes and edges represent atoms and bonds , respectively . A representation that is closer to the physical system is one in which a molecule is described by its geometry or conformation . A conformation x of a molecule is defined by a set of atoms { ( i , ri ) } Nvi=1 , where Nv is the number of atoms in the molecule , i ∈ { H , C , O , ... } is the chemical element of the atom i , and ri ∈ R3 is its position in Cartesian coordinates . Importantly , the relative positions of the atoms are restricted by the bonds in the molecule and the angles between them . Due to thermal fluctuations resulting in stretching of and rotations around bonds , there exist infinitely many conformations of a molecule . A molecule ’ s graph representation and a set of its conformations are shown in Fig . 1 . Under a wide range of conditions , the probability p ( x ) of a conformation x , is governed by the Boltzmann distribution and is proportional to exp { −E ( x ) /kBT } , where E ( x ) ∈ R is the conformation ’ s energy , kB is the Boltzmann constant , and T is the temperature . To compute a molecular property for a molecule , one must sample from p ( x ) . The main approach is to start with one conformation and make small changes to it over time , e.g. , by using Markov chain Monte Carlo ( MCMC ) or molecular dynamics ( MD ) . These methods can be used to accurately sample equilibrium states of molecules , but they become computationally expensive for larger ones ( Shim & MacKerell , 2011 ; Ballard et al. , 2015 ; De Vivo et al. , 2016 ) . Other heuristic approaches exist in which distances between atoms are set to fixed idealized values ( Havel , 2002 ; Blaney & Dixon , 2007 ) . Several methods based on statistical learning have also recently been developed to tackle the issue of conformation generation . However , they are mainly geared towards studying proteins and their folding dynamics ( AlQuraishi , 2019 ) . Some of these models are not targeting a distribution over conformations but the most stable folded configuration ( Evans et al. , 2018 ; Ingraham et al. , 2019 ) , while others are not transferable between different molecules ( Lemke & Peter , 2019 ; Noé et al. , 2019 ) . This work includes the following key contributions : • We introduce a novel probabilistic model for learning conformational distributions of molecules with graph neural networks . • We create a new , challenging benchmark for conformation generation , which is made publicly available . To the best of our knowledge , this is the first benchmark of this kind . • By combining a conditional variational autoencoder ( CVAE ) with an Euclidean distance geometry ( EDG ) algorithm we present a state-of-the-art approach for generating one-shot samples of molecular conformations for unseen molecules that is independent of their size and shape . • We develop a rigorous experimental approach for evaluating and comparing the accuracy of conformation generation methods based on the mean maximum deviation distance metric . • We show how this generative model can be used as a proposal distribution in an importance sampling ( IS ) scheme to estimate molecular properties . 2 METHOD . Our goal is to build a statistical model that generates molecular conformations in a one-shot fashion from a molecule ’ s graph representation . First , we describe how a molecule ’ s conformation can be represented by a set of pairwise distances between atoms and why this presentation is advantageous over one in Cartesian coordinates ( Section 2.1 ) . Second , we present a generative model in Section 2.2 that will generate sets of atomic distances for a given molecular graph . Third , we explain in Section 2.3 how a set of predicted distances can be transformed into a molecular conformation and why this transformation is necessary . Finally , we detail in Section 2.4 how our generative model can be used as a proposal distribution in an IS scheme to estimate molecular properties . 2.1 EXTENDED MOLECULAR GRAPHS AND DISTANCE GEOMETRY . In this study , a molecule is represented by an undirected graph which is defined as a tuple G = ( V , E ) . V = { vi } Nvi=1 is the set of nodes representing atoms , where each vi ∈ RFv holds atomic attributes ( e.g. , the element type i ) . E = { ( ek , rk , sk ) } Nek=1 is the set of edges , where each ek ∈ RFe holds an edge ’ s attributes ( e.g. , the bond type ) , and rk and sk are the nodes an edge is connecting . Here , E represents the molecular bonds ( and the auxiliary edges which are explained below ) in the molecule . We assume that , given a molecular graph G , one can represent one of its conformations x by a set of atomic distances d = { dk } Nek=1 , where dk = |rrk − rsk | is the Euclidean distance between the positions of the atoms rk and sk in this conformation . As the set of edges between the bonded atoms ( Ebond ) alone would not suffice to describe a conformation , we expand the traditional graph representation of a molecule by adding auxiliary edges . Auxiliary edges between atoms that are second neighbors in the original graph fix angles between atoms , and those between third neighbors fix dihedral angles ( denoted Eangle and Edihedral , respectively ) . In this work , Eangle consists of edges between all second neighbors in the original graph . Edges between third neighbors are added according to a heuristic ( see Appendix A.1 ) . From now on we are always referring to this extended molecular graph when talking about molecular graphs . In Fig . 2 , the process of extending the molecular graph and the extraction of d from x and G are illustrated . A key advantage of a representation in terms of distances is its invariance to rotation and translation ; by contrast , Cartesian coordinates depend on the ( arbitrary ) choice of origin , for example . In addition , it reflects pair-wise physical interactions and their generally local nature . Auxiliary edges can be placed between higher-order neighbors depending on how far the physical interactions dominating the potential energy of the system reach . We have a set of NG molecular graphs { Gl } NGl=1 . Further , for each Gl , we have Sl conformational samples { xl , j } Slj=1 from the ground-truth distribution resulting in Sl sets of distances { dl , j } Sl j=1 . With this data , we will train a generative model which we detail in the following section . 2.2 GENERATIVE MODEL . We employ a CVAE ( Kingma & Welling , 2014 ; Pagnoni et al. , 2018 ) to model the distribution over distances d given a molecular graph G. A CVAE first encodes G together with d into a latent space z ∈ RkNv , where k ∈ N+ , with an encoder qφ ( z|d , G ) . Subsequently , the decoder pθ ( d|z , G ) decodes z back into a set of distances . A graphical model is shown in Fig . 2 C ) . A conformation has , in general , 3Nv − 6 spatial degrees of freedom ( dofs ) : one dof per spacial dimension per atom minus three translational and three rotational dofs . Therefore , the latent space should be proportional to the number of atoms in the molecule . In addition , the latent space should be smaller than 3Nv as it is the role of the encoder to project the conformation into a lower-dimensional space . As a result , we set k = 1 to avoid overfitting.1 Here , qφ ( z|d , G ) and pθ ( d|z , G ) are Gaussian distributions , the mean and variance of which are modeled by two artificial neural networks . At the center of this model are message-passing neural networks ( MPNNs ) ( Gilmer et al. , 2017 ) with multi-head attention ( Veličković et al. , 2018 ) . In short , an MPNN is a convolutional neural network that allows end-to-end learning of prediction pipelines whose inputs are graphs of arbitrary size and shape . In a convolution , neighboring nodes exchange so-called messages between neighbors to update their attributes . Edges update their attributes with the features of the nodes they are connecting . The MPNN is a well-studied technique that achieves state-of-the-art performance in representation learning for molecules ( Kipf & Welling , 2017 ; Duvenaud et al. , 2015 ; Kearnes et al. , 2016 ; Schütt et al. , 2017b ; Gilmer et al. , 2017 ; Kusner et al. , 2017 ; Bradshaw et al. , 2019a ) . In the following , we describe the details of the model.2 In Fig . 3 , an illustration of the model is shown . In the encoder qφ ( z|d , G ) , each dk is concatenated with the respective edge feature ek to give e′k ∈ RFe+1 . Then , each vi and each e′k are passed to Fenc , v and Fenc , e ( two multilayer perceptrons , 1Experiments showed that our model performs similarly with a latent space of R2Nv and overfits with latent spaces of R3Nv . 2The full model ( including all parameters ) is available online https : //figshare.com/s/ 1b42bf865bd78c457354 MLPs ) , respectively , to give G ( 0 ) enc , where G ( t ) enc = ( { v ( t ) i , enc } Nv i=1 , { ( e ( t ) k , enc , rk , sk ) } Ne k=1 ) , v ( t ) i , enc ∈ RLv , and e ( t ) k , enc ∈ RLe . Then , T MPNNs of depth 1 , { MP ( t ) enc } Tt=1 , are consecutively applied to obtain G ( T ) enc . Finally , the read-out function Renc ( an MLP ) takes each v ( T ) i , enc to predict the mean µzi ∈ R and the variance σ2zi ∈ R of the Gaussian distribution for zi . The so-called reparametrization trick is employed to draw a sample for zi . In summary , v ( 0 ) i , enc = Fenc , v ( vi ) , e ( 0 ) k , enc = Fenc , e ( e ′ i ) , G ( 1 ) enc = MP ( 0 ) enc ( G ( 0 ) enc ) , G ( t+1 ) enc = MP ( t ) enc ( G ( t ) enc ) , G ( T ) enc = MP ( T−1 ) enc ( G ( T−1 ) enc ) , µzi , σ 2 zi = Renc ( v ( T ) i , enc ) . ( 1 ) In the decoder pθ ( d|z , G ) , each zi is concatenated with the respective node feature vi to give v′i ∈ RFv+1 . Each v′i and each ek are passed to Fdec , v and Fdec , e ( two MLPs ) , respectively , to give G ( 0 ) dec , where G ( t ) dec = ( { v ( t ) i , dec } Nv i=1 , { ( e ( t ) k , dec , rk , sk ) } Ne k=1 ) , v ( t ) i , dec ∈ RLv , and e ( t ) k , dec ∈ RLe . Then , T MPNNs of depth 1 , { MP ( t ) dec } Tt=1 , are consecutively applied to obtain G ( T ) dec . Finally , the read-out functionRdec ( an MLP ) takes each e ( T ) k , dec to predict the mean µdk ∈ R and the variance σ2dk ∈ R of the Gaussian distribution for dk . In summary , v ( 0 ) i , dec = Fdec , v ( v ′ i ) , e ( 0 ) k , dec = Fdec , e ( ei ) , G ( 1 ) dec = MP ( 0 ) dec ( G ( 0 ) dec ) , G ( t+1 ) dec = MP ( t ) dec ( G ( t ) dec ) , G ( T ) dec = MP ( T−1 ) dec ( G ( T−1 ) dec ) , µdk , σ 2 dk = Rdec ( e ( T ) k , dec ) . ( 2 ) The sets of parameters in the encoder and decoder , φ and θ ( i.e. , parameters in Fenc , v , Fenc , e , { MP ( t ) enc } Tt=1 , Renc , Fdec , v , Fdec , e , { MP ( t ) dec } Tt=1 , Rdec ) , respectively , are optimized by maximizing the evidence lower bound ( ELBO ) : L = Ez∼qφ ( z|d , G ) [ log pθ ( d|z , G ) ] −DKL [ qφ ( z|d , G ) ||pθ ( z|G ) ] , ( 3 ) where the prior pθ ( z|G ) consists of factorized Gaussians . The optimal values for the hyperparameters for the network dimensions , number of message passes , batch size , and learning rate of the Adam optimizer ( Kingma & Ba , 2014 ) were tuned by maximizing the validation performance ( ELBO ) with a Bayesian optimizer and are reported in Appendix A.1.3 .

Summary: The authors propose a generative model designed for molecules, which is essentially a conditional variational auto-encoder. The model learns to generate, when conditioning on a molecule graph, the distribution of distances between each of the atoms and its second and third neighbour. Finally, using these information new molecules can be generated, which satisfy the generated distances. In the experiments the authors show the effectiveness of the method.