Year: 2021

Abstract:
We consider representation learning of 3D molecular graphs in which each atom is associated with a spatial position in 3D. This is an under-explored area of research, and a principled message passing framework is currently lacking. In this work, we conduct analyses in the spherical coordinate system (SCS) for the complete identification of 3D graph structures. Based on such observations, we propose the spherical message passing (SMP) as a novel and powerful scheme for 3D molecular learning. SMP dramatically reduces training complexity, enabling it to perform efficiently on large-scale molecules. In addition, SMP is capable of distinguishing almost all molecular structures, and the uncovered cases may not exist in practice. Based on meaningful physically-based representations of 3D information, we further propose the SphereNet for 3D molecular learning. Experimental results demonstrate that the use of meaningful 3D information in SphereNet leads to significant performance improvements in prediction tasks. Our results also demonstrate the advantages of SphereNet in terms of capability, efficiency, and scalability.

heading: 1 INTRODUCTION
text: 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.

heading: 2 COMPLETE REPRESENTATIONS OF MOLECULES
text: 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 cannot 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.

heading: 3 SPHERICAL MESSAGE PASSING

heading: 3.1 MESSAGE PASSING SCHEME
text: 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.

heading: 3.2 COMPLETENESS VERSUS EFFICIENCY
text: The identification criteria described in Sec. 2 can fully determine the structure of a 3D molecule, but involves edge-based 2-hop information. Hence, the computational complexity is as sizeable as O(nk3), where n is the number of atoms, and k denotes the average number of neighboring atoms for each center atom. Unfortunately, such design can hardly generalize to large molecular graphs. To this end, we propose SMP as an efficient and scalable scheme to realize message passing in SCS. Our SMP only involves edge-based 1- hop information, thus the
time complexity is reduced to O(nk2). This enables the application of SMP to large molecules, like the newly released OC20 data (Chanussot et al., 2020). We rigorously investigate the completeness of SMP and show that it can distinguish even complex geometric properties such as chirality, as indicated by Fig. 3 (a) and Fig. 3 (b). As SMP uses the last atom as the reference atom (like q2 is the reference atom for q3 in Fig. 2 (b)) in a predefined direction, the relative order between adjacent atoms is considered while the absolute order is neglected. Hence, SMP can not distinguish the two molecules illustrated by Fig. 3 (b) and Fig. 3 (c). However, this scenario may not exist in nature. This is also demonstrated in experiments that our SMP achieves comparable performance with the complete representations, while the latter induces huge time complexity and severe memory issues.

heading: 3.3 RELATIONS WITH PRIOR MESSAGE PASSING METHODS
text: When developing message passing methods for 3D graphs, sphere message passing is an advanced scheme where the relative location of each atom is more specified. The development for 3D graphs
with relative information is still in early stage. To our best knowledge, there exist several notable methods in literature, and they can be viewed as special cases of the SMP, as they capture partial 3D position information. For example, the SchNet and PhysNet consider distance, and the DimeNet encodes directional information. Formally, these methods can be perfectly fit the Spherical scheme defined in Eq. 1. We describe the details of these methods in Appendix C. Notably, compared with prior models, SMP provides a rigorous justification on its completeness with failure cases clearly described. Importantly, SMP is developed based on the identification analyses of 3D molecular graphs. Hence, it aims at learning complete data representations for 3D molecular graphs, rather than simply including extra 3D information (like angle or torsion).

