Generating Realistic 3D Molecules with an Equivariant Conditional Likelihood Model

1 INTRODUCTION . Identifying molecules with desirable characteristics is of fundamental importance in many fields , including drug discovery . The astronomically large number of possible drug-like compounds , however , makes an exhaustive experimental or virtual screening intractable ( Polishchuk et al. , 2013 ) . As a result , a large body of work has used machine learning to explore chemical space and propose molecules with specific characteristics . Organic molecules can be represented as graphs in which nodes are individual atoms , and edges are covalent bonds through which atoms share electrons . Each node in a molecular graph is labelled with the atomic number of its corresponding atom , and edges are labelled with the number of electrons shared in that covalent bond . Because each atom has a predetermined number of electrons with which to form bonds , only a subset of the possible edge labels result in chemically valid molecules . The connectivity-based description of a molecule provided by molecular graphs is important for many applications , including chemical synthesis , virtual screening , and molecular dynamics simulation , and has thus motivated extensive research into generative models for molecular graphs ( Gómez-Bombarelli et al. , 2018 ; Maragakis et al. , 2020 ; Liu et al. , 2018 ; Shi et al. , 2020 ; Jin et al. , 2018 ; 2020 ) . The functional characteristics of a given molecule , however , arise not only from its connectivity , but from its configuration in 3D space . Physical limitations on the lengths and angles of covalent bonds among specific types of atoms impose constraints on the set of geometric configurations that are compatible with a given molecular graph . Because most organic molecules contain rotatable bonds , however , these constraints are generally not sufficient to unambiguously reconstruct the 3D coordinates of all atoms in the molecule . The probability of observing a particular 3D geometry for a given molecule is in general a function of the quantum mechanical energy of that geometry , with low-energy geometries being more likely . Graph-based generative models are capable of producing a wide range of molecules whose atoms have the correct number and type of bonds , but their outputs do not contain any geometric information , and may include molecular graphs for which stable 3D geometries ( i.e. , low-energy geometries with physically realistic inter-atomic distances and bond angles ) do not exist . Although it is possible that a model could learn to assess the geometric feasibility of a molecular graph without training on any geometric data , a model that is trained explicitly on 3D molecules would likely be better at generating molecular graphs with corresponding low-energy 3D geometries . The lack of geometric information is also problematic because many downstream applications require 3D information . For instance , many drug discovery efforts screen for molecules with high predicted affinity for a target protein pocket using a computational docking process to score molecule poses . It would be extremely valuable if a machine learning model could directly generate such molecules , as this would expedite the screening process . A model that is explicitly trained to recreate the geometric poses of bound molecules , rather than just molecular graphs , has the potential to produce molecules and geometries that are more conducive to binding . To address the above shortcomings , we created GEN3D—a graph-generative model that proposes molecules with 3D coordinates . The model is rotationally and translationally equivariant , providing an inductive bias that exploits the symmetries of chemical space . GEN3D creates molecular graphs through a sequential sampling process , like other graph-based generative models ( Shi et al. , 2020 ; Liu et al. , 2018 ) , but at each step GEN3D also calculates a likelihood function over positions in 3D space , and uses that function to sample coordinates for each new atom . We show that GEN3D generates novel , chemically valid molecular graphs that have realistic low-energy geometries . In particular , we found that our model outperforms existing graph generators in its ability to create chemically valid molecular graphs , and outperforms existing 3D-generative models in the ability to create realistic geometries . In Appendix E , we also show that our pre-trained model can be tuned to generate novel molecules in geometries that score well in a virtual screening tool , which illustrates the potential application of our method in drug discovery . Finally , we demonstrate the geometric accuracy of our model by using it to sample geometries for fixed molecular graphs , and show that it achieves state-of-the-art results on the previously established benchmark of Xu et al . ( 2021b ) . 2 RELATED WORK . Different models construct molecular graphs using Graph Neural Networks ( GNNs ) in a number of different ways . CGVAE , for example , uses a GNN to create VAE latent representations for each node in a molecular graph ( Liu et al. , 2018 ) . The nodes and edges of the graph are then reconstructed in a sequential decision process , guided by another GNN , in which atoms are connected one at a time to a growing molecular graph . Another model , GraphAF , also generates molecular graphs one atom at a time , but it samples molecules using a Gaussian Autoregressive Flow ( Shi et al. , 2020 ) . Other models , like JT-VAE ( Jin et al. , 2018 ) and HierVAE ( Jin et al. , 2020 ) , generate molecules by sequentially adding multi-atom motifs rather than individual atoms . Because of their importance in classification tasks on point clouds and 3D molecules , there has been a great deal of interest in machine learning problems involving 3D structured data . When working with such data , it is often desirable that models be invariant or equivariant to rotations or translations of the input data , as these transformations are only an arbitrary change of coordinate systems . Approaches like SchNet ( Schütt et al. , 2017 ) and EGNN ( Satorras et al. , 2021b ) achieve equivariance by using invariant features like pairwise distances as inputs to the model . These approaches have proven effective in many domains , but they are reflection-invariant , so they can not distinguish between mirror-image isomers . Other approaches , like DimeNet ( Klicpera et al. , 2020 ) and SphereNet ( Liu et al. , 2021 ) , achieve equivariance through message-passing schemes that encode distance and angular information , while Tensor Field Networks ( Thomas et al. , 2018 ) and SE ( 3 ) -Transformers ( Fuchs et al. , 2020 ) process data using a basis of equivariant spherical harmonic functions . There are several prior works on generative models for 3D molecular data . Two models from the same group generate 3D molecules by outputting a voxelized grid of atomic densities , which is then converted to a set of atoms and bonds in a secondary optimization step ( Ragoza et al. , 2020 ; Masuda et al. , 2020 ) . While efficient to implement using convolutions , the voxelized grid approach imposes a practical limit on the size of generated molecules , and is not equivariant . The process of reconstructing a molecule from a density field also suffers from relatively high rates of chemical and geometric invalidity , and the one-shot nature of the generative process does not allow for valence constraints to be enforced . Another pair of recent papers aims to generate stable 3D molecules using reinforcement learning ( Simm et al. , 2019a ; b ) . These models , however , are only capable of generating molecules by placing atoms from a set corresponding to a pre-determined stoichiometry , have only been demonstrated on very small molecules , and do not generate bonding information . Two other 3D-generative models are closely related to GEN3D : E-NF ( Satorras et al. , 2021a ) and G-SchNet ( Gebauer et al. , 2019 ) . E-NF uses an equivariant , EGNN-based normalizing flow to convert random initial atom positions into realistic molecular geometries . This model , however , has only been demonstrated on relatively small molecules , is very expensive to train , and produces molecules with low rates of chemical validity . G-SchNet is an equivariant model that autoregressively generates 3D molecules , and has recently been extended to conditionally generate molecules with desired electronic properties and molecular fingerprints ( Gebauer et al. , 2021 ) . G-SchNet , however , exclusively generates atomic positions , and does not generate molecular graphs that provide the chemical bonding information required by many downstream applications . In addition , without explicitly generating bonds , one can not use simple rules to constrain chemical properties , like atomic valences and ring sizes , during the generative process . There are also machine learning methods that predict molecular geometries for a given molecular graph . A number of works have approached this problem by predicting inter-atomic distances from the molecular graph , and then using the predictions to generate a 3D geometry in a secondary optimization step ( Xu et al. , 2021a ; Simm & Hernández-Lobato , 2019 ) . Most recently , Xu et al . ( 2021b ) proposed the ConfVAE architecture , which estimates inter-atomic distances and then optimizes molecular geometry using an end-to-end differentiable optimization procedure . 3 PROBABILISTIC MODEL . We will represent a molecule as a 3D-dimensional graph G = ( V , A , X ) . For a molecule with n atoms , V ∈ Rn×d is a list of d-dimensional atom features , A ∈ Rn×n×b is an adjacency matrix with b-dimensional edge features , and X ∈ Rn×3 is a list of 3D atomic coordinates for each atom . In practice , V simply encodes the atomic number of each atom , and A encodes the number of shared electrons in each covalent bond . To model a chemical space of interest , we consider the distribution p ( V , A , X ) . Previous works have aimed at calculating various marginal and conditional densities with respect to this joint distribution . For example , graph-based generative models learn the marginal distribution p ( V , A ) = ∫ X p ( V , A , X ) dX , molecular geometry prediction amounts to learning the conditional distribution p ( X|V , A ) , and 3D generative models like G-SchNet learn the distribution p ( V , X ) = ∫ A p ( V , A , X ) dA . To learn the joint distribution p ( V , A , X ) , it is necessary to factorize the density . We choose to use the following factorization : p ( V , A , X ) = n∏ i=1 p ( V : i , X : i , A : i|V : i−1 , A : i−1 , X : i−1 ) · p ( stop|V , A , X ) = n∏ i=1 p ( X : i|V : i , A : i , X : i−1 ) p ( A : i|V : i , A : i−1 , X : i−1 ) p ( V : i|V : i−1 , A : i−1 , X : i−1 ) · p ( stop|V , A , X ) Here , n is the number of atoms in the input graph , and V : i , A : i and X : i indicate the graph ( V , A , X ) restricted to the first i atoms . Computing p ( V : i|V : i−1 , A : i−1 , X : i−1 ) is relatively simple because it amounts to predicting a single atom type based on a 3D graph ( V : i−1 , A : i−1 , X : i−1 ) . Calculating p ( A : i|V : i , A : i−1 , X : i−1 ) is more complex because it involves a prediction over a new row of the adjacency matrix . More concretely , computing the conditional density of A : i ∈ Ri×i×b amounts to computing a joint density over the new entries of the adjacency matrix Ai,1 , . . . , Ai , i−1 ∈ Rb . To solve this problem , we further decompose this distribution : p ( A : i|V : i , A : i−1 , X : i−1 ) = p ( Ai,1 , . . . , Ai , i−1|V : i , A : i−1 , X : i−1 ) = i−1∏ j=1 p ( Ai , j |Ai , :j−1 , V : i , A : i−1 , X : i−1 ) Intuitively , Ai,1 , . . . , Ai , i−1 represent the edges from atom i to atoms 1 . . . i− 1 . Finally , estimating the density p ( X : i|V : i , A : i , X : i−1 ) involves modeling a continuous distribution over positions Xi ∈ R3 for atom i . To accomplish this , we assume Xi belongs to a finite set of points X , and model its probability mass as a product of distributions over angles and distances : p ( Xi|V : i , A : i , X : i−1 ) = 1 C i−1∏ j=1 p ( ||Xi −Xj || | V : i , A : i , X : i−1 ) · ∏ ( j , k ) ∈I p ( Angle ( Xi −Xk , Xj −Xk ) | V : i , A : i , X : i−1 ) Intuitively , p ( ||Xi−Xj || | V : i , A : i , X : i−1 ) predicts the distances from each existing atom to the new atom , and p ( Angle ( Xi − Xk , Xj − Xk ) | V : i , A : i , X : i−1 ) predicts the bond angles of connected triplets of atoms involving atom i. I is a set of pairs ( j , k ) where atom k is connected to atom i , and atom j is connected to atom k. “ Angle ” denotes the angle between two vectors . C is a normalizing constant derived from summing this density over all of X . To increase the computational tractability of estimating this factorized density , we assume that the nodes in the molecular graph ( V , A , X ) are listed in the order of a breadth-first traversal over the molecular graph .

Summary: This paper proposes a generative model to sample a new molecule with its 3d coordinates. The generative process is divided into three different steps: 1) the atom network generates a new atom 2) the edge network connects the new atom to the existing atoms 3) the distance and angle network generates the pairwise distances and angles from existing atoms to the new atoms. EGNN is used for each of these sub-networks to impose the equivariance to translations and rotations. Experiments with QM9 datasets show the geometric correctness of the proposed generative process.