ComENet: Towards Complete and Efficient Message Passing for 3D Molecular Graphs
Anonymous Author(s)
AffiliationAddress email
Abstract
Many real-world data can be modeled as 3D graphs, but learning representationsthat incorporates 3D information completely and efficiently is challenging.
Existingmethods either use partial 3D information, or suffer from excessive computationalcost.
To incorporate 3D information completely and efficiently, we propose a novelmessage passing scheme that operates within 1-hop neighborhood.
Our methodguarantees full completeness of 3D information on 3D graphs by achieving globaland local completeness.
Notably, we propose the important rotation angles to fulfillglobal completeness.
Additionally, we show that our method is orders of magnitudefaster than prior methods.
We provide rigorous proof of completeness and analysisof time complexity for our methods.
As molecules are in essence quantum systems,we build the complete and efficient graph neural network (ComENet) by combingquantum inspired basis functions and the proposed message passing scheme.
Experimental results demonstrate the capability and efficiency of ComENet, especiallyon real-world datasets that are large in both numbers and sizes of graphs.
Introduction
In machine learning, structured objects such as molecules [14, 50], proteins [9, 16, 33, 37, 26],materials [51] and quantum systems [31] are usually modeled as graphs.
Original modeling showsbasic connections between units and the resulted data type is known as 2D graphs.
Accordingly, 2Dgraph neural networks (GNNs) have been intensively studied [14, 21, 19, 20, 52, 56, 32] and themessage passing scheme [22, 7, 49] is shown effective to realize 2D GNNs.
However, in the modernmachine learning era, it is increasingly accepted that modeling real-world data like molecules as2D graphs leads to inborn defects for succeeding learning models.
In practice, 3D information iscrucial, based on which some important geometries can be derived, such as chemical bond lengths inmolecular modeling.
This essentially raises the need of a new data type, known as 3D graphs.
Formally, a 3D graph contains the original 2D graph as well as Cartesian coordinates for all nodes.
Inthis work, we follow invariant 3D GNNs [41, 29, 28, 43, 34, 30] that take relative 3D information, likedistances and angles, as inputs to networks.
Such relative geometries are naturally SE(3)-invariant.
The main challenge comes from what geometries should be computed such that 3D information isincorporated completely.
A most recent work SphereNet [34] shows that distance, angle, and torsioninformation is necessary to incorporate more comprehensive 3D information.
However, SphereNet isonly complete in local neighborhood, failing to achieve global completeness and distinguish a widerange of molecular structures such as conformers.
In addition, existing methods exhibit excessivetime complexity of O ( nk 2 ) or even O ( nk 3 ) , severely preventing their scalability in real-worldapplications. Here n and k denote the number of nodes and the average degree in a 3D graph.
In this work, we propose ComENet as a complete and efficient graph neural network for 3D moleculargraph learning.
We first formally provide the definition of completeness .
Intuitively, a geometric transformation is considered as complete if it generates distinct representations for any two different3D graphs.
Based on this, we propose a novel message passing scheme by faithfully fulfilling globalcompleteness via the important rotation angles.
In addition, we design novel strategies to achievelocal completeness, largely reducing the computing complexity to O ( nk ) .
To elucidate the meritsof the proposed methods, we provide rigorous proof of geometric completeness achieved by ourmethod.
Combining the novel message passing scheme and quantum inspired features, ComENet isdeveloped for 3D molecular graphs.
We apply ComENet to two large-scale datasets including OC20and Molecule3D, and a commonly used dataset QM9.
Experiments show that ComENet performssimilar to existing best methods, but accelerates the training and inference by 6-10 times on variousdatasets.
We summarize the contributions of ComENet as below .
(i).
To our best knowledge, it isthe first rigorously complete pipeline for 3D molecular graph learning.
Theoretically, it is guaranteedto incorporate 3D information completely without information loss.
Practically, it can distinguish allmolecular structures in nature.
(ii).
It is highly efficient.
The message passing is shown to be ordersof magnitude faster than existing methods in terms of time complexity.
(iii).
The great capabilityand efficiency of ComENet allow its scalability to real-world molecule datasets that are large in bothnumbers and sizes of graphs.
(iv).
It achieves similar or better performance compared with existingmethods, and dramatically accelerates the training and inference by 6-10 times.
Relations with Prior Work.
SphereNet [34] is a recent method that achieves local completenesswith a complexity of O ( nk 2 ) .
In summary, SphereNet is not complete and not efficient enough forprocessing large-scale molecular graphs.
However, ComENet is complete and much more efficientwith a complexity of O ( nk ) .
Technically, a primary difference is that, ComENet proposes to use theimportant rotation angles to achieve global completeness.
As a result, ComENet is provably completeas shown in Sec. 2.2.
Practically, ComENet is able to identify a wide range of real-world structuressuch as conformers, achieving completeness at the conformer level, as introduced in Sec. 2.3.
Inaddition, ComENet and SphereNet both use ( d, θ, ϕ ) in the spherical coordinate system (SCS) toobtain local completeness.
Indeed, the fact that ( d, θ, ϕ ) can specify the location of a point in SCS iswidely known.
The key difference lies in that, SphereNet operates within 2-hop neighborhood, while
ComENet operates within 1-hop neighborhood.
When coupled with rotation angles, ComENet canachieve provable completeness with a reduced complexity of O ( nk ) .
This different message passingfor local completeness in ComENet entails many differences with SphereNet, including buildingcoordinate systems, defining z -axis, choosing reference nodes, and computing ( d, θ, ϕ ) , as detailedin Sec. 2.4.
All the computing procedures for ComeNet are described in detail in Algorithm 1 of
Appendix A.1.
The Proposed Message Passing Scheme2.1 Notations & Definitions
We first formally define notations and the concept of completeness used in this paper.
Notations.
A 3D graph G can be represented as G =
( V, A, P ) .
The node feature matrix V =[ v 1 , v 2 , ··· , v n ] T ∈ R n × d v with each v i ∈
R d v .
The adjacency matrix A ∈ R n × n , based on which weadditionally define there is an edge e ij if A
[ i ][ j ] = 1 .
The position matrix P =
[ p 1 , p 2 , · · · , p n ] T ∈ 75 R n ×, where p i =
( x i , y i , z i ) ∈ R 3 is the position vector for node i given in the Cartesian coordinatesystem (CSC).
The relative position p ij of node j to node i is defined as p ij = p
j − p
i .
Particularly,throughout this paper, we define k as the average degree for G .
We then formally define completeness given a geometric transformation T .
In particular, we aim atincorporating 3D information in 3D molecular graphs.
Hence, our definition of completeness is setfrom the geometric view.
Generally, T maps a 3D graph G =
( V, A, P ) to a geometric representationwith size m × d .
T can be different dependent on different methods, resulting in different m ∈ N + ord ∈ N + .
For example, SchNet generates a distance for each edge, thus, it maps G to a representationwith size m × d , where m is the number of edges in G and d = 1 .
However, such geometrictransformation is not complete.
We provide the definition of completeness as below:
Definition 1 (Completeness) .
For two 3D graphs G 1 = ( V, A, P 1 ) and G 2 = ( V, A, P 2 ) , a geometrictransformation T : ( R n × d v , R n × n , R n × 3 ) (cid:55)→ R m ×
d is considered as complete when
T ( G 1 ) =
T ( G 2 )
if and only if
[Equation]
89 Intuitively, a complete geometric transformation T can distinguish any two different 3D graphs.
Thisis to say, as long as two 3D graphs differ in structures, their outputs from T would be different.
2.2 Global Completeness via Rotation Angles
[Figure]
Existing studies focus on the complete representation learning of local neighborhood.
Earlier methodslike SchNet and DimeNet cannot achieves local completeness.
In a more recent method SphereNet,completeness is guaranteed within local neighborhood, but fails to hold in the whole 3D graph.
Inthis section, we move a step forward to formally fulfill global completeness for a given 3D moleculargraph.
Particularly, for the purpose of clear illustration, we safely assume local completeness isalready obtained by exiting methods like SphereNet.
[Figure]
Without loss of generality, we
[Figure]
start by describing our method to
[Figure]
attain full completeness within 2-
[Figure]
hop neighborhood, as illustratedin Fig. 1(a).
Formally, for a center node i , we let N i and N 2 i denote two sets of indices of i ’s 1hop and 2-hop neighboring nodes, [Figure] respectively.
We also define any [Figure] node i and its 1-hop neighbor-
[Figure]
hood as a local structure. Then [Figure] the whole 2-hop neighborhood of
[Figure] node i can be viewed as 1 + |N i |local structures centered in i and N i , defined as L i and L ij,j ∈N i , respectively.
Apparently, each local structure L ij,j ∈N i shares the 113 common edge e ij with the local structure L i .
Given the complete representation for each localstructure, for the structure L
i ∪ L ij , the only remaining degree of freedom is the rotation angle ofedge e ij , denoted as τ ij .
With the rotation angles for all the |N i | common edges specified, we canobtain a complete representation for 2-hop neighborhood of node i .
Achieving completeness beyond2-hop neighborhood is similar.
Overall, after considering rotation angles, the global completenesscan be easily guaranteed when it gradually generalizes from n - to ( n + 1) -hop neighborhood.
Essentially, each edge in an input graph can be treated as a common edge between different localstructures.
We reveal how to compute the rotation angle τ ij for each edge e ij in Fig. 1(b).
Specifically,we choose two reference nodes whose indices are f i \ j and f j \ i for node i and j , respectively.
f i \ jdenotes the index of i ’s nearest neighboring node except j , and f j \ i denotes j ’s nearest neighboringnode except i .
Then τ ij for edge e ij is the angle from the plane formed by f i \ j , i, j to the planeformed by i, j, f j \ i .
As analyzed previously, the global conformation of the input graph can beidentified based on all local structures and rotation angles.
As a result, given fixed local structures,the global completeness the input 3D molecular graph is fulfilled by additionally considering therotation angle of each edge, as introduced in this section.
2.3 Rotation Angles for Conformer Identification
The proposed rotation angles in Sec.
2.2 play a crucial role in identifying some important molecularstructures, such as conformers.
In nature, a real molecule exits as an ensemble of interconventing 3Dstructures, known as conformers [5, 4, 18].
Different conformers posses the same 2D molecular graph,but differ in 3D structures.
Generally, a conformer for a molecule exits with a certain probability andmay exhibit distinct properties [5, 4].
As shown in Fig.
2(a), different conformers for the moleculebutane (C 4 H 10 ) show varying conformation energy.
To this end, it is important to design complete3D GNNs for identifying molecules at the conformer level.
From the geometry perspective of view,a conformer distinguishes itself from others mainly through varying rotation angles of chemicalbonds [18].
As shown in Fig. 2(b), given the fixed ethyl (-C 2 H 5 ) of both sides, the only degree offreedom is the rotation angle of the C-C bond.
In literature, the ethyl is formulated as a local 3D graph.
Existing studies focus on the complete representation of such local structures, failing to identify thewhole 3D graph globally.
Essentially, they can only distinguish different molecules, trying to achieve 𝐂 𝐂𝐇𝐇 𝐇 0. 180° 120° 60° 0° Energy/eV
A C B D 𝐂𝐇𝐂 𝐇 𝐇 𝐂 𝐂𝐇𝐇 𝐇 𝐇 𝟑 𝐂 𝐂 𝐇 𝐇 𝐂 𝐂𝐇𝐇 𝐇 𝐇 𝟑 𝐂 𝐂 𝐇 𝐇 𝐂 𝐂𝐇𝐇 𝐇 𝐇 𝟑 𝐂 𝐂 𝐇
A 𝐂 𝐂𝐇𝐇 𝐇 𝐂 𝐂𝐇𝐇 𝐇
C
B
D 𝐇 𝐂 𝐂𝐇𝐇 𝐇 𝐂 𝐇 𝐇 𝐇 𝟑 𝐂 𝐂 𝐂𝐇𝐇 𝐇 𝐂 𝐇 𝐇 𝐇 𝟑 𝐂 𝐂 𝐂𝐇𝐇 𝐇 𝐂 𝐇 𝐇 𝐂𝐇
[Figure]
completeness at the molecule level rather that the finer comformer level.
By integrating the rotationangles as in Sec.
2.2 into the message passing scheme, our methods can fulfill rigorous completenessat the conformer level and can distinguish all conformers in nature.
2.4 Local Completeness with Improved Efficiency
Formally, each node and its local neighborhood can be viewed as a local structure.
Existing studiesfocus on the learning of local structures and SphereNet achieves local completeness.
However,
SphereNet induces the complexity of O ( nk 2 ) , restricting its scalability on large molecules in practice.
Here, we design a novel strategy to guarantee local completeness with the computing cost of O ( nk ) .
Specifically, we follow SphereNet and perform on the spherical coordinate system.
It is commonlyknown that the location of each node can be completely determined using the tuple ( d, θ, ϕ ) in
SCS.
SphereNet employs the directional message passing (DMP) fashion that operates within 2hop neighborhood.
It first updates messages over edges thus the center edge is z -axis in SCS.
Fornode i , the computing of the tuple ( d, θ, ϕ ) involves 2-hop information.
However, we view 1-hopneighborhood as a local structure, which requires all strategies in our local completeness to bedifferent from SphereNet, including building local coordinate systems, defining z-axis, pickingreference nodes, and computing ( d, θ, ϕ ) .
First, we build a light local coordinate system for any nodei ’s corresponding local structure.
Similarly, the center node i serves as the origin.
Then z -axis isdefined as the direction from i to its nearest neighbor f i , and xz -plane is further formed by z -axis andi ’s second nearest neighbor s i .
Finally, the tuple ( d, θ, ϕ ) is computed within 1-hop neighborhoodwith a complexity of O ( nk ) .
Particularly, we analyze efficiency versus model expressiveness in
Sec. 3.2.
We show that compared with the DMP fashion used by SphereNet, our method operatingwithin 1-hop neighborhood hurts the model expressiveness a bit by largely improves the efficiency.
2.5 Message Passing Scheme
Based upon global completeness achieved in Sec. 2.2 and improved local completeness introduced in
Sec. 2.4, the complete geometric transformation T required by Def. 1 should be formulated basedon a 4-tuple as ( d, θ, ϕ, τ ) .
Specifically, we build such transformation within 1-hop neighborhood,and a 4-tuple is computed for each edge.
Hence, for a 3D graph G =
( V, E, P ) , the full expressionfor T is T ( G ) =
[( d ij , θ ij , ϕ ij , τ ij )] i =1 ,...,n ; j ∈N i ∈ R m × 4 , where m is the number of edges in G .
Especially, as T converts absolute Cartesian coordinates in P to relative information, it is naturally
SE(3)-invariant as required in Def. 1.
Thus, we formally build our message passing scheme as
[Equation]
where g and f can be implemented by neural networks or mathematical operations.
Intuitively, ourmessage passing is established in 1-hop local neighborhood and all edges connecting regions beyond.
Essentially, d ij , θ ij , and ϕ ij specify the 1-hop local neighborhood, and τ ij determines the orientationof the local neighborhood.
By doing this, the complete representation for a whole 3D moleculargraph is eventually achieved.
The formulas of computing of d ij , θ ij , and ϕ ij , and τ ij are shown in
Algorithm 1 in Appendix A.1, along with detailed description of the complexity of O ( nk ) .
Merits of Our Methods
3.1 Geometric Completeness
Proposition 1.
For a strongly connected 3D graph G =
( V, E, P ) , its geometric transformation
T ( G ) =
[( d ij , θ ij , ϕ ij , τ ij )] i =1 ,...,n ; j ∈N i is complete.
Proof.
We employ mathematical induction and assume the number of nodes in a 3D graph is n .
Base case: It is obvious that the 3D structure of G can be identified when n = 1 , 2 .
Hence, we letn = 3 be the base case, where the completeness can be achieved by only considering d and θ in T .
Inductive hypothesis: The claim that T is complete holds for the node numbers of n up to k ≥ 3 .
Inductive step: Let n = k + 1 .
Without loss of generality, among the existing k nodes, we safelyassume i and its neighboring nodes
c k =1 , 2 ,... form the local region of interest.
Then j is the index ofthe newly ( k + 1) -th node connected to the center node i .
To show global completeness of the wholegraph, based on Def. 1, we only need to prove that the relative position of the new node j is uniquelydetermined given T .
We propose the following lemma for this as
Lemma 1.
Assume a strongly connected 3D graph G =
( V, E, P ) with more than 2 nodes is fullyidentified.
If a new node j is connected to a node i of G following the geometric transformation
T ( G ) =
[( d ij , θ ij , ϕ ij , τ ij )] i =1 ,...,n ; j ∈N i , then p ij is uniquely determined.
The proof of Lemma 1 is provided in Appendix A.2.
With Lemma 1 successfully proved, we showthat such geometric transformation T can determine a unique 3D graph.
Hence, the if condition in
Def. 1 holds.
In addition, as T renders purely relative 3D information like distance and angle, it’snaturally SE(3) invariant.
Hence, the only if condition in Def. 1 holds.
Overall, based on Def. 1, wecomplete the proof of Prop. 1.
Intuitively, since a molecular graph is strongly connected, two arbitrary nodes are connected by atleast one path.
Hence, starting from the existing structure, we can restore the relative position ofany new node step by step along a path with finite length.
As a result, geometric completeness forthe whole 3D graph with any number of nodes can be guaranteed in our message passing scheme.
Theoretically, 3D information of the input 3D molecular graph is fully captured without informationloss.
In practice, our method can distinguish all structures in nature.
3.2 Efficiency
Efficiency versus Model Expressiveness: Our method induces the complexity of O ( nk ) by operatingwithin 1-hop neighborhood.
Existing methods, such as like DimeNet and SphereNet, employ the DMPfashion that update edges within 2-hop neighborhood, inducing the complexity of O ( nk 2 ) .
Notably, DMP [46, 54, 29, 34] incorporates 2-hop information in one single layer but ( n + 1) -hop whenstacking n layers.
For each center edge In DMP, stacking n DMP layers incorporates informationfrom its n -hop neighboring edges. To further aggregate information from edges to nodes, studieslike Stokes et al. [46], Yang et al. [54] use an aggregation layer after the n -th DMP layer while
Klicpera et al.
[29], Liu et al.
[34] add such aggregation layer after each DMP layer.
Then each nodeis updated based on the nodes within ( n + 1) -hop edges.
Thus, stacking n DMP layers integrates( n + 1) -hop information.
However, stacking n proposed message passing layers is already able toincorporate information from n -hop neighborhood.
Obviously, compared with a network containingseveral DMP layers, a network with the same number of our proposed message passing layers merelyhurts the model expressiveness a bit, but significantly improves the model efficiency.
Efficiency via Less Torsion Angles: In addition, our method achieves completeness by computing
O ( n ) torsion angles, which is efficient especially compared with methods including Klicpera et al.[30], Adams et al. [1], Ganea et al.
Given a scenario where a global region is the union of twolocal regions with n 1 and n 2 nodes.
As introduced in Sec. 2.2 and Sec. 2.4, our method computes thesame number of torsion angles as nodes in a local region, then employs a rotation angle (torsion angleessentially). Hence, the number of torsion angles that our method computes is n 1 + n 2 + 1 . Ganeaet al. [18] computes a torsion angle based on one pair of nodes, each of which is from a separate localregion.
Hence, the number of torsion angles needed is n 1 × n 2 .
Apparently, our method reduces thenumber of torsion angles from O ( n 2 ) to O ( n ) , which is significant considering the computing oftorsion is excessively expensive.
[Figure]
Embedding
Interaction
Interaction
Interaction
Interaction
Self-Atom LocalConv GlobalConv MLP
Down-Project
I n t e r a c t i o n
B l o c k
TBF + SBF Concat MLP 𝑑𝑑 , 𝜃𝜃 , 𝜙𝜙 𝑑𝑑 , 𝜏𝜏 𝐯𝐯 𝑧𝑧
Output Pooling 𝑃𝑃 [Figure]
Based upon the message passingscheme introduced in Sec. 2.5, we propose the complete and efficient graphneural network (ComENet) as shownin Fig. 3.
Generally, ComENet consists of an embedding layer, multipleinteraction layers, a self-atom layer,and a pooling layer.
To be in linewith the message passing scheme in Sec.
2.5, we take the updating processfor node i ’s feature vector v i as anexample to describe the network.
Inparticular, we omit all indices in Fig.
3and below for clear presentation.
Embedding Layer converts eachatom type z to an initial node featurevector via learnable atom type embeddings [41, 29].
Interaction Layer updates each nodefeature vector v based on features ofthe neighboring nodes and the corresponding 3D information in P . Firstly,it converts 3D information in P to a set of geometries based on the proposed complete geometrictransformation T and message passing scheme. Since distance is the most important geometry, wealso consider d in global representation and split the output of T into two tuples ( d, θ, ϕ ) and ( d, τ ) ,for local and global representations, respectively. Importantly, the tuples ( d, θ, ϕ ) and ( d, τ ) cannot serve as immediate inputs to the network.
Theyneed to be transformed into physically meaningful vectors based on quantum-based basis functions. As in previous studies [25, 29, 34, 30], we test different basis functions including MLP, Gaussianand sine functions, spherical Bessel functions, and spherical harmonics. We found spherical Besseland spherical harmonics perform best. Formally, the basis function for tuple ( d, θ, ϕ ) is TBFj ℓ (cid:16) β ℓn c d (cid:17)
Y mℓ ( θ, ϕ ) , where j ℓ ( · ) is a spherical Bessel function of order ℓ , Y mℓ is a spherical harmonicfunction of degree m and order ℓ , c is the cutoff, and β ℓn is the n -th root of the Bessel function oforder ℓ . The basis function for tuple ( d, τ ) is SBF j ℓ (cid:16) β ℓn c d (cid:17) Y 0 ℓ ( τ ) . These two basis functions arealso used in SphereNet [34] and GemNet [30].
The two physically meaningful vectors from TBF and SBF are then imported into a local convolutionlayer and a global convolution layer, respectively.
For both convolution layers, we use the GraphConv [38] implemented in the PyTorch Geometric library [15].
The vectors from the basis functionsare used as edge weights in the convolution layers.
The outputs of local and global convolution layersare concatenated to generate a new node feature vector. Then the concatenated vector is forwardedinto several linear layers to generate the updated feature vector v ′ . Self-Atom and Pooling layers. The self-atom layer is used to update each node feature and projectthe feature dimension into 1, similar to the atom-wise layer used in SchNet.
The pooling layer is asum-pooling performing on all node features to obtain final predictions.
Related Work
We consider how to represent 3D information in 3D graphs for facilitating molecular learningtasks [3, 10].
One category of methods are equivariant 3DGNNs that directly use coordinates in the CSC as inputs to networks [47, 2, 17, 42, 8].
These methods are efficient and do not need to considernodes beyond 1-hop neighborhood. However, they suffer from several setbacks.
Firstly, each networkcomponent needs to be carefully designed to be rotation equivariant of input graphs.
For example,
NequIP [8] employs equivariant convolutions based on spherical harmonics and Clebsch-Gordon transforms.
Non-linearities also need careful design and common techniques like ReLu cannot beused.
Secondly, the reason why this category of methods are efficient lies in that they only use thetype-1 basis in Spherical harmonics, which is an approximation essentially.
It is proved in tensor fieldnetworks [47] that theoretically, l is infinite in terms of type- l basis.
In practice, l should be at least 2for achieving satisfactory performance.
Type-1 basis essentially coerces the conv kernel to be in anarrow learning space, which imposes a hard constraint on the network capability.
However, whenusing type-2 basis, the conv kernel can be more expressive while the efficiency issue would emergeas a new bottleneck.
Thirdly, the performance of such equivariant GNNs is shown to be worse thanthe other category that is invariant 3D GNNs [34].
In this work, we follow invariant 3D GNNs to inherit the merit of SE(3)-invariance by using relative 3D information, such as bond distances, angles between bonds, and angles between planes. Such relative information is used in both representation learning tasks [41, 29, 28, 43, 34, 30]and coordinates generation tasks [18, 44, 45, 53, 26, 6].
Since the relative information is SE(3)-invariant of input 3D graphs, the employed networks favorably achieve the invariance merit.
Wefocus on the 3D graph learning problem and existing methods either capture partial 3D information or suffer from high computational cost.
For example, SchNet [41] only considers distanceinformation and DimeNet [29] further incorporates angles between bonds.
They both integrateincomplete 3D information that the network capacity is limited in practice.
SphereNet [34] generatesapproximate complete 3D representations by using distance, angle, and torsion information but thecomplexity is O ( nk 2 ) .
GemNet [30] is based on quadruplets of nodes, which is more expensive.
Our objective is to build a fully complete 3D graph net with a much lower computational budget.
[Table] There also exist some methods in literature using both absolute and relative3D information [23, 25, 55].
In thiswork, we only use relative informationas input to avoid inherent limitationsof equivariant 3D GNNs.
Moreover,we design a novel message passingsuch that the computational cost iscomparable with that of equivariant3D GNNs.
Experiments
We examine the power and efficiency of ComENet on two large-sacle datasets including Open Catalyst 2020 (OC20) [12] and Molecule3D [53], and the mostly commonly used datastet QM9 [40].
The statistics of three datasets are provided in Table 1.
Detailed descriptions of the datasets areprovided in Appendix A.3.
In particular, the Molecule3D contains about 4 million 3D molecular graphs, and the OC20 has the average graph size of 77.75.
Baseline methods include GINVirtual [24], CGCNN [51], SchNet [41], PhysNet [48], MGCN [35], DimeNet [29], DimeNet++ [28],
SphereNet [34], PaiNN [42], GemNet [30].
Unless otherwise specified, the values for baselinemethods are taken from the referred papers.
For the ComENet, we use data loader in the PyTorch
Geometric library [15] to load the datasets.
All the models are trained using the Adam optimizer [27]and the optimal hyperparameters are obtained on validation sets using grid search.
Experimentalsetup and search space for all models are provided in Appendix A.4.
Code will be released after theanonymous review period.
[Table]
6.1 OC20
The Open Catalyst 2020 (OC20) dataset is a newly released large-scale dataset with millions of
DFT relaxations to model and discover catalysts.
In this work, we focus on Initial Structure to
Relaxed Energy (IS2RE) task, which is the most common task in catalyst discovery.
Descriptionsof the data and tasks are provided in Appendix A.3.
The ground truth of the test set is not publiclyavailable, therefore, we compare the results of different methods on the validation set.
The evaluationmetrics include the energy MAE and the percentage of Energies within a Threshold (EwT) of theground truth energy.
The values for the baseline methods are taken from Chanussot et al.
[11], Liuet al.
Notably, we aim to predict relaxed energy directly from initial structure and do notcompare with some methods using relaxation [12], trajectory information, or relaxed structures.
Using relaxation [12, 30, 43] is computationally expensive during prediction while the relaxationtrajectory and relaxed structures [23, 55] are hard to obtain in practice.
Table 2 shows that ComENet outperforms all the baseline methods in terms of energy MAE, which isalso used as the main evaluation metric in the Open Catalyst Challenge [13].
ComENet achieves bestperformance on two splits and the second best on the other two splits in terms of EwT. Specifically,
ComENet reduces the average energy MAE by 0.0135, which is 2.2% of the second best model.
Notethat ComENet achieves the best results on the OOD Both split in terms of both energy MAE and
EwT. In practice, it is common that test data is in the different domain with the training data.
Hence,
OOD Both can test the generalization capability of learning models.
More importantly, ComENet ismuch more efficient than methods like DimeNet++ and SphereNet.
For example, SphereNet needs 5hours per epoch while ComENet only requires 20 minutes using the same computing infrastructure(NVIDIA RTX A6000 48GB).
Overall, compared with existing best methods, ComENet achievesbetter performance and largely reduces training by at least 10 times.
6.2 Molecule3D
[Table]
The Molecule3D dataset [53] is anewly proposed large-scale dataset, including around 4 million moleculeswith precise ground-state 3D information derived from DFT and molecular properties.
We focus on the prediction of the HOMO-LUMO gapas it is one of the most practicallyrelevant quantum chemical propertiesof molecules.
A detailed description of the Molecule3D dataset is provided in Appendix A.3.
As this is anewly proposed dataset, we run baseline methods including GIN-Virtual [24], SchNet [41], DimeNet++[28] and SphereNet [34], among [Table] which GIN-Virtual is a powerful baseline for 2D graphs while the others are for 3D graphs.
All themodels are trained using the same computing infrastructure (Nvidia GeForce RTX 2080 Ti 11GB).
[Figure]
Table 3 shows that ComENet dramatically reduces training time by 69 times compared with DimeNet++and SphereNet, and costs similar timeas GIN-virtual and SchNet that onlyconsiders distance information.
Interms of performance, our ComENetperforms much better than SchNetfor both random and scaffold splitswith similar time and computing costs,but a little worse than DimeNet++and SphereNet.
This may be dueto the molecules are relatively smallin
Molecule3D compared with OC20,the structures that our complete strategy can distinguish may not exist inthe dataset.
However, considering thecomparable performance and high efficiency, our ComENet is more practically useful than other methods on such large datasets.
Inaddition, Fig. 4 also shows ComENet either converges much faster in terms of total training time orperforms much better compared with other baselines.
6.3 QM9
The QM9 dataset is a widely used dataset for predicting various properties of molecules.
Theevaluation metrics include the MAE for each property and the overall mean standardized MAE (std.
MAE) for all the 12 properties.
Notably, we do not list results for PPGN [36] and Cormorant [2]since they use different training/val/test sizes.
Table 4 shows that ComENet is much better thanthe methods operating in 1-hop neighborhood like SchNet, PhyNet, and MGCN.
Compared with
DimeNet, DimeNet++, PaiNN, and SphereNet, our ComENet achieves similar results on all propertiesand the overall mean std.
MAE.
Consistently, compared with those methods operating within 2-hopneighborhood including DimeNet and SphereNet, ComENet is much more efficient.
6.4 Ablation Study
[Table]
The proposed ComENet considers both local and global representations with two tuples( d, θ, ϕ ) and ( d, τ ) to achieve full completeness.
We investigate contributions of differentrepresentations to demonstrate the advances ofour complete ComENet.
We remove local andglobal convolutional layers from ComENet anddenote as "w/o ( d, θ, ϕ ) " and "w/o ( d, τ ) ", respectively.
The other parts of the network areremaining the same.
Table 5 shows that
the
ComENet performs the best and the other two models achieve similar results. These results demonstrate the effectiveness of both local and global representations in ComENet.
Conclusion
3D information is crucial for 3D molecular graph learning.
Existing methods either learn partial3D information or induce high time complexity.
We propose ComENet that is both complete inincorporating 3D information and efficient with time complexity of O ( nk ) .
Particularly, we proposethe novel rotation angles to fulfill global completeness.
ComENet can generalize to large-scaledatasets, accelerating training and inference by 6-10 times with superior or comparable performance.
A.1 Algorithm for the Proposed Message Passing
1: for i = 1 , ..., n do 2: for j = 1 , ..., k do 3:
Compute d ij = ||
p
i − p
j ||4: end for 5:
Get reference nodes via f i = argmin
k ∈N i ( d ik ) , s i = argmin
k ∈N i \{ f i } ( d ik ) 6: end for 7: for i = 1 , ..., n do 8: for j = 1 , ..., k do 9:
Get reference nodes via f i \ j = (cid:26) f i , if f i ̸
=
j s i , otherwise , f j \ i = (cid:26) f j , if f j ̸
= i s j , otherwise 10:
Compute angles via θ ij = angle( f i , i, j ) , ϕ ij =
angle(plane f
i ,i,s i , plane f
i ,i,j ) , τ ij =
angle(plane f
i \ j ,i,j , plane i,j,f j \ i ) , 11: end for 12:
Update node features via Eq.13: end for
As rigorously shown in Algorithm 1 , there are two nested loops in the message passing.
For eachnested loop, the complexity for the outer loop is O ( n ) , and the complexity for the inner loop is
O ( k ) .
Particularly, within the inner loop, the operations are picking proper reference nodes, thusthe complexity is simply O (1) .
Overall, the total complexity of our message passing is O ( nk ) .
Importantly, the efficiency of our method is also demonstrated in Sec. 6 that ComeNet is 6-10 timesfaster than SphereNet.
Actually, the training time of ComENet is similar to SchNet whose complexityis also O ( nk ) .
A.2 Proofs of Lemma
[Equation]
Nodes
Case (1): i has more than one neighboring nodes c 1 , 2 ,...
𝑗 𝑖 𝑐 ! 𝑑
"# 𝜃 "#
N Nodes [Figure] [Figure] Proof.
We employ contradiction to prove it.
Basically, there exist two cases when adding this newnode j , as illustrated in Fig.
5 and described as following.
Case (1): i has more than one neighboringnodes c 1 , 2 ,... ; Case (2): i only has one neighboring node c 1 .
Generally, in case (1), we need to prove p ij is uniquely defined as ( d ij , θ ij , ϕ ij ) ; in case (2), we needto prove p ij is uniquely defined as ( d ij , θ ij , τ c 1 i ) .
We assume there exists another location for j such that there is a different relative location vector (cid:101) p ij , and this assumption leads to contradiction.
Theproofs of both cases are provided below.
Proof of case (1) :
In case (1), we need to prove p ij is uniquely defined as ( d ij , θ ij , ϕ ij ) .
Based onthe notations defined in Sec. 2.4, we have
[Equation]
where T denotes an operator of projection to the plane perpendicular to p if i .
Apparently, all thequantities on the right side are known based on Sec.
2.4 and Fig.
Assume the solution setcontains at least two different solutions p ij and (cid:101) p ij .
Eq.
2 can simply imply that
[Equation]
Apparently, both T ( p is i ) and p ij − (cid:101) p ij are in the plane perpendicular to p if i .
Hence, T ( p is i ) ×
( p ij − (cid:101) p ij ) =
α p if i holds for some α ̸ = 0 .
However, from Eq. 3, we have
[Equation]
Since α ̸ = 0 and p if i ̸ = 0 , Eq.
4 causes a contradiction.
Thus, the solution set contains a uniquesolution.
Proof of case (2) :
In case (2), we need to prove p ij is uniquely defined by ( d ij , θ ij , τ c 1 i ) .
Based onthe notations defined in Sec.
2.4 and Sec. 2.2, we have
[Equation]
where T still denotes an operator of projection to the plane perpendicular to p if i .
Similarly, allthe quantities on the right side are known based on Sec.
2.4 and Sec. 2.2.
Assume the solution setcontains at least two different solutions p ij and (cid:101) p ij .
Eq.
5 can simply imply that
[Equation]
Both T ( p f i f fi \ i ) and p ij − (cid:101) p ij are in the plane perpendicular to p if i .
Hence, T ( p f i f fi \ i ) ×
( p ij −(cid:101) p ij ) =
α p if i holds for some α ̸ = 0 .
However, from Eq. 6, we have
[Equation]
Since α ̸ = 0 and p if i ̸ = 0 , Eq.
7 causes a contradiction.
Thus, the solution set contains a uniquesolution.
A.3 Data Description
OC20.
The Open Catalyst 2020 (OC20) dataset [12] is a newly released dataset to model anddiscover catalysts.
Specifically, the goal is efficient DFT approximation of structure relaxation,which is a fundamental calculation in catalysis to determine a structure’s activity and selectivity.
All the structures in the dataset contain a surface and an adsorbate, and the surface is defined by aunit cell that is periodic in all directions.
There are three tasks including Structure to Energy and
Forces (S2EF), Initial Structure to Relaxed Structure (IS2RS), and Initial Structure to Relaxed Energy(IS2RE).
In this work, we focus on Initial Structure to Relaxed Energy (IS2RE) task, which is the most commontask in catalysis as the relaxed energies are often correlated with catalyst activity and selectivity.
Thedataset for IS2RE is originally split into training, validation, and test sets.
The training set contains460,328 structures and the validation set has four splits including in-domain (ID), out-of-domainadsorbate (OOD Ads), out-of-domain catalyst (OOD Cat), and out-of-domain adsorbate and catalyst(OOD Both), with 24,733, 24,961, 24,738, 24,971 structures respectively.
Molecule3D.
The Molecule3D dataset [53] is a newly proposed large-scale dataset, including around4 million molecules with precise ground-state geometries derived from DFT.
The dataset is collectedfrom PubChemQC [39] and designed for predicting 3D geometries from molecular graphs [53] whilewe aim to learn representations and predict properties for molecules based on their geometries.
Thedataset contains 3,899,647 molecules and is split into training, validation, and test sets via randomand scaffold split with ratio 6:2:2.
Random split ensures the training, validation, and test data aresampled from the same distribution while scaffold split leads to a distribution shift between trainingand test data.
The evaluation metric is the MAE between the predictions and the ground truth.
A.4 Experimental Setup
ComENet.
The values/search space of model and training hyperparameters for ComENet on OC20,
Molecule3D, and QM9 are provided in Table 6 and Table 7.
For Molecule3D and QM9, the optimalhyperparameters are chosen by the performance on the validation sets.
For OC20, since the finalcomparison results are on validation set, we firstly use 10% of the training data to choose optimalhyperparameters, then train our model on whole training data.
Specifically, we use a larger cutoffvalue for OC20 to generate graphs and larger hidden dimensions for OC20 and Molecule3D. Allmodels are trained on NVIDIA GeForce RTX 2080 Ti 11GB GPU for Molecule3D and QM9.
For the
OC20 dataset, we use NVIDIA RTX A6000 48GB GPU.
Note that one experiment is only conductedon one GPU.
Baselines for Molecule3D. As Molecule3D is a newly proposed dataset, we run experiments andprovide results for baseline methods including GIN-Virtual [24], SchNet [41], DimeNet++[28] and
SphereNet [34].
All the models are trained on one GPU (Nvidia GeForce RTX 2080 Ti 11GB).
Themodel hyperparameters and training hyperparameters are listed in Table 8 and Table 9.
The optimalhyperparameters are chosen by the performance on validation set.
Note that for DimeNet++[28] and
SphereNet [34], the maximum batch size is 32 due to the GPU memory limitation.
[Table]
[Table]
[Table]
[Table]