heading: 4 SPHERENET
text: The obtained 3-tuple (d, θ, φ) indicates the relative location of any atom in a 3D molecular graph. However, it cannot serve as the direct input to neural networks as it lacks meaningful representations. Essentially, molecules are quantum systems thus the representation design needs to follow physics laws. An important aspect is to choose appropriate basis functions that transform the 3-tuple (d, θ, φ) into physically-based representations. Several basis functions have been explored in Hu et al. (2021); Klicpera et al. (2020b), including MLP, Gaussian and sine functions, spherical Bessel basis, and spherical harmonics. Especially, spherical Bessel is shown to be the best basis for encoding distance, and spherical harmonics is the most appropriate one for encoding angle (Hu et al., 2021; Klicpera et al., 2020b). We denote the final representation as Ψ(d, θ, φ). Referring to theories in Griffiths & Schroeter (2018); Cohen et al. (2019); Klicpera et al. (2020b), one form of the representation can be expressed as Ψ(d, θ, φ) = jℓ ( βℓn c d ) Y mℓ (θ, φ), where jℓ(·) is a spherical Bessel function of order ℓ, Y mℓ is a spherical harmonic function of degree m and order ℓ, c denotes the cutoff, βℓn is the n-th root of the Bessel function of order ℓ. We also have ℓ ∈ [0, · · · , L− 1], m ∈ [−ℓ, · · · , ℓ] and n ∈ [1, · · · , N ]. L and N denote the highest orders for the spherical harmonics and spherical Bessel functions, respectively. They are hyperparameters in experimental settings. In addition, we can derive two simplified representations Ψ(d) and Ψ(d, θ) from Ψ(d, θ, φ).
Based upon the spherical message passing scheme described in Sec. 3 and physical representations, we build the SphereNet for molecular learning. Apparently, SphereNet can produce data representations that are both accurate and physically meaningful. By incorporating the position information in spherical coordinates, SphereNet also generates predictions invariant to translation and rotation of input molecules. Following the architecture design in the research line stated in Sec. 5.1.2, our network is composed of an input block, several interaction blocks, and an output block. For clear description, we assume the message ek for the edge k in Fig. 2 and Eq. (1) is the message for update. The update process and detailed architecture for the SphereNet are provided in Appendix B.

heading: 5 RELATED WORK

heading: 5.1 METHODS FOR 3D MOLECULAR GRAPHS

heading: 5.1.1 EQUIVARIANT GRAPH NEURAL NETWORKS
text: One research line for 3D molecular graphs is equivariant graph neural networks (EGNNs) including tensor field networks (TFNs) (Thomas et al., 2018), SE(3)-transformers (Fuchs et al., 2020), PaiNN (Schütt et al., 2021), NequIP (Batzner et al., 2021), etc. The raw input of these methods usually contains the absolute information, such as coordinates in the Cartesian coordinate system. In intermediate layers, absolute information could be decomposed into partial absolute information and partial relative information as needed. A simple example is that a vector can be decomposed into its orientation (absolute) and length (relative) (Thomas et al., 2018; Schütt et al., 2021). Apparently, network components of these methods should be carefully designed to be equivariant. The preliminary work like TFNs were developed for 3D point clouds. However, it is demonstrated that for molecules whose downstream tasks usually require the systems to be invariant, the complicated SE(3) group representations are not necessary but S2 representations are enough (Klicpera et al., 2021). Moreover, their performance on molecular tasks is not satisfactory.

heading: 5.1.2 INVARIANT GRAPH NEURAL NETWORKS
text: Another category of methods purely take relative 3D information as input, such as distances between atoms, angles between bonds, angles between planes, etc. Hence, the network is naturally invariant. The development of these methods is in early stage, and existing studies focus on leveraging different geometries. The SchNet (Schütt et al., 2017) incorporates the distance information during the information aggregation stage by using continuous-filter convolutional layers. The PhysNet (Unke & Meuwly, 2019) integrates both the atom features and distance information in the proposed interaction block. The DimeNet (Klicpera et al., 2020b) is developed based on the PhysNet and moves a step forward by considering directional information in the interaction block. The GemNet (Klicpera et al., 2021) is proposed recently for universal molecular representations. The OrbNet (Qiao et al., 2020) combines distance information with the atomic orbital theory to design important SAAO features as inputs to GNNs. Generally, the use of 3D position information usually results in improved performance. However, existing methods simply include additional geometries such as distance and angle, and there lacks a rigorous justification on how different geometries contribute to the information aggregation process. We conduct formal analysis and show that all the distance, angle, and torsion are necessary for 3D molecular identification, based on which we propose SphereNet to generate more powerful molecular representations.

