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, 53], proteins [9, 16, 34, 38, 29],materials [54], and quantum systems [32] 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, 20, 19, 55, 59, 33] and the messagepassing scheme [23, 7, 52] is shown effective to realize 2D GNNs.
However, in the modern machinelearning era, it is increasingly accepted that modeling real-world data like molecules as 2D graphsleads to inborn defects for succeeding learning models.
In practice, 3D information is crucial, basedon which some important geometries can be derived, such as chemical bond lengths in molecularmodeling.
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 [44, 21, 31, 46, 35, 22] 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 [35] 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.
Let n and k denote the number of nodes and theaverage degree in a 3D graph.
Existing methods also exhibit excessive time complexity of O ( nk 2 )
oreven O ( nk 3 ) , severely preventing their scalability in real-world applications.
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 [35] 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 × h , where m is the number of transformed geometric features and h is the feature size.
T can be different dependent on different methods, resulting in different m ∈ N + or h ∈ N + .
Forexample, SchNet [44] only computes the distance for each edge based on the coordinates of the twonodes connected by this edge.
Thus, SchNet maps G to a representation with size m × h , where m isthe number of edges in G and h = 1 .
However, such geometric transformation is not complete.
Weprovide 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
[Equation]
90 Here SE(3) denotes the Special Euclidean group in 3 dimensions.
It include all rotations andtranslations in 3D [1, 25, 35, 42].
Thus R is a transformation that combines rotation and translation.
Arotation transformation can be represented with a 3 × 3 rotation matrix, and a translation transformationcan be represented by a 3 × 1 vector.
Matrix form of SE(3) is provided in Appendix A.2.
The reasonwhy we introduce SE(3) lies in that, a combination of rotation and translation will not change the3D conformation of a 3D graph.
In Def. 1, if P 1 and P 2 are in the same SE(3) group, then G 1 and
G 2 would share the same 3D conformation.
As a result, G 1 and G 2 would be the same 3D graph.
Intuitively, a complete geometric transformation T can distinguish any two different 3D graphs.
This is to say, as long as two 3D graphs differ in 3D conformations, their outputs from T would bedifferent.
2.2 Global Completeness via Rotation Angles
[Figure]
Existing studies focus on the complete representation learning of local neighborhood.
Earlier methodslike SchNet [44] and DimeNet [21] cannot achieve local completeness.
In a more recent method
SphereNet, completeness is guaranteed within edge-based 1-hop local neighborhood, but fails to holdin the whole 3D graph.
In this section, we move a step forward to formally fulfill global completenessfor a given 3D molecular graph.
Particularly, for the purpose of clear illustration, we safely assumelocal completeness is already 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 illustrated
[Figure]
in 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,respectively.
We also define anynode i and its 1-hop neighborhood as a local structure.
Thenthe whole 2-hop neighborhood ofnode 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.
As shown in Fig. 1(a), L i is the local structure centeredin i , and L ij is the local structure centered in j .
Apparently, each local structure L ij,j ∈N i sharesthe 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, theglobal completeness the input 3D molecular graph is fulfilled by additionally considering the rotationangle of each edge, as introduced in this section.
Note that in terms of the selection of referencenodes for computing a rotation angle, we think a selection strategy is valid as long as it is applied toevery edge consistently.
In this work, for an edge e ij , we choose the nearest neighboring nodes for iand j as reference nodes, which are f i \ j and f j \ i , respectively.
We apply this selection strategy to alledges in a 3D graph for computing corresponding rotation angles.
With such selection strategy, wecan prove that our method is complete as shown in Sec. 3.1.
We also show in Appendix A.3 that it iseasier to prove completeness using nearest neighboring nodes as reference nodes.
𝐂 𝐂𝐇𝐇 𝐇 0.
180° 120° 60° 0°
Energy/eV
A
C
B
D 𝐂𝐇𝐂 𝐇 𝐇 𝐂 𝐂𝐇𝐇 𝐇 𝐇 𝟑 𝐂 𝐂 𝐇 𝐇 𝐂 𝐂𝐇𝐇 𝐇 𝐇 𝟑 𝐂 𝐂 𝐇 𝐇 𝐂 𝐂𝐇𝐇 𝐇 𝐇 𝟑 𝐂 𝐂 𝐇
A 𝐂 𝐂𝐇𝐇 𝐇 𝐂 𝐂𝐇𝐇 𝐇
C
B
D 𝐇 𝐂 𝐂𝐇𝐇 𝐇 𝐂 𝐇 𝐇 𝐇 𝟑 𝐂 𝐂 𝐂𝐇𝐇 𝐇 𝐂 𝐇 𝐇 𝐇 𝟑 𝐂 𝐂 𝐂𝐇𝐇 𝐇 𝐂 𝐇 𝐇 𝐂𝐇
[Figure]
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 achievecompleteness 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.
To this end, we formally build our message passing scheme as
[Equation] [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 ) .
Overall,our formal analysis in Sec.
2.2, Sec. 2.3, and Sec.
2.4 lead to the proposed message passing schemedefined in Eq.
It is the first fully complete scheme with great efficiency of O ( nk ) .
We also providerigorous proof on completeness and analysis on efficiency of our message passing in Sec. 3.
Notethat to achieve efficiency, our message passing scheme adopts a novel strategy that computes allthe needed geometries within 1-hop neighborhood.
Hence, it can not be directly applied to existingarchitectures built in 2-hop neighborhood, such as DimeNet++ [31] and SphereNet [35].
To this end,we design a new network to implement the proposed message passing scheme, as detailed in Sec. 4.
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 .
Notethat the 3D graph we consider is strongly connected, which means that there exist a path between anytwo nodes in the graph.
All the molecules in nature can be constructed as strongly connected graphs.
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.3.
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.
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
DMP fashion that update edges within 2-hop neighborhood, inducing the complexity of O ( nk 2 ) .
Notably, DMP [49, 57, 21, 35] incorporates 2-hop information in one single layer but ( n + 1) -hopwhen stacking n layers.
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 ( nk ) torsion angles, which is efficient especially compared with existing methods [22, 1, 18] thatcompute O ( nk 3 ) torsion angles.
Detailed analysis is provided in Appendix A.4.
[Figure]
Based message passingscheme introduced in Sec. 2.5, propose the complete and efficientgraph neural network (ComENet) asshown in Fig. 3.
Existing invariant 3D
GNN methods [43, 21, 31, 35, 22, 45]share the similar architecture pipeline,which contains an input block, several interaction blocks, and an output block.
Our ComENet also follows such architecture fashion alongwith several novel components, suchas self-atom layers and specificallydesigned local and global graph convolution layers, to better fulfill ourproposed message passing scheme in
Eq. 1.
Generally, ComENet consistsof an embedding layer, multiple interaction layers, a self-atom layer, and apooling layer.
To be in line with themessage passing scheme in Sec.
2.5,we take the updating process for nodei ’s feature vector v i as an example todescribe the network.
In practice, feature vectors for all the nodes in a graph are updated simultaneously.
In particular, we omit all indices in Fig.
3 and below for clear presentation.
Specifically,embedding Layer converts atom type z to an initial node feature vector v via learnable atom typeembeddings [44, 21].
Interaction Layer updates node feature vector v based on features of the neighboring nodes and geometric features ( d, θ, ϕ, τ ) in Eq.
1 using the local and global graph convolutionlayers.
A detailed description of interaction layer is provided in Appendix A.5.
Self-atom layer isused to update each node feature and project the feature dimension into 1.
And 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 molecular graphs [3, 10].
One category ofmethods are equivariant 3D GNNs that directly use coordinates in the CSC as inputs to networks [50,2 , 17, 45, 8].
These methods are efficient but suffer from several setbacks.
Firstly, each networkcomponent needs to be carefully designed to be rotation equivariant of input graphs.
Secondly,the reason why this category of methods are efficient lies in that they only use the type-1 basis in
Spherical harmonics, which is an approximation essentially.
It is proved in [50] that theoretically,l is infinite in terms of type- l basis.
In practice, l should be at least 2 for achieving satisfactoryperformance.
Type-1 basis essentially coerces the conv kernel to be in a narrow learning space, which imposes a hard constraint on the network capability.
However, when using type-2 basis, the convkernel can be more expressive while the efficiency issue would emerge as a new bottleneck.
Thirdly,the performance of such equivariant GNNs is shown to be worse than invariant 3D GNNs [35].
In this work, we follow invariant 3D GNNs to inherit the merit of SE(3)-invariance by investigatingrelative 3D information, which is used in both representation learning tasks [44, 21, 31, 46, 35, 22]and coordinates generation tasks [18, 47, 48, 56, 29, 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 orsuffer from high computational cost.
For example, SchNet [44] only considers distance informationand DimeNet [21] further incorporates angles between bonds.
They both integrate incomplete 3Dinformation that the network capacity is limited in practice.
SphereNet [35] generates approximatecomplete 3D representations by using distance, angle, and torsion information but the complexity is
O ( nk 2 ) .
GemNet [22] is based on quadruplets of nodes, which is more expensive.
Our objective isto build a fully complete 3D graph net with a much lower computational budget.
There also exist some methods in literature using both absolute and relative 3D information [24,27, 58].
In this work, we only use relative information as input to avoid inherent limitations ofequivariant 3D GNNs.
Moreover, we design a novel message passing such that the computationalcost is comparable with that of equivariant 3D GNNs.
Experiments
[Table]
We examine the power and efficiency of
ComENet on two large-sacle datasets including Open Catalyst 2020 (OC20) [12]and Molecule3D [56], and the mostlycommonly used datastet QM9 [41].
Thestatistics of three datasets are providedin Table 1.
Detailed descriptions of thedatasets are provided in Appendix A.6.
In particular, the Molecule3D contains about 4 million 3Dmolecular graphs, and the OC20 has the average graph size of 77.75.
Baseline methods include GINVirtual [26], CGCNN [54], SchNet [44], PhysNet [51], MGCN [36], DimeNet [21], DimeNet++ [31],
SphereNet [35], PaiNN [45], GemNet [22].
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 [30]and the optimal hyperparameters are obtained on validation sets using grid search.
Experimentalsetup and search space for all models are provided in Appendix A.7.
Code will be released after theanonymous review period.
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.6.
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 not
[Table]
compare with some methods using relaxation [12], trajectory information, or relaxed structures.
Using relaxation [12, 22, 46] is computationally expensive during prediction while the relaxationtrajectory and relaxed structures [24, 58] 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 [56] 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.6.
As this is anewly proposed dataset, we run baseline methods including GIN-Virtual [26], SchNet [44], DimeNet++[31] and SphereNet [35], amongwhich 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 6-9 times compared with
DimeNet++ and SphereNet, and costs similartime as GIN-virtual and SchNet that only considers distance information.
In terms of performance, our ComENet performs much better than
SchNet for both random and scaffold splits withsimilar time and computing costs, but a littleworse than DimeNet++ and SphereNet.
Thismay be due to the molecules are relatively smallin
Molecule3D compared with OC20, the structures that our complete strategy can distinguishmay not exist in the dataset.
However, considering the comparable performance and high efficiency, our ComENet is more practically usefulthan other methods on such large datasets.
In addition, Fig. 4 also shows ComENet either convergesmuch faster in terms of total training time or performs 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.
A detailed description of the dataset is provided in Appendix A.6.
Comparison results are provided in Table 9 of Appendix A.8.
Notably, we do not list results for PPGN [37] and Cormorant [2] since they use different train/val/test sizes. We use the same split as
DimeNet and SphereNet and provide extensive explanations in Appendix A.8.
Table 9 shows that ComENet is much better than the methods operating in 1-hop neighborhood like SchNet, PhyNet, and
MGCN.
Compared with DimeNet, DimeNet++, PaiNN, and SphereNet, ComENet achieves similarresults on all properties and the overall std.
MAE.
Consistently, compared with methods operating in2-hop neighborhood like DimeNet, DimeNet++, and SphereNet, ComENet is much more efficient.
6.4 Ablation Study for Identifying Conformers
[Figure]
For y-axis, frequency counts the number of conformers for each individual adsorbate and catalyst.
As mentioned in Sec. 2.3, rotationangles are the main difference between different conformers [18, 28].
We investigate the contribution ofour proposed rotation angles τ todemonstrate the effectiveness of ourglobal complete representations.
Weconduct experiments on the OC20dataset since there exist different conformers for molecules in this dataset.
Specifically, there are 660,010 datasamples in the OC20 dataset (IS2RE), where each sample is a combination of two parts, namely,adsorbate and catalyst.
There are 82 adsorbates and 11,500 catalysts used in the datasets.
Eachadsorbate inevitably corresponds to different conformers in the dataset, and it is similar to catalyst.
We show the number of conformers for each adsorbate and each catalyst in Fig.
We remove therotation angle τ from ComENet and denote it as "ComENet w/o τ ".
The results in Table 4 show thatremoving rotation angles τ can harm the performance of ComENet, demonstrating the effectivenessof our global complete representations for identifying conformers.
Additional ablation studies on
QM9 are in Appendix A.8.
[Table]
Conclusions, Limitations, Outlook, and Societal Impacts
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.
Even though ComENet is the first complete and the most efficient 3D GNN model, there exists onemajor limitation, which is not only for ComENet but also for existing 3D GNNs.
Basically, existing3D GNN models are centered on the fact that 3D information is given in data.
However, acquiring3D information itself is difficult and expensive in practice.
Current methods rely on experiments or
DFT-based computing, which is extremely time-consuming.
It would be significant that machinelearning models can be developed to tackle this problem.
Looking forward, we can derive twodirections for fulfilling such objective.
Firstly, we can study to generate 3D graphs either from 2Dgraphs or from scratch by developing generative models, such as VAE, flow and diffusion models.
Especially, in some real-world applications like drug discovery, 2D molecules are usually not given.
This raises the need of developing new generation methods from scratch.
Secondly, we can targetat a research case where we have a minimal set of training data with 3D information, but the vastunseen data or new data lack such 3D information.
We may develop novel contrastive learningcomponents to force correspondence and consistency between 2D graphs and their 3D geometricviews, then integrate such components into end-to-end learning systems for application based on 2Dgraph data.
ComENet can facilitate a plethora of important real-world applications, such as drugdiscovery and material discovery.
It can be used in several research domains including quantumchemistry and physics, material sciences, molecular dynamics simulations, etc.
Any negative societalimpact associated with those applications and domains can be applied to our method.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes](b) Did you describe the limitations of your work?
[Yes](c)
Did you discuss any potential negative societal impacts of your work?
[Yes](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results
[Table]
Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix A.3.3.
If you ran experiments...
[Table]
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] We will includethese details upon acceptance.
[Table]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Appendix A.7.
[Table]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[N/A] Error bars are not necessary for baseline methods andour method.
[Table]
(d) Did you include the total amount of compute and the type of resources used (e.g.,type of GPUs, internal cluster, or cloud provider)?
[Yes] All of these are reported inAppendix A.7.
[Table]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
[Table]
(a) If your work uses existing assets, did you cite the creators?
[Yes]
[Table]
(b) Did you mention the license of the assets?
[N/A]
[Table]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
[Table]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
[Table]
using/curating?
[N/A]
[Table]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
[Table]
information or offensive content?
[N/A]
[Table]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
[Table]
Board (IRB) approvals, if
applicable?
[N/A](c)
Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]
A
Appendix
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 SE(3)
SE(3) is the Special Euclidean group in 3 dimensions, including all rotations and translations in 3D.
It is the set of 4 × 4 real matrices of the form
[Equation]

