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 paper proposed a new VAE-based generative model for generating molecular conformations from graphs. The whole formulation and training objective are similar to existing VAE-based models like CVGAE, and the main contrition lies in the architecture of the decoder. The architecture borrows the idea from the recent advance of AlphaFold: the proposed decoder is composed of several identical updating blocks, where they iteratively update the node/edge embeddings as well as the coordinates in each block. And the blocks are implemented with advanced graph neural networks.