heading: 5.2 METHODS FOR OTHER OBJECTS MODELED AS GRAPHS
text: Besides molecules, many other data objects are also represented as graphs, such as 3D point clouds (Guo et al., 2020; Simonovsky & Komodakis, 2017; Shen et al., 2018; Landrieu & Simonovsky, 2018) and meshes (Bronstein et al., 2021; De Haan et al., 2020; Perraudin et al., 2019). When modeling 3D point clouds as 3D graphs, points are represented as nodes and connections between points are directed edges. Existing methods mainly capture distance information from local neighborhood in 3D space. In DGCNN (Wang et al., 2019b), a novel layer namely EdgeConv is proposed to aggregate distance-based edges features for node learning. In Landrieu & Boussaha (2019), neighborhood radius along with spatial orientation is incorporated in local point embedding. The work Wang et al. (2019a) proposes a graph attention convolution for 3D point clouds. Generally, these methods can be fit into our message passing scheme defined in Eq. 1. The work De Haan et al. (2020) is an exemplary study that formulates meshes as graphs with considering geometrical information. The used convolutional kernel depends on the angle between a predefined reference edge and any other edge projected to the tangent plane for each vertex. It focuses on the design of gauge equivariance rather than the learning of complete geometry information. In this work, we study the complete learning of 3D molecules and leave extensive studies on other data types as future work.

heading: 6 EXPERIMENTAL STUDIES

heading: 6.1 EXPERIMENTAL SETUP
text: We apply our SphereNet to three benchmark datasets, including Open Catalyst 2020 (OC20) (Chanussot et al., 2020), QM9 (Ramakrishnan et al., 2014), and MD17 (Chmiela et al., 2017). Baseline methods include PPGN (Maron et al., 2019), SchNet (Schütt et al., 2017), PhysNet (Unke & Meuwly, 2019), Cormorant (Anderson et al., 2019), PaiNN (Schütt et al., 2021), NequIP (Batzner et al., 2021), MGCN (Lu et al., 2019), DimeNet (Klicpera et al., 2020b), DimeNet++ (Klicpera et al., 2020a), GemNet (Klicpera et al., 2021), CGCNN (Xie & Grossman, 2018), and sGDML (Chmiela et al., 2018). Detailed configurations of all the models used in the following sections are provided in the supplementary material. Unless otherwise specified, for all the baseline methods, we report the results taken from the referred papers or provided by the original authors. For the SphereNet, all models are trained using the Adam optimizer (Kingma & Ba, 2014). The optimal hyperparameters are obtained by grid search. Network configurations and search space for all models are provided in Appendix D. Code is contained in supplementary material.

heading: 6.2 OC20
text: The Open Catalyst 2020 (OC20) dataset is a newly released large-scale dataset for catalyst discovery and optimization (Chanussot et al., 2020). It comprises millions of DFT relaxations across huge
chemical structure space such that machine learning models can be fully trained. We focus on the IS2RE task in this work, and description of the data is provided in Appendix E. Results for CGCNN, SchNet, and DimeNet++ are provided in Chanussot et al. (2020). The original GemNet paper does not contain results on the OC20 dataset, We use the publicly available code from the OC Project website1 to produce results for GemNet-T. We report evaluation results of fixed epochs for SphereNet. Following a setting in Chanussot et al. (2020), we use the direct approach and all the training data for training models. The used metrics are the energy MAE and the percentage of Energies within a Threshold (EwT) of the ground truth energy. Table 1 shows that the SphereNet achieves the best performance on 3 out of 4 splits and the average in terms of energy MAE. For EwT, SphereNet is the best on all the 4 splits. Specifically, it reduces the average energy MAE by 0.019, which is 3.10% of the second best model. In addition, it improves the average EwT from 3.42% to 3.64%, which is a large margin considering the inherently low EwT values.
Notably, ForceNet (Hu et al., 2021) and GemNet (Klicpera et al., 2021) are recently proposed for quantum system learning. A prominent advantage for ForceNet is its high efficiency and scalability to large molecules. ForceNet focuses on S2EF thus there are no original results for the IS2RE task. However, DimeNet++ is slightly better than ForceNet in terms of performance, and our SphereNet outperforms DimeNet++ significantly. GemNet has two variants GemNet-T and GemNet-Q. GemNetT considers distance and angle information as input, and contains an effective architecture with novel network components, such as bilinear layers and scaling factors. We can see GemNet-T is similar as DimeNet++ in terms of performance. GemNet-Q is claimed to be able to capture universal representations of molecules. However, it considers edge-base 2-hop information and the time complexity is extremely high. It cannot be configured properly to run experiments on the large catalyst molecules.
1https://github.com/Open-Catalyst-Project/ocp

