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: This paper provides a novel geometric GNN: SMP. SMP is based on the SCS, and can reduce the complexity by an order of magnitude. Besides, it is a unified framework that can cover some mainstream geometric GNN methods. The empirical results support the effectiveness of SMP.