Spherical Message Passing for 3D Molecular Graphs

1 INTRODUCTION . In many real-world studies , structured objects such as molecules are naturally modeled as graphs ( Gori et al. , 2005 ; Wu et al. , 2018 ; Shervashidze et al. , 2011 ; Fout et al. , 2017 ; Liu et al. , 2020 ; Wang et al. , 2020 ) . With the advances of deep learning , graph neural networks ( GNNs ) have been developed for learning from graph data ( Kipf & Welling , 2017 ; Defferrard et al. , 2016 ; Veličković et al. , 2018 ; Zhang et al. , 2018 ; Xu et al. , 2019 ; Gao & Ji , 2019 ; Gao et al. , 2018 ; 2020 ) . Currently , the message passing scheme ( Gilmer et al. , 2017 ; Sanchez-Gonzalez et al. , 2020 ; Vignac et al. , 2020 ; Battaglia et al. , 2018 ) is one of the commonly used architectures for realizing GNNs . In this work , we aim at developing a novel message passing method for 3D graphs . Generally , a 3D molecular graph contains 3D coordinates for each atom given in the Cartesian system along with the graph structure ( Liu et al. , 2019 ; Townshend et al. , 2019 ; Axelrod & Gomez-Bombarelli , 2020 ) . Different types of relative 3D information can be derived from 3D molecular graphs , and they can be important in molecular learning , such as bond lengths , angles between bonds ( Schütt et al. , 2017 ; Klicpera et al. , 2020b ) . We first investigate complete representations of 3D molecules . This requires the graph structure to be uniquely defined by relative 3D information . To this end , we conduct formal analyses in the spherical coordinate system ( SCS ) , and show that relative location of each atom in a 3D graph is uniquely determined by three geometries , including distance , angle , and torsion . However , such completeness needs to involve edge-based 2-hop information , leading to excessively high computational complexity . To circumvent the computational cost , we propose a novel message passing scheme , known as the spherical message passing ( SMP ) , for fast and accurate 3D molecular learning . Our SMP is efficient and approximately complete in representing 3D molecules . First , we design a novel strategy to compute torsion , which only considers edge-based 1-hop information , thus substantially reducing training complexity . This enables the generalization of SMP to large-scale molecules . In addition , we show that our SMP can distinguish almost all 3D graph structures . The uncovered cases are clearly stated and they seem rarely appear in practice . By naturally using relative 3D information and a novel torsion , SMP yields predictions that are invariant to translation and rotation of input graphs . We apply the SMP to real-world molecular learning , where meaningful physical representations are needed . Geometries ( d , θ , φ ) specified by SMP are then physically represented by Ψ ( d , θ , φ ) , which can be a solution to the Schrödinger equation , as described in Sec . 4 . Based on this , we develop the spherical message passing neural networks , known as the SphereNet , for 3D molecular learning . We conduct experiments on various types of datasets including OC20 , QM9 , and MD17 . Results show that , compared with baseline methods , SphereNet achieves the best performance without increasing the computing budget . Ablation study reveals contributions and necessity of different types of 3D information , including distance , angle , and torsion . Particularly , we compare with a complete message passing scheme that can distinguish all 3D graph structures but involves edge-based 2-hop information . Experimental results show that SphereNet achieves comparable performance but reduces running time by 4 times . This suggests the use of SphereNet in practice rather than the complete message passing scheme , whose computational complexity prevents its use on large molecules . 2 COMPLETE REPRESENTATIONS OF MOLECULES . Equivariant graph neural networks ( EGNNs ) represent one research area for 3D molecular graphs , as introduced in Sec . 5.1 . These methods usually take coordinates in the Cartesian coordinate system ( CCS ) for all atoms as the raw input . Hence , all the network layers need to be carefully designed to be equivariant . The computing of some equivariant components is expensive , like spherical harmonics and Clebsh-Gordan coefficients ( Thomas et al. , 2018 ; Fuchs et al. , 2020 ) . In addition , the complicated SE ( 3 ) group representations may not be necessary for molecular learning where final representations are generally required to be invariant . In this work , we focus on the other category of methods that take relative position information purely as input to graph learning models . Relative 3D information could be distance or angle , which is inherently invariant to translation and rotation of input molecules . It is natural to consider such information in the spherical coordinate system ( SCS ) . We start by investigating the structure identification of 3D molecules in the SCS . For any point in the SCS , its location is specified by a 3-tuple ( d , θ , φ ) , where d , θ , and φ denote the radial distance , polar angle , and the azimuthal angle , respectively . When modeling 3D molecular graphs in the SCS , any atom i can be the origin of a local SCS , and d , θ , and φ naturally become the bond length , the angle between bonds , and the torsion angle , respectively . Thus , the relative location of each neighboring atom of atom i can be specified by the corresponding tuple ( d , θ , φ ) . Similarly , the relative location of each atom in the 3D molecular graph can be determined , leading to the identified structure , which is naturally invariant to translation and rotation of the input graph . The SCS can be easily converted from the Cartesian coordinate system , thus , the tuple ( d , θ , φ ) can be easily obtained . As in Fig . 1 , we use the chemical structure of the hydrogen peroxide ( H2O2 ) to show how d , θ , and φ are vital for the molecular structure identification . It is obvious that the structure is uniquely defined by the three bond lengths d1 , d2 , d3 , the two bond angles θ1 , θ2 , and the torsion angle φ . Note that the input may not contain all pairwise distances ( all possible bond lengths ) . This is because the atomic connectivity is usually based on real chemical bonds and cut-off distances . The cut-off distance is usually set as a hyperparameter . It is hard to guarantee that the cut-off is larger than any pairwise distance in a molecule . Hence , in this example , H-H bond length may not be considered as input if the cut-off is small . Setting a proper cut-off is even harder for other complicated and large molecules where a distance between two atoms could be large . In addition , considering all pairwise distances will cause severe redundancies , dramatically increasing the computational complexity . The model also easily gets confused by excessive noise , leading to unsatisfactory performance . From the perspective of completeness , using all pairwise distance is not capable of recognizing the chirality property . The chiral H2O2 molecules can not be distinguished even though the H-H bond is taken into account . To overcome the above challenges , we use a combination of distance , angle , and torsion for rigorous design and accurate learning . Apparently , the two O-H bonds can rotate around the O-O bond without changing any of the bond lengths and bond angles . In this situation , however , the torsion angle φ changes and the structure of the H2O2 varies accordingly . The importance of torsion angle has also been demonstrated in related research domains . Garg et al . ( 2020 ) formally shows that the torsion along with the port numbering can improve the expressive power of GNNs in distinguishing geometric graph properties , such as girth and circumference , etc . Other studies ( Ingraham et al. , 2019 ; Simm et al. , 2020 ) reveal that protein sequences and molecules can be accurately generated by considering the torsion in the given 3D structures . In this work , we propose SMP to systematically consider distance , angle , and torsion for approximately complete representation learning of 3D molecular graphs . Note that by using angle and torsion , SMP can easily recognize the chirality property . 3 SPHERICAL MESSAGE PASSING . 3.1 MESSAGE PASSING SCHEME . Currently , the class of message passing neural networks ( MPNNs ) ( Gilmer et al. , 2017 ) are one of the most widely used architectures for GNNs . Based upon the completeness analyses in Sec . 2 , we propose to perform message passing in the spherical coordinate system ( SCS ) , resulting in a novel and efficient scheme known as spherical message passing ( SMP ) . We show that message passing schemes used in existing methods , such as SchNet and DimeNet , are special cases of SMP . We first formally define a 3D molecular graph , which is usually represented as a 4-tuple G = ( u , V , E , P ) . The u ∈ Rdu is a global feature vector for the molecular graph G. V = { vi } i=1 : n is the set of atom features , where each vi ∈ Rdv is the feature vector for the atom i. E = { ( ek , rk , sk ) } k=1 : m is the set of edges , where each ek ∈ Rde is the feature vector , rk is the index of the receiver atom , and sk is the index of the sender atom for the edge k. P = { rh } h=1 : n is the set of 3D Cartesian coordinates that contains 3D spatial information for each atom . In addition , we let Ei = { ( ek , rk , sk ) } rk=i , k=1 : m denote the set of edges that point to the atom i , and Ni denote the indices of incoming nodes of atom i . The outputs after a message passing process include the updated global feature vector u′ ∈ Rdu , the updated atom features V ′ = { v′i } i=1 : n , and the updated edges E′ = { ( e′k , rk , sk ) } k=1 : m. An illustration of the message aggregation scheme for SMP is provided in Fig . 2 ( a ) . Apparently , the embedding of the atom rk is obtained by aggregating each incoming message ek . The message ek is updated based on Esk , the set of incoming messages pointing to the atom sk . Let q denote the sender atom of any message in Esk . Hence , we can define a local SCS , where sk serves as the origin , and the direction of the message ek naturally serves as the z-axis . We define a neighboring atom o of sk as the reference atom . Thus , the reference plane is formed by three atoms sk , rk , and o . For atom q , its location is uniquely defined by the tuple ( d , θ , φ ) , as shown in Fig . 4 ( a ) . Specifically , d determines its distance to the atom sk , θ specifies its direction to update the message ek . The torsion angle φ is formed by the defined reference plane and the plane spanned by sk , rk , and q . Intuitively , as an advanced message passing architecture in spherical coordinates for 3D graphs , SMP specifies relative location for any neighboring atom q by considering all the distance , angle , and torsion information , leading to more comprehensive representations for 3D molecular graphs . Generally , the atom sk may have several neighboring atoms , which we denote as q1 , ... , qt . It is easy to compute the corresponding bond lengths and bond angles for these t atoms . The SMP computes torsion angles by projecting all the t atoms to the plane that is perpendicular to ek and intersect with sk . Then on this plane , the torsion angles are formed in a predefined direction , such as the anticlockwise direction . By doing this , any atom naturally becomes the reference atom for its next atom in the anticlockwise direction . Notably , the sum of these t torsion angles is 2π . A simplified case is illustrated in Fig . 2 ( b ) . The atom sk has three neighboring atoms q1 , q2 , and q3 ; q3 is the reference atom for q1 , and they form φ1 ; q1 is the reference atom for q2 , and they form φ2 ; similarly , q2 is the reference atom for q3 , and they form φ3 . It is obvious that the sum of φ1 , φ2 , and φ3 is 2π . As the torsion is defined relatively , q1 can be picked arbitrarily , which will not affect the output of the message passing scheme , as we perform summation when aggregating information to sk from its neighbors q1 , q2 , and q3 . Notably , by designing each atom to be the reference atom of the next one in the predefined direction , invariance is effectively achieved because reference atom is naturally relative . In addition , our method computes torsion within edge-based 1-hop neighborhood . Even though a torsion angle involves four atoms , our design avoids the number of torsion angles to be exponential , but makes it the same as the number of neighboring atoms . Hence , it is efficient and will not cause time or memory issues . Formally , the proposed SMP can be defined in the SCS as e′k = ϕ e ( ek , vrk , vsk , Esk , ρ p→e ( { rh } h=rk∪sk∪Nsk ) ) , v′i = ϕ v ( vi , ρ e→v ( E′i ) ) , u ′ = ϕu ( u , ρv→u ( V ′ ) ) , ( 1 ) where ϕe , ϕv , and ϕu are three information update functions on edges , atoms , and the whole graph , respectively . ρe→v and ρv→u aggregate information between different types of geometries . Especially , in SMP , the 3D information in P is converted and incorporated to update each message ek . Hence , SMP employs another position aggregation function ρp→e for message update . Notably , the main difference between our SMP scheme defined in Eq . 1 and the GN framework in Battaglia et al . ( 2018 ) is the inclusion of 3D information P . In line with the research area described in Sec . 5.1.2 , we focus on such 3D information and develop a systematic solution to incorporate it completely and efficiently . Detailed description of these functions is given in Appendix A .

Summary: The paper proposes a message passing scheme using spherical coordinates. It is tested on three datasets of 3D moleclular graphs. The paper offer an in depth analysis of different aspects, with an extensive experimentation of the method.