heading: 6.3 QM9
text: We apply the SphereNet to the QM9 dataset, which is widely used for predicting various properties of molecules. It consists organic molecules composed of up to 9 heavy atoms. Thus, this test can examine the power of the SphereNet for similar quantum chemistry systems. The dataset is original split into three sets, where the training set contains 110,000, the validation set contains 10,000, and the test set contains 10,831 molecules. For energy-related properties, the training processes use the unit eV. All hyperparameters are tuned on the validation set and applied to the test set. We compare our SphereNet with baselines using mean absolute error (MAE) for each property and the overall mean standarized MAE (std. MAE) for all the 12 properties. The comparison results are summarized in Table 2. SphereNets achieves best performance on 5 properties and the second best performance on 3 properties. It also improves the overall mean std. MAE of the QM9 dataset from 0.98 to 0.91 and sets the new state of the art. Notably, the most recent method PaiNN uses the same data splits as SphereNet in terms of sample numbers. Its final performance is the average of three different runs on three random splits. We follow such settings and run SphereNet on four properties including ϵHOMO, ϵLUMO, U0, and µ. The corresponding results are 22.9 ± 0.2, 18.8 ± 0.2, 6.28 ± 0.05, and 0.0243± 0.00, respectively. It is obvious that these results are highly close to those in Table 2, thus, we can draw consistent conclusions.

heading: 6.4 MD17
text: The MD17 dataset is used to examine the expressive power of SphereNet for molecular dynamics simulations. Following the settings in Schütt et al. (2017); Klicpera et al. (2020b), we train a separate model for each molecule to predict atomic forces. We use 1000 samples for training, and each of the eight molecules has both the validation and test sets. Note that all the baseline models employ a joint loss of forces and conserved energy during training. In the original SchNet (Schütt et al., 2017) and DimeNet (Klicpera et al., 2020b) papers, the authors set the weight of force over energy (WoFE) to 100, while the NequIP (Batzner et al., 2021) and GemNet (Klicpera et al., 2021) papers use a weight of 1000. As the WoFE tends to affect the force prediction significantly, we perform SphereNet with both WoFE values for fair comparisons. PaiNN (Schütt et al., 2021) uses neither 100 nor 1000 as WoFE, so we do not compare with it on MD17. The results for forces are reported in Table 3. Note that for Benzene, all models are evaluated on Benzene17, thus, the result for sGDML is 0.20 rather than 0.06 (Benzene18). We can observe from the table that when WoFE is 100 for all models, SphereNet consistently outperforms SchNet and DimeNet by largin margins. Notably, sGDML is one of the original work that created the MD17 dataset with carefully-designed features. Compared with sGDML, SphereNet performs better on four and worse on the other four molecules, which is similar to DimeNet. One reason is sGDML incorporates molecular symmetries to boost precision, and different molecules have different symmetries. However, sGDML has poorer generalization power to larger datasets without hand-engineered features. In addition, SphereNet achieves much
better overall std. MAE than sGDML. When using the same WoFE that is 1000, SphereNet achieves similar results with GemNet in spite of that GemNet-T is of high complexity and contains carefully designed network components for performance boost.

