Direct Molecular Conformation Generation

1 INTRODUCTION . Molecular conformation generation aims to generate the 3D coordinates of all the atoms of molecules , which then can be used in molecular property prediction ( Axelrod & Gomez-Bombarelli , 2021 ) , docking ( Roy et al. , 2015 ) , structure-based virtual screening ( Kontoyianni , 2017 ) , etc . Molecular conformation can be physically observed using X-ray crystallography , but it is prohibitively costly for industry-scale tasks ( Mansimov et al. , 2019 ) . Ab initio methods , such as density functional theory ( DFT ) ( Parr , 1980 ; Baseden & Tye , 2014 ) , can accurately predict the molecular shapes , but take up to several hours per small molecule ( Hu et al. , 2021 ) .To handle large-scale molecules , people turn to leverage the classical force fields , like UFF ( Rappe et al. , 1992 ) or MMFF ( Halgren , 1996 ) , to estimate conformations , which is efficient but unacceptably inaccurate ( Kanal et al. , 2018 ) . Recently , people start to explore machine learning methods to generate conformation . Simm & Hernández-Lobato ( 2020 ) , Shi et al . ( 2020 ) and Shi et al . ( 2021 ) leveraged variational auto-encoder ( VAE ) , flow-based models and score-based methods for conformation generation , respectively . With learned models , molecule conformations can be sampled independently ( Simm & HernándezLobato , 2020 ; Shi et al. , 2020 ) or using Langevin dynamics ( Xu et al. , 2021a ; Shi et al. , 2021 ) . The common part of these methods is that they are all built upon the interatomic distances among atoms ( i.e. , the distance between atom pairs ) . Specifically , Simm & Hernández-Lobato ( 2020 ) ; Shi et al . ( 2020 ) ; Xu et al . ( 2021a ; b ) use various generative methods to model the distribution of interatomic distances , and then reconstruct conformations based on distances . Shi et al . ( 2021 ) leveraged a score-matching network to model the density gradient of interatomic distances . Moreover , Winter et al . ( 2021 ) use variational auto-encoder to first predict bond length , bond angle and dihedral angle and then reconstruct the coordinates based on the intermediate results . A major reason of using distance-based methods is that the interatomic distances are invariant to rotation and translation of the conformation . However , a potential drawback is that the predicted distances might conflict with each other . For example , they did not explicitly consider the triangle inequality for distances among any three atoms . The underlying degree-of-freedom of these distances is only 3N − 6 ( N is the number of atoms in a molecule ) while practically , these methods often generate distances with the degree-of-freedom of roughly N ( N − 1 ) /2 , which usually lead to violations of the triangle inequality . For example , we found that in GraphDG ( Simm & Hernández-Lobato , 2020 ) , a prevailing representative of distance-based method , 8.65 % of molecular graphs in the GEOM-Drugs test set will produce distance matrix that violates the triangle inequality . In this work , we explore the possibility to directly generate the coordinates of all atoms . This is a straightforward and more natural choice and avoids violations of the triangle inequality , and has demonstrated remarkable performance on protein structure prediction by the AlphaFold 2 ( Jumper et al. , 2021 ) . We design a model that generates atom coordinates directly . The generator in our model stacks multiple blocks , and each block outputs a conformation which is then refined by the following block . A block consists of several modules that can encode the previous conformation as well as the features of bonds , atoms and global information . At the end of each block , we add a normalization layer that centers the coordinates at the origin . Since a molecule may have multiple conformations , we use the variational auto-encoder ( VAE ) framework which allows diverse generation . To realize roto-translation equivariance , i.e. , rotating and translating a set of atom coordinates do not change the conformation , we adopt a loss that is invariant to rotation and translation of atom coordinates . We conduct experiments on four benchmark settings , which are GEOM-QM9 and GEOM-Drugs with small-scale setting ( Shi et al. , 2021 ) and large-scale setting ( Axelrod & Gomez-Bombarelli , 2021 ) . Compared with previous methods , ours achieves state-of-the-art results on GEOM-Drugs ( both small-scale and large-scale settings ) , and outperforms almost all previous baselines on GEOMQM9 , demonstrating the effectiveness of our method . Our contributions are summarized as follows : ( 1 ) We explore a new direction of conformation generation , that directly generates the coordinates of a molecular conformation without generating interatomic distances . We empirically show that directly generating conformation achieves state-of-the-art results on several tasks . ( 2 ) We leverage a fine-grained loss function for training , that is invariant to the roto-translation of the conformation . ( 3 ) We propose a new model that iteratively refines the conformations . Our model is inspired by multiple advanced architectures like GATv2 ( Brody et al. , 2021 ) , graph network ( GN ) block ( Battaglia et al. , 2018 ) , that can effectively model molecules . 2 PRELIMINARIES . In this section , we introduce the notations used in this work , give a formal definition of the molecular conformation generation problem , and briefly introduce how to align two conformations under rotation and translation . Notations : Let G = ( V , E ) denote a molecular graph , where V and E are collections of atoms and bonds , respectively . Specifically , V = { v1 , v2 , · · · , v|V | } with the i-th atom vi . Let eij denote the bond between atom vi and vj . LetN ( i ) denote the neighbors of atom i , i.e. , N ( i ) = { vj | eij ∈ E } . We useR to represent the conformation ofG , whereR ∈ R|V |×3 . The i-th row ofR is the coordinate of atom vi . Let ρ ( · ) denote a roto-translation operation , i.e . an affine transformation ρ ( R ) = Rt + b where the orientation-preserving orthogonal transformation t ∈ SO ( 3 ) ⊂ R3×3 represents a rotation and the vector b ∈ R1×3 represents a translation . Let MLP ( · · · ) denote the multi-layer perception network , where the inputs are concatenated together and then processed by linear mapping , Batch Normalization ( Ioffe & Szegedy , 2015 ) and ReLU activation sequentially . Problem definition : Given a graph G = ( V , E ) , our task is to learn a mapping , that can output the coordinates R of all atoms in V , i.e. , R ∈ R|V |×3 . 1 Matching loss : Let R1 and R2 denote two conformations , both of which are N × 3 matrices ( N is the number of rows in R1 and R2 ) . The matching loss between R1 and R2 is defined as follows : ` M ( R1 , R2 ) = min ρ ‖ρ ( R1 ) −R2‖2F , ( 1 ) 1A molecule corresponds to different conformations at different energy level . To model such cases , the problem is to generate R based on G and a random noise z. where ‖ · ‖F denotes the Frobenius norm , i.e. , ‖A‖2F = ∑ i , j |Aij |2 . The matching loss is invariant to the rotation and translation of either of the two input conformations . It is obviously nonnegative , and is zero if R1 is obtained by a roto-translational operation of R2 . Karney ( 2007 ) proposed to use quaternions to solve Eqn. ( 1 ) . A quaternion q is an extension of complex numbers , q = q0u + q1i + q2j + q3k , where q0 , q1 , q2 , q3 are scalars and u , i , j , k are orientation vectors . With quaternions , any rotation operation is specified by a 3 × 3 matrix , where each element in the matrix is the basic algebra of q0 to q3 and some constants . The solution to Eqn . ( 1 ) is the minimal eigenvalue of a 4× 4 matrix obtained by the basic algebra on R1 and R2 , the corresponding optimal rotation and translation ρ∗ can be obtained from the eigenvectors . Note that the best ρ∗ is related to both R1 and R2 . To stabilize training , we set the gradients of ρ∗ w.r.t . the model parameter as zero . More details can be found in Karney ( 2007 ) . 3 OUR METHOD . We introduce our proposed method in this section . 3.1 CONFORMATION GENERATION FRAMEWORK . Our goal can be formalized as learning to generate samples from the conditional distribution p ( R|G ) . We follow the variational auto-encoder ( VAE ) scheme ( Kingma & Welling , 2014 ; Rezende et al. , 2014 ; Sohn et al. , 2015 ) as it enables flexible model architecture ( vs. flow-based models ) , efficient ( i.e. , i.i.d ) generation ( vs. score-matching methods ) , and stable training and diverse sample generation ( vs. adversarial methods ) . Given a molecule graph G , the conditional variant of VAE generates R using a likelihood model p ( R|z , G ) where z is a latent code drawn from a prior p ( z ) = N ( 0 , I ) . The corresponding p ( R|G ) = ∫ p ( z ) p ( R|z , G ) dz is however intractable to evaluate and maximize to fit data . VAE handles this by introducing a ( conditional ) inference model q ( z|R , G ) which leads to the following identity : log p ( R|G ) = Eq ( z|R , G ) [ log p ( R|z , G ) ] −DKL ( q ( z|R , G ) ‖p ( z ) ) +DKL ( q ( z|R , G ) ‖p ( z|R , G ) ) , ( 2 ) where DKL denotes the Kullback-Leibler ( KL ) divergence , and p ( z|R , G ) = p ( z ) p ( R|z , G ) ∫ p ( z ) p ( R|z , G ) dz is the true posterior . Due to the nonnegativity of the third term , the first two terms form a lower bound of log p ( R|G ) . When q ( z|R , G ) is properly chosen , the bound is tractable to optimize . Moreover , as the left-hand-side is independent of q ( z|R , G ) , tightening the bound ( i.e. , minimizing the third term ) can be achieved by maximizing the bound w.r.t q ( z|R , G ) . Therefore , as a loss function to be minimized , the objective is : Eq ( z|R , G ) [ − log p ( R|z , G ) ] +DKL ( q ( z|R , G ) ‖p ( z ) ) . ( 3 ) Now we develop our specification for the conformation generation task . We construct the likelihood model with regard to a decoded reference conformation R̂ ( z , G ) : p ( R|z , G ) ∝ exp { − 1 γ ` M ( R̂ ( z , G ) , R ) } , ( 4 ) where γ > 0 is a variance parameter that we choose to fix as 1 ( see Sec . 3.4 ( 1 ) for explanation ) . Note that the usage of ` M guarantees all roto-translation equivalent conformations have the same probability . For the inference model , we choose q ( z|R , G ) = N ( z|µR , G , ΣR , G ) , ( 5 ) where the conditional mean and variance are outputs from some encoder . It enables tractable loss optimization via reparameterization : z ∼ q ( z|R , G ) is equivalent to z = µR , G + ΣR , G where ∼ N ( 0 , I ) . With the above specification , the loss function becomes : minE ∼N ( 0 , I ) ` M ( R̂ ( µR , G + ΣR , G , G ) , R ) + βDKL ( N ( µR , G , ΣR , G ) ‖N ( 0 , I ) ) , ( 6 ) where the minimization is applied over all model parameters . Note that in Eqn . ( 6 ) , an additional hyperparameter β is introduced in the spirit of β-VAE ( Higgins et al. , 2016 ) to handle the strength of the prior regularization . The second term in Eqn . ( 6 ) has a closed-form expression for optimization .

Summary: The authors present a method that is able to accurately generate molecule confirmations, achieving competitive results across four datasets.  Unlike previous methods that often rely on directly predicting various inter-atomic distances (through a distance matrix) (which in some cases results in physically impossible molecules), the presented work predicts the location of all atoms. The loss function is invariant to rotation, an important property of many molecules.