Directional Message Passing on Molecular Graphs via Synthetic Coordinates

1 Introduction . Graph neural networks ( GNNs ) have set the state of the art on many tasks of molecular machine learning , such as the prediction of quantum mechanical properties ( Gilmer et al. , 2017 ) , solubility ( Wu et al. , 2018 ) , or the generation of new molecules ( Jin et al. , 2020 ) . Thanks to their fast inference time , good generalization and scalability , GNNs are thus promising to revolutionize large parts of chemistry , from ab-initio quantum mechanical simulations and reaction kinetics to synthesis planning and drug discovery . Atom positions are central to many of these tasks , but unavailable in most cases . Many tasks in chemistry instead use a more coarse-grained representation : The molecular graph . Unfortunately , this representation makes many predictive tasks substantially harder , and GNNs have performed significantly better when they have access to the exact molecular configuration ( Gilmer et al. , 2017 ) . Missing atom positions furthermore preclude the use of many advanced GNNs that were developed with coordinates in mind . In this work we aim to fill in this information with well-defined coordinates constructed purely from the molecular graph . Regular approximation methods for generating atom positions often do not benefit model performance , due to the fundamental ambiguity of molecular configurations . The energy landscape of molecules can have multiple local minima , and a molecule can be in any of multiple different minima , known as conformers . In this work , we propose to circumvent this problem by incorporating the conformational ambiguity via empirical distance bounds . Instead of yielding a potentially wrong configuration , these bounds only estimate the range of viable molecular geometries . They are thus valid regardless of which state the molecule was in when generating the data . 1https : //www.daml.in.tum.de/synthetic-coordinates 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . A molecular configuration is fully specified by the pairwise distances between all atoms , due to rotational , translational , and reflectional invariance . We can thus obtain a molecular geometry from any method that provides pairwise distances between atoms . Since directional message passing does not require the full molecular geometry , these distances do not need to correspond to an actual threedimensional configuration . We leverage this generality and propose purely graph-based distances calculated from a symmetric variant of personalized PageRank ( PPR ) as a second set of coordinates . This distance performs surprisingly well , despite incorporating no chemical knowledge . Both the distance bounds and the symmetric PPR distance require no hand-tuning and can be calculated efficiently , even for large molecules . We leverage these two variants of synthetic coordinates to transform regular GNNs into directional message passing , as illustrated in Fig . 1 . We first calculate the synthetic , pairwise distances for the given molecular graph . Based on these , we calculate the edge distances and angles between edges . Finally , we compute the molecule ’ s line graph . Executing a GNN on the line graph improves its expressivity ( Garg et al. , 2020 ) and allows us to incorporate angular information . We use the original node and edge attributes together with the distances as node attributes , and the obtained angles as edge attributes . The GNN is then executed on this featurized line graph instead of the original graph . Our experiments show this transformation can significantly improve the performance of the underlying GNN , across multiple models and datasets . Incorporating synthetic coordinates reduces the error of a normal GNN by 55 % , putting it on par with the current state of the art . Our enhanced version of the SMP model ( Vignac et al. , 2020 ) improves upon the current state of the art on ZINC by 21 % , and DimeNet++ ( Klicpera et al. , 2020a ) with synthetic coordinates outcompetes previous methods on coordinate-free QM9 by 20 % . In summary , our core contributions are : • Well-defined synthetic coordinates based on node distances and simple molecular bounds , which significantly improve the performance of GNNs for molecules . • A general scheme of converting a normal GNN into a directional MPNN , which can improve performance and allows incorporating both distance and angular information . 2 Directional message passing . Graph neural networks . To use GNNs for molecules we represent them as graphs G = ( V , E ) , where the atoms define the node set V and the interactions the edge set E . These interactions are usually the bonds of the molecular graph , but they can also be all atoms pairs within a given cutoff of e.g . 5 Å . In this work we focus on an extension of message passing neural networks ( MPNNs ) ( Gilmer et al. , 2017 ) . MPNNs embed each atom u separately as hu ∈ RH , and can additionally use interaction embeddings e ( uv ) ∈ RHe . These embeddings are updated in each layer using messages passed between neighboring nodes , starting with the atom features h ( 0 ) u = x ( V ) u ( e.g . its type ) and the interaction features e ( 0 ) ( uv ) = x ( E ) ( uv ) ( e.g . the bond type or a distance representation ) . Extended MPNNs can be expressed via the following two equations : h ( l+1 ) u = fupdate ( h ( l ) u , Agg v∈Nu [ fmsg ( h ( l ) u , h ( l ) v , e ( l ) ( uv ) ) ] ) , ( 1 ) e ( l+1 ) ( uv ) = fedge ( h ( l+1 ) u , h ( l+1 ) v , e ( l ) ( uv ) ) . ( 2 ) The atom and interaction update functions fnode and fedge and the message function fmsg are learnable functions , such as simple linear layers or arbitrarily complex neural networks . The aggregation Agg over the atom ’ s neighbors Nu is usually a simple summation . Line graph . The directed line graph L ( G ) = ( VL , EL ) expresses the adjacencies between the directed edges in G. Its nodes are the directed edges of the original graph VL = { ( u , v ) | u ∈ V , v ∈ Nu } . For undirected graphs like molecular graphs , every undirected edge { u , v } is split into two directed edges ( u , v ) and ( v , u ) . Two nodes in L ( G ) are connected if the corresponding edges in G share a node , i.e . EL = { ( ( u , v ) , ( v , w ) ) | ( u , v ) , ( v , w ) ∈ VL } . We obtain node features for the line graph by embedding the original node and edge features as x ( VL ) ( uv ) = femb ( x ( V ) u , x ( V ) v , x ( E ) ( uv ) ) . The line graph can furthermore incorporate additional features for atom triplets as edge features x ( EL ) ( uvw ) , such as the angle between bonds or interactions . Directional message passing . Directional MPNNs improve upon regular MPNNs in two ways . First , they embed the directed messages instead of the nodes in the graph , essentially operating on the directed line graph . Models using only this first step are also known as directed MPNNs or line graph neural networks ( Dai et al. , 2016 ; Yang et al. , 2019 ; Chen et al. , 2019 ) . Directed MPNNs are strictly more expressive than regular MPNNs ( Morris et al. , 2020 ) . We can transform any MPNN to a directed MPNN simply by executing it on the directed line graph instead of the original graph . Second , for graphs with nodes that are embedded in an inner product space ( such as molecules in 3D space ) the directed edges correspond to directions in that space , via x ( VL ) ( uv ) = x ( V ) u −x ( V ) v . Directional MPNNs leverage this connection to better represent the molecular configuration , usually by using the angles in x ( EL ) ( uvw ) ( Klicpera et al. , 2020b ) . To fully leverage both aspects of directional MPNNs we therefore need some form of coordinates . Expressivity of GNNs . A central limitation of GNNs is their inability of distinguishing between certain non-isomorphic graphs . For example , GNNs are not able to distinguish between a hexagon and two triangles if all nodes and edges have the same features . More specifically , Xu et al . ( 2019 ) ; Morris et al . ( 2019 ) have shown that GNNs are only as powerful as the 1-Weisfeiler-Lehman ( WL ) test of isomorphism . While it is still possible to construct indistinguishable examples for directional MPNNs , this is significantly more difficult ( Garg et al. , 2020 ) . Dym & Maron ( 2021 ) have shown that MPNNs using SO ( 3 ) group representations and atom positions are even universal , i.e . able to approximate any continuous function to arbitrary precision . This demonstrates that coordinates can alleviate and even solve this central limitation of GNNs . 3 Molecular configurations . To prevent any pitfalls when constructing synthetic coordinates for GNNs based on chemical knowledge we first need to consider the properties of atomic positions in a molecule and how they are obtained . At first glance these positions might seem like an obvious and straightforward property . However , molecular configurations are actually ambiguous and difficult to obtain , even for small molecules . This misconception has even led some works to suggest semi-supervised learning methods leveraging positions , effectively treating them as abundant input features ( Hao et al. , 2020 ) . To clarify this issue we will next describe the complexity behind molecular configurations and how to approximate them efficiently . Finding molecular configurations . The atoms of a molecule can in principle be at any arbitrary position . However , most of these configurations will lead to an extremely high energy and are thus very unlikely to be observed in nature . A molecular configuration thus usually refers to the atom positions at or close to equilibrium , i.e . at the molecule ’ s energy minimum . To find these positions we have to search the molecule ’ s energy landscape and solve a non-convex optimization problem . This is in fact a bilevel optimization problem , where the atom positions are optimized in the outer and the electron wavefunctions in the inner task . These wave functions can then be ignored in the outer task ; they only influence the energy and the forces Fi = − ∂E∂xi acting on each atomic nucleus . We can then use these forces for gradient-based optimization , and avoid saddle points by using quasi-Newton methods . Difficulties . The above optimization process is very expensive due to the quantum mechanical ( QM ) computations required for optimizing the electron wavefunction at each gradient step . It is orders of magnitude more expensive than calculating the energy of a given molecular configuration , since we need to calculate the energy ’ s gradient for each optimization step . Furthermore , the optimization will only converge to a local , and not the global minimum . And in fact , the global minimum is not the only state of interest — any reasonably low local minimum of the energy landscape is a valid configuration , known as a conformer . A molecule thus does not have a unique configuration ; it can be in any of these states . Their statistical distribution and the interaction between them is central for many molecular properties . This ambiguity of atom positions poses a fundamental limit on how precise molecular predictions can be without knowing the exact ( ensemble of ) configurations . For example , without knowing the molecule ’ s conformer we can not reasonably predict its energy at a precision below roughly 60 meV ( Grimme , 2019 ) — except for small , rigid molecules that do not have multiple conformers ( e.g . benzene ) . Approximating energies and forces . The most prominent way of accelerating the process of finding a valid molecular configuration is by approximating its most expensive part : The quantum mechanical optimization of the electron wavefunction . There is a large hierarchy of methods with various runtime versus accuracy trade-offs ( Folmsbee & Hutchison , 2021 ) . The cheapest class of methods are force fields . They allow running molecular dynamics simulations with millions of atoms , and can estimate the equilibrium structure of a small molecule in less than one second . Force fields approximate the quantum-mechanical interactions via a closed-form , differentiable function that only depends on the atom positions . One common example is the Merck Molecular Force Field ( MMFF94 ) ( Halgren , 1996 ) . MMFF94 calculates the molecular energy based on interatomic distances , angles , dihedral angles , and long-range interaction terms . Each term is approximated using an analytic equation with empirically chosen coefficients that depend on the involved atom types . Forces are obtained via the analytical gradients Fi = − ∂E∂xi , and conformers via gradient-based optimization . Generating configurations with force fields is fast enough to even generate a large ensemble of conformers . However , the resulting conformers are highly biased and require corrections based on expensive QM-based methods for reasonably approximating the molecule ’ s true distribution ( Ebejer et al. , 2012 ; Kanal et al. , 2018 ) . Directly predicting the configuration . There are multiple methods that circumvent the optimization process to quickly generate low-energy conformers for a given molecular graph . Distance geometry methods generate conformers using an experimental database of ideal bond lengths , bond angles , and torsional angles ( Havel , 2002 ) . The ETKDG method combines this with empirical torsional angle preferences ( Riniker & Landrum , 2015 ) . Multiple machine learning methods for generating conformers have also recently been proposed ( Weinreich et al. , 2021 ; Lemm et al. , 2021 ) . Restrictions for ML . All of the above methods yield reasonable molecular configurations . However , they often require many initializations and a considerable amount of hand-tuning to yield a good result for every molecule in a dataset . Furthermore , the obtained conformer might not even be the correct one for the data of interest . The data could have been generated by a different conformer or by a statistical ensemble of multiple conformers . The configuration of a wrong conformer can cause our model to overfit to the false training data and cause bad generalization ( see Sec . 6 ) . To solve this issue we could try to generate an ensemble of conformers and embed their distribution . However , cheap generation methods yield strongly biased ensembles and would thus require expensive postprocessing , defeating the purpose of fast and scalable machine learning ( ML ) methods ( Ebejer et al. , 2012 ; Kanal et al. , 2018 ) . We propose to instead solve this issue by using less precise synthetic coordinates that are easier and cheaper to obtain .

Summary: The authors use a generic algorithm to embed molecular graph into 3D space using a distance matrix computed from PageRank. Then they use this 3D information as additional features to do property prediction. Results show that a) this approach is better than using GNNs without this information b) the augmented structures can even be applied to on Euclidean GNNs such as DimNet and get better results compared to simple GNNs.