heading: 6.5 COMPLETENESS VERSUS EFFICIENCY
text: The message passing scheme Q-MP in GemNet represents the edge-based 2-hop geometric message passing and can generate complete representations of 3D molecular graphs. We study the capability and efficiency of the proposed SMP by comparing with Q-MP. Specifically, We use the same backbone network for these two MP methods for fair comparisons. We extensively use two backbones, which are the SphereNet backbone introduced in Sec. 4 and the GemNet backbone proposed in Klicpera et al. (2021). We conduct experiments on MD17, re-
porting performance and average running time for all the 8 molecules per epoch using the same computing infrastructure (Nvidia GeForce RTX 2080 TI 11GB). Results are shown in Table 4, from which we can observe that on either backbone network, SMP achieves very similar results with Q-MP. However, the time cost is much less than SMP, which indicates it is much more efficient than Q-MP. Based on analyses in Sec. 3.2, SMP can distinguish almost all molecular structures, and the failure cases may not exist in nature. Hence, SMP performs similarly with Q-MP even though the latter is complete theoretically but not scalable in practice. We further compare the efficiency between SphereNet and other models in terms of parameters and time cost in Appendix F. SphereNet uses similar computing budget with others but achieves the best performance.
6.6 ABLATION STUDY
The proposed SMP considers all the distance, angle, and torsion, leading to more powerful data representations. We investigate contributions of different geometries to demonstrate the advances of our SMP. We remove torsion information from SMP which we denote as “SMP w/o φ”; we further remove angle information which we
denote as “SMP w/o (θ, φ)”. The three message passing strategies are integrated to the same architecture with other network parts remaining the same. We evaluate these models on four molecules of MD17. Table 5 shows that SMP outperforms SMP w/o φ, and SMP w/o φ outperforms “SMP w/o (θ, φ)”. These results demonstrate the effectiveness of angle and torsion information used in the SMP. The best performance of SMP further reveals that SMP represents an accurate scheme for 3D graphs. In addition, we provide visualization results for SphereNet filters in Appendix G to further show that all the distance, angle, and torsion information determine the structural semantics of filters.

heading: 7 CONCLUSIONS
text: 3D information is important for molecules but there lacks a principled message passing framework to consider it. We first propose the spherical message passing as a unifying and efficient scheme that can achieve approximately complete representations of molecules without increasing computing budget. Based on SMP and meaningful physical representations, SphereNet is presented, and experiments on various types of datasets demonstrates its capability, efficiency, and scalibility.
REPRODUCIBILITY STATEMENT
Detailed experimental setup is provided in Appendix D. Implementation hyper-parameters of SphereNet on all the three datasets OC20, QM9, and MD17 are given in Table 6, Table 7, and Table 8, respectively. Code is attached in supplementary material.