[Equation] [Equation]
where Rot ∈ SO(3) and t ∈ R 3 .
The SO(3) is the set of 3 × 3 real matrices Rot satisfying
Rot T Rot =
I and det( Rot )
= 1 .
A.3 Proofs of Lemma 1
Proof.
We employ contradiction to prove it.
Basically, there exist two cases when adding this newnode j , as illustrated in Fig.
6 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 suchthat there is a different relative location vector (cid:101) p ij , and this assumption leads to contradiction.
Theproofs of both cases are provided below.
𝑗𝑗 𝑖𝑖 𝑐𝑐𝑐𝑐𝜙𝜙 𝑖𝑖𝑖𝑖 𝑑𝑑 𝑖𝑖𝑖𝑖 𝜃𝜃 𝑖𝑖𝑖𝑖
N Nodes
Case (1): i has more than one neighboring nodes c 1 , 2 ,...
𝑗 𝑖 𝑐 ! 𝑑
"# 𝜃 "#
N Nodes 𝜏 $ ! "
Case (2): i only has one neighboring node c
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]
(3) 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.
3 can simply imply that
[Equation] [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. 4, we have
[Equation]
Since α ̸ = 0 and p if i ̸ = 0 , Eq.
5 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 set contains at least two different solutions p ij and (cid:101) p ij .
Eq.
6 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 −
[Equation]
p
[Equation]
=
α p if i holds for some α ̸ = 0 .
However, from Eq. 7, we have
[Equation]
Since α ̸ = 0 and p if i ̸ = 0 , Eq.
8 causes a contradiction.
Thus, the solution set contains a uniquesolution.
A.4 Efficiency via Less Torsion Angles
Our method achieves completeness by computing O ( nk ) torsion angles, which is efficient especiallycompared with methods including Gasteiger et al.
[22], Adams et al.
[1], Ganea et al.
Given ascenario where a global region is the union of two local regions with n 1 and n 2 nodes.
As introducedin Sec. 2.2 and Sec. 2.4, our method computes the same number of torsion angles as nodes in a localregion, then employs a rotation angle (torsion angle essentially).
Hence, the number of torsion anglesthat our method computes is n 1 + n 2 + 1 .
Ganea et al.
[18] computes a torsion angle based on onepair of nodes, each of which is from a separate local region.
Hence, the number of torsion anglesneeded is n 1 × n 2 .
Apparently, our method reduces the number of torsion angles from O ( nk 3 ) to
O ( nk ) , which is significant considering the computing of torsion is excessively expensive.
A.5 Model architecture
Interaction Layer updates each node feature vector v based on features of the neighboring nodes andthe corresponding 3D information in P .
Firstly, it converts 3D information in P to a set of geometriesbased on the proposed complete geometric transformation T and message passing scheme.
Sincedistance is the most important geometry, we also consider d in global representation and split theoutput 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 [27, 21, 35, 22], 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 [35] and GemNet [22].
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 [39] 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 ′ .
A.6 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 [56] is a newly proposed large-scale dataset, including around4 million molecules with precise ground-state geometries derived from DFT.
The dataset is collectedfrom PubChemQC [40] and designed for predicting 3D geometries from molecular graphs [56] 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.
QM9.
The QM9 dataset [41] is a widely used dataset for predicting various properties of molecules.
It includes geometric, energetic, electronic, and thermodynamic properties for 134k stable smallorganic molecules.
The dataset is split into three sets, where the training set contains 110,000, thevalidation set contains 10,000, and the test set contains 10,831 molecules.
The twelve properties aredipole moment ( µ ), isotropic polarizability ( α ), highest occupied molecular orbital energy ( ϵ HOMO ),lowest unoccupied molecular orbital energy ( ϵ LUMO ), gap between ϵ HOMO and ϵ LUMO , electronicspatial extent ( (cid:10) R 2 (cid:11) ), zero point vibrational energy (ZPVE), internal energy at 0K ( U 0 ), internalenergy at 298.15K ( U ) , enthalpy at 298.15K ( H ), free energy at 298.15K ( G ) , and heat capavity at298.15K ( c v ).
A.7 Experimental Setup
ComENet.
The values/search space of model and training hyperparameters for ComENet on OC20,
Molecule3D, and QM9 are provided in Table 5 and Table 6.
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 [26], SchNet [44], DimeNet++[31] and
SphereNet [35].
All the models are trained on one GPU (Nvidia GeForce RTX 2080 Ti 11GB).
Themodel hyperparameters and training hyperparameters are listed in Table 7 and Table 8.
The optimalhyperparameters are chosen by the performance on validation set. Note that for DimeNet++[31] and SphereNet [35], the maximum batch size is 32 due to the GPU memory limitation. [Table]
[Table]
A
[Table]
8 Additional results
[Table]
Results on QM9.
Experiment results on QM9 are listed in Table 9.
Baseline methods include
SchNet [44], PhysNet [51], MGCN [36], PaiNN [45], DimeNet [21], DimeNet++ [31] [Table] and SphereNet [35].
All baseline methods use 110,000 molecules for training, 10,000 for validation, and10,831 for testing.
Specifically, we use exactly the same data split as DimeNet, DimeNet++, and
SphereNet with the same random seed of 42.
PaiNN uses the same sizes of training, validation, andtest sets.
But its result is the average of three runs with three random splits.
Therefore, we follow thissetting and run our ComENet on α and ϵ HOMO .
The results are 0 .
045 ± 0 . 000 and 23 .
2 ± 0 .
2 , whichare similar to the values in Table 9.
For SchNet, PhysNet, and MGCN, they use different randomseeds.
Considering that their results are much worse than DimeNet, DimeNet++, SphereNet, andour ComENet, we do not reimplement these methods.
The reported values for these three methodsare taken from the referred papers.
Specifically, the results of PhysNet and MGCN are taken from
DimeNet [21] and DimeNet++ [31].
The results of SchNet are taken from the original paper [44].
Ablation study on QM9.
The proposed ComENet considers both local and global representationswith two tuples ( d, θ, ϕ ) and ( d, τ ) to achieve full completeness.
We investigate contributionsof different representations to demonstrate the advances of our complete ComENet.
We removelocal and global convolutional layers from ComENet and denote as "w/o ( d, θ, ϕ ) " and "w/o ( d, τ ) ",respectively.
The other parts of the network are remaining the same.
Table 10 shows that the ComENetperforms the best and the other two models achieve similar results.
These results demonstrate theeffectiveness of both local and global representations in ComENet.
[Table] [Table]
[Table] [Table]