heading: C RELATIONS WITH PRIOR MESSAGE PASSING METHODS
text: Our SMP is a specific architecture for 3D graphs and is formally defined in Eq. (1). Especially, the message passing schemes used in some existing models can be viewed as special cases of SMP as they only encode partial 3D information. In this section, we clearly give the used functions as well as their inputs for each existing method. Basically, we describe how each method realizes Eq. (1).
C.1 SCHNET, SCHÜTT ET AL. (2017)
In SchNet, the used aggregation function to encode 3D position information is ρp→e ({rh}h=rk∪sk) = Ψ (∥rrk − rsk∥), which converts the position information to an embedding of distance with radial basis functions. In addition to the ρp→e function, the ϕe function used is NN (NN (vrk)⊙ NN (Ψ (∥rrk − rsk∥))), where NN denotes a neural network and ⊙ denotes the element-wise multiplication. The ρe→v function is ∑ (e′k,rk,sk)∈E′i e′k. The ϕ v function is
vi + ∑
(e′k,rk,sk)∈E′i e′k. The global feature u is updated based on the final node features V T and the
function is ϕu = ∑ i=1:n NN ( vTi ) . Formally, the update process is expressed as
e′k =ϕ e (vrk , ρ p→e ({rh}h=rk∪sk)) =ϕe (vrk ,Ψ(∥rsk − rrk∥)) =NN (NN (vrk)⊙ NN (Ψ (∥rrk − rsk∥))) , v′i =ϕ v (vi, ρ e→v (E′i))
=ϕv vi, ∑ (e′k,rk,sk)∈E′i e′k  =vi +
∑ (e′k,rk,sk)∈E′i e′k,
u =ϕu ( ρv→u ( V T )) =
∑ i=1:n NN ( vTi ) .
(2)
C.2 PHYSNET, UNKE & MEUWLY (2019)
PhysNet uses distance between atoms as an important feature and proposes more powerful neural networks for chemical applications. The position aggregation function is ρp→e ({rh}h=rk∪sk) = Ψ (∥rrk − rsk∥), where Ψ is any radial basis function with a smooth cutoff. For the information update functions, the ϕe function is σ (W1)σ (vsk) ⊙ W2Ψ(∥rrk − rsk∥), the ϕv
function is NN ( W3 ⊙ vi + NN ( σ (W4)σ (vi) + ∑ (e′k,rk,sk)∈E′i e′k )) and the ϕu function is
u+ ∑
i=1:n NN (v ′ i). Here NN denotes a neural network, W1,W2,W3,W4 are learnable weight
matrices, σ is an activate function, and ⊙ denotes the element-wise multiplication. PhysNet is expressed as
e′k =ϕ e (vsk , ρ p→e ({rh}h=rk∪sk)) =ϕe (vsk ,Ψ(∥rrk − rsk∥)) =σ (W1)σ (vsk)⊙W2Ψ(∥rrk − rsk∥) , v′i =ϕ v (vi, ρ e→v (E′i))
=ϕv vi, ∑ (e′k,rk,sk)∈E′i e′k  =NN W3 ⊙ vi + NN σ (W4)σ (vi) + ∑
(e′k,rk,sk)∈E′i
e′k   ,
u′ =ϕu (ρv→u (V ′) ,u) =u+ ∑ i=1:n NN (v′i) .
(3)
C.3 DIMENET, KLICPERA ET AL. (2020B)
DimeNet explicitly considers distances between atoms and directions of directed edges. The aggregation functions on the position information is ρp→e = (Ψ (d) ∥Ψ(d, θ)), where ∥ denotes concatenation, Ψ(d) and Ψ(d, θ) are the same basis functions used in SphereNet as introduced in Sec. 4. Specifically, Ψ(d) denotes the representation of the distance based on spherical Bessel function, and Ψ(d, θ) denotes the representation of distance and angle based on spherical Bessel function and spherical harmonics. For other functions, the ϕe function used is e′k = ( e′k,1∥e′k,2 ) with
e′k,1 = NN ( ek,1 + NN ( σW1ek,1 + ∑ (ej ,rj ,sj)∈Esk W2Ψ(dj , θjk) (W3Ψ(dj)⊙ σW4ej,1) ))
and e′k,2 = W5Ψ(dj) ⊙ e′k,1, where NN denotes a neural network, W1,W2,W3,W4,W5 are different weight matrices, and σ is an activation function. The ρe→v function is ∑ (e′k,rk,sk)∈E′i e′k,2
and the ϕv is NN (∑
(e′k,rk,sk)∈E′i e′k,2
) . The ρv→u is ∑ i=1:n v ′ i and the ϕ u is u+ ∑ i=1:n v ′ i. Note
that ρp→v , ρp→u, ρe→u functions are not required in DimeNet. The whole model is expressed as
ek =(ek,1∥ek,2) , ρp→e =(Ψ (d) ∥Ψ(d, θ)) ,
e′k,1 =ϕ e ( ek, Esk , ρ p→e ( {rh}h=rk∪sk∪Nsk )) =NN ek,1 + NN σW1ek,1 + ∑
(ej ,rj ,sj)∈Esk
W2Ψ(dj , θjk) (W3Ψ(dj)⊙ σW4ej,1)  , e′k,2 =W5Ψ(dj)⊙ e′k,1, v′i =ϕ v (ρe→v (E′i))
=NN  ∑ (e′k,rk,sk)∈E′i e′k,2  , u′ =ϕu (u, ρv→u (V ′))
=u+ ∑ i=1:n v′i.
(4)

heading: D EXPERIMENTAL SETUP
text: For all the models used in three datasets, we set input embedding size = 256 and output embedding size = 64 for both LB2 and LB blocks. For each separate model, we first perform warmup on initial learning rate. Then two learning rate strategies, including ReduceLROnPlateau and StepLR, are used for training. For StepLR, the learning rate is decayed by the decay ratio every fixed epochs represented as step size. We do not use weight decay or dropout for all models. Some hyperparameters are fixed values, and some are tuned by grid search. Values/search space of hyperparameters for OC20, QM9, and MD17 are provided in Table 6, Table 7, and Table 8, respectively. Optimized hyperparameters are tuned on validation sets and applied to test sets for QM9 and MD17. For OC20, optimized hyperparameters are obtained on the ID split within max epochs, and then applied to the other three splits. Pytorch is used to implement all methods. For QM9 and MD17 datasets, all models are trained using one NVIDIA GeForce RTX 2080 Ti 11GB GPU. For the OC20 dataset, all models are trained using four NVIDIA RTX A6000 48GB GPUs.

heading: E OC20 DATA DESCRIPTION
text: There exist three tasks including S2EF, IS2RS, and IS2RE. In this work, we focus on IS2RE that predicts structure’s energy in the relaxed state. It is the most common task in catalysis as relaxed energies usually influence the catalyst activity. The dataset for IS2RE is originally split into training/validation/test sets. There are 460,318 structures in the training dataset in total. The test label is not publicly available. Performance is evaluated on the validation set, which has four splits including In Domain (ID), Out of Domain Adsorbates (OOD Ads), Out of Domain catalysts (OOD cat), and Out of Domain Adsorbates and catalysts (OOD Both), where numbers of structures are 24,943, 24,961, 24,963, 24,987, respectively. The average number of atoms per structure is 77.75.

heading: F EFFICIENCY STUDY OF SPHERENET
text: We study the efficiency of SphereNet by comparing with other models regarding number of parameters and time cost per epoch using the same computing infrastructure (Nvidia GeForce RTX 2080 TI
11GB). Experiments are conducted on the property U0 of QM9 and results are shown in Table 9. It is obvious that SphereNet uses similar computational resources as DimeNet++ and GemNet-T, and is much more efficient than DimeNet. The main reason could be we develop an efficient way to compute torsion, as introduced in Sec. 3 and Fig. 2 (b). Moreover, GemNet-Q cannot run on QM9 using the infrastructure as mentioned above.

heading: G SPHERENET FILTER VISUALIZATION
text: We visualize SphereNet filters from a learned SphereNet model. Specifically, we port learned weights for the block LB2 after the torsion embedding Ψ(d, θ, φ) in Fig. 4. For each location represented by a tuple (d, θ, φ), the initial embedding size is N2 × L. The computation for the above LB2 is W1 (W2Ψ(d, θ, φ)), which results in the new embedding size of 64 for each location (d, θ, φ). We then perform sampling on locations in 3D space for visualizing weights as SphereNet filters. The visualization results are provided in Fig. 6. We set sampling rate in the torsion direction to be π/4, thus, there are eight samples in the torsion direction. There are totally 64 elements for each location, and we randomly pick 6 elements. Apparently, among the distance, angle and torsion, considering any one when fixing the other two, the structural value of filters will be different when the one of interest changes. It essentially shows that all the distance, angle, and torsion information determine the structural semantics of filters. This further demonstrates that SMP enables the learning of different 3D information for improved representations.

Review summaries:
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.

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.

This work aims at the representation learning of 3D molecular graphs and proposes a principled message passing framework to unify existing works. Besides, a novel message passing approach is presented by incorporating 3D geometric information including distance, angle, and torsion. In experiments, the performance of the proposed model is validated on three benchmark datasets.