Title: A new perspective on building efficient and expressive 3D equivariant graph neural networks
Venue: NeurIPS 2023
Accepted: yes
Keywords: Graph Neural Networks

Geometric Deep Learning

Equivariance

Symmetry
Abstract: Geometric deep learning enables the encoding of physical symmetries in modeling 3D objects. Despite rapid progress in encoding 3D symmetries into Graph Neural Networks (GNNs), a comprehensive evaluation of the expressiveness of these network architectures through a local-to-global analysis lacks today. In this paper, we propose a local hierarchy of 3D isomorphism to evaluate the expressive power of equivariant GNNs and investigate the process of representing global geometric information from local patches. Our work leads to two crucial modules for designing expressive and efficient geometric GNNs; namely local substructure encoding (\textbf{LSE}) and frame transition encoding (\textbf{FTE}). To demonstrate the applicability of our theory, we propose LEFTNet which effectively implements these modules and achieves state-of-the-art performance on both scalar-valued and vector-valued molecular property prediction tasks. We further point out future design space for 3D equivariant graph neural networks. Our codes are available at \url{https://github.com/yuanqidu/LeftNet}.

A new perspective on building efficient and expressive 3D equivariant graph neural networks

Abstract

Geometric deep learning enables the encoding of physical symmetries in modeling 3D objects. Despite rapid progress in encoding 3D symmetries into Graph Neural Networks (GNNs), a comprehensive evaluation of the expressiveness of these network architectures through a local-to-global analysis lacks today. In this paper, we propose a local hierarchy of 3D isomorphism to evaluate the expressive power of equivariant GNNs and investigate the process of representing global geometric information from local patches. Our work leads to two crucial modules for designing expressive and efficient geometric GNNs; namely local substructure encoding ( LSE ) and frame transition encoding ( FTE ). To demonstrate the applicability of our theory, we propose LEFTNet which effectively implements these modules and achieves state-of-the-art performance on both scalar-valued and vector-valued molecular property prediction tasks. We further point out future design space for 3D equivariant graph neural networks. Our codes are available at https://github.com/yuanqidu/LeftNet .

1 Introduction

The success of many deep neural networks can be attributed to their ability to respect physical symmetry, such as Convolutional Neural Networks (CNNs) [1] and Graph Neural Networks (GNNs) [2]. Specifically, CNNs encode translation equivariance, which is essential for tasks such as object detection. Similarly, GNNs encode permutation equivariance, which ensures that the node ordering does not affect the output node representations, by aggregating neighboring messages. Modeling 3D objects, such as point clouds and molecules, is a fundamental problem with numerous applications [3], including robotics [4], molecular simulation [5, 6], and drug discovery [7-11]. Different from 2D pictures and graphs that only possess the translation [1] and permutation [2] symmetry, 3D objects intrinsically encode the complex SE (3) /E (3) symmetry [12], which makes their modeling a nontrivial task in the machine learning community.

To tackle this challenge, several approaches have been proposed to effectively encode 3D rotation and translation equivariance in the deep neural network architectures, such as TFN [13], EGNN [14], and SphereNet [15]. TFN leverages spherical harmonics to represent and update tensors equivariantly,

Zhi-Ming Ma

Chinese Academy of Sciences

while EGNN processes geometric information through vector update. On the other hand, SphereNet is invariant by encoding scalars like distances and angles. Despite rapid progress has been made on the empirical side, it's still unclear what 3D geometric information can equivariant graph neural networks capture and how the geometric information is integrated during the message passing process [16-18]. This type of analysis is crucial in designing expressive and efficient 3D GNNs, as it's usually a trade-off between encoding enough geometric information and preserving relatively low computation complexity. Put aside the SE (3) /E (3) symmetry, this problem is also crucial in analysing ordinary GNNs. For example, 1-hop based message passing graph neural networks [19] are computationally efficient while suffering from expressiveness bottlenecks (comparing with subgraph GNNs [20, 21]). On the other hand, finding a better trade-off for 3D GNNs is more challenging, since we must ensure that the message updating and aggregating process respects the SE (3) /E (3) symmetry.

In this paper, we attempt to discover better trade-offs between computational efficiency and expressive power for 3D GNNs by studying two specific questions: 1. What is the geometric expressiveness of 3D GNNs through a local 3D graph isomorphism lens? 2. What is expressiveness of 3D GNNs in representing global geometric information through local patches? The first question relates to the design of node-wise geometric messages, and the second question relates to the design of equivariant (or invariant) aggregation. To tackle these two problems, we take a local-to-global approach. More precisely, we first define three types of 3D graph isomorphism to characterize local 3D structures: tree, triangular, and subgraph isomorphism, following a local hierarchy. Our local hierarchy lies between the 1-hop and 2-hop geometric isomorphism defined in Joshi et al. [22], detailed in Appendix G.2; thus, it can be used to measure the expressive power of 3D GNNs by their ability of differentiating non-isomorphic 3D structures similar to the geometric WL test proposed in Joshi et al. [22]. Under this theoretical framework, we summarize one essential ingredient for building expressive geometric messages on each node: local 3D substructure encoding ( LSE ), which allows an invariant realization. To answer the second question, we analyze whether local invariant features are sufficient for expressing global geometries by message aggregation, and it turns out that frame transition encoding ( FTE ) is crucial during the local to global process. Although FTE can be realized by invariant scalars, we further demonstrate that introducing equivariant messaging passing is more efficient. After presenting LSE and FTE modules, we are able to present a modular overview of 3D GNNs designs. In realization of our theoretical findings, we propose LEFTNet that efficiently implements LSE and FTE without sacrificing expressiveness. Empirical experiments on real-world scenarios, predicting scalar-valued property (e.g. energy) and vector-valued property (e.g. force) for molecules, demonstrate the effectiveness of LEFTNet.

2 Preliminary

In this section, we provide an overview of the mathematical foundations of E (3) and SE (3) symmetry, which is essential in modeling 3D data. We also summarize the message passing graph neural network framework, which enables the realization of E (3) /SE (3) equivariant models.

Euclidean Symmetry. Our target is to incorporate Euclidean symmetry to ordinary permutationinvariant graph neural networks. The formal way of describing Euclidean symmetry is the group E (3) = O (3) ⋊ T (3) , where O (3) corresponds to reflections (parity transformations) and rotations. For tasks that are anti-symmetric under reflections (e.g. chirality), we consider the subgroup SE (3) = SO (3) ⋊ T (3) , where SO (3) is the group of rotations. We will use SE (3) in the rest of the paper for brevity except when it's necessary to emphasize reflections.

Equivariance. A tensor-valued function φ ( x ) is said to be equivariant with respect to SE (3) if for any translation or rotation g ∈ SE (3) acting on x ∈ R 3 , we have where M ( · ) is a matrix representation of SE (3) acting on tensors. See Appendix A for a general definition of tensor fields. In this paper, we will use bold letters to represent an equivariant tensor, e.g., x as a position vector. It is worth noting that when φ ( x ) ∈ R 1 and M ( g ) ≡ 1 (the constant group representation), the equivariant function φ ( x ) is also called an invariant scalar function.

Message Passing Scheme for Geometric Graphs. A geometric graph G is represented by G = ( V, E ) . Here, v i ∈ V denotes the set of nodes (vertices, atoms), and e ij ∈ E denotes the set of edges. For brevity, the edge feature attached on e ij is also denoted by e ij . X = ( x 1 , . . . , x n ) ∈ R n × 3 denotes the node positions which determine the geometric structure of G .

x

i

S

i

A common machine learning tool for modeling graph-structured data is the Message Passing Neural Network (MPNN) [16]. A typical 1-hop MPNN framework consists of two phases: (1) message passing; (2) readout. Let h l i , h l j denote the l -th layer's node features of source i and target j that also depend on the 3D positions ( x i , x j ) , then the aggregated message is

and ⊕ j ∈N ( i ) is any permutation-invariant pooling operation between the 1-hop neighbors of i . We also include the edge features e l ij into the message passing phase for completeness. 3D equivariant GNNs (3D GNNs for short) require the message m i to be equivariant with respect to the geometric graph. That is, for an arbitrary edge e ij :

where g ∈ SE (3) is acting on the whole geometric graph simultaneously: ( x 1 , . . . , x n ) → ( g x 1 , . . . , g x n ) . For example, the invariant model ComENet [23] satisfies Eq. 2 by setting M ( g ) ≡ 1 , and MACE [24] realized Eq. 2 for nonconstant irreducible group representations M ( g ) through spherical harmonics and Clebsch-Gordan coefficients.

3 A local hierarchy of 3D graph isomorphism

As presented in Section 2, defining expressive messages is an essential component for building powerful 3D GNNs. In this section, we develop a fine-grained characterization of local 3D structures and build its connection with the expressiveness of 3D GNNs.

Since the celebrated work [25], a popular expressiveness test for permutation invariant graph neural networks is the 1-WL graph isomorphism test [26], and Wijesinghe and Wang [27] has shown that the 1-WL test is equivalent to the ability to discriminate the local subtree-isomorphism. It motivates us to develop a novel (local) 3D isomorphism for testing the expressive power of 3D GNNs. However, this task is nontrivial, since most of the previous settings for graph isomorphism are only applicable to 2D topological features. For 3D geometric shapes, we should take the SE (3) symmetry into account. Formally, two 3D geometric graphs X , Y are defined to be globally isomorphic, if there exists g ∈ SE (3) such that

In other words, X and Y are essentially the same, if they can be transformed into each other through a series of rotations and translations. Not that Eq. 3 is up to the permutation of nodes. Inspired by Wijesinghe and Wang [27], now we introduce a novel hierarchy of SE (3) equivariant local isomorphism to measure the local similarity of 3D structures.

Let S i represent the 3D subgraph associated with node i . This subgraph contains all the 1-hop neighbors of i as its node set, along with all edges in E where both end points are one-hop neighbors of i . For each edge e ij ∈ E , the mutual 3D substructure S i -j is defined by the intersection of S i and S j : S i -j = S i ∩ S j .

Given two local subgraphs S i and S j that correspond to two nodes i and j (not necessarily adjacent), we say S i is { tree-, triangular-, subgraph-} isometric to S j , if there exists a bijective function

f : S i → S j such that h f ( u ) = h u for every node u ∈ S i , and the following conditions hold respectively:

Tree Isometric: If there exists a collection of group elements g iu ∈ SE (3) , such that ( x f ( u ) , x f ( i ) ) = ( g iu x u , g iu x i ) for each edge e iu ∈ S i ;

Subgraph Isometric: for any two adjacent nodes u, v ∈ S i , f ( u ) and f ( v ) are also adjacent in S j , and there exist a single group element g i ∈ SE (3) such that g i S i = S j .

Triangular Isometric: If there exists a collection of group elements g iu ∈ SE (3) , such that the corresponding mutual 3D substructures satisfy: S f ( u ) -f ( i ) = g iu S u -i for each edge e iu ∈ S i ;

Note that tree isomorphism only considers edges around a central node, which is of a tree shape. On the other hand, the mutual 3D substructure can be decomposed into a bunch of triangles (since it's contained in adjacent node triplets), which explains the name of triangular isomorphism.

In fact, the three isomorphisms form a hierarchy from micro to macro, in the sense that the following implication relation holds:

Subgraph Isometric ⇒ Triangular Isometric ⇒ Tree Isometric

This is an obvious fact from the above definitions. To deduce the reverse implication relation, we provide a visualized example. Fig. 1 shows two examples of local 3D structures: 1. the first one shares the same tree structure, but is not triangular-isomorphic; 2. the second one is triangular-isomorphic but not subgraph-isomorphic. In conclusion, the following diagram holds:

Tree Isometric ̸⇒ Triangular Isometric ̸⇒ Subgraph Isometric

One way to formally connect the expressiveness power of a geometric GNN with their ability of differentiating geometric subgraphs is to define geometric WL tests, the reader can consult [22]. In this paper, we take an intuitive approach based on our nested 3D hierarchy. That is, if two 3D GNN algorithms A and B can differentiate all non-isomorphic local 3D shapes of tree (triangular) level, while A can differentiate at least two more 3D geometries which are non-isomorphic at triangular(subgraph) level than B, then we claim that algorithm A's expressiveness power is more powerful than B.

Since tree isomorphism is determined by the one-hop Euclidean distance between neighbors, distinguishing local tree structures is relatively simple for ordinary 3D equivariant GNNs. For example, the standard baseline SchNet [28] is one instance of Eq. 1 by setting e t ij = RBF ( d ( x i , x j )) , where RBF ( · ) is a set of radial basis functions. Although it is powerful enough for testing tree nonisomorphism (assuming that RBF ( · ) is injective), we prove in Appendix D that SchNet cannot distinguish non-isomorphic structures at the triangular level.

On the other hand, Wijesinghe and Wang [27] has shown that by leveraging the topological information extracted from local overlapping subgraphs, we can enhance the expressive power of GNNs to go beyond 2D subtree isomorphism. In our setting, the natural analogue of the overlapping subgraphs is exactly the mutual 3D substructures. Now we demonstrate how to merge the information from 3D substructures to the message passing framework (1). Given an SE (3) -invariant encoder ϕ , define the 3D structure weights A ij := ϕ ( S i -j ) for each edge e ij ∈ E . Then, the message passing framework (1) is generalized to:

Formula 4 is an efficient realization of enhancing 3D GNNs by injecting the mutual 3D substructures. However, a crucial question remains to be answered: Can the generalized message passing framework boost the expressive power of 3D GNNs? Under certain conditions, the following theorem provides an affirmative answer:

Theorem 3.1. Suppose ϕ is a universal SE(3)-invariant approximator of functions with respect to the mutual 3d structures S i -j . Then, the collection of weights {{ A ij := ϕ ( S i -j ) } e ij ∈ E } enables the differentiation of local structures beyond tree isomorphism. Furthermore, under additional injectivity assumptions (as described in Eq. 16), 3D GNNs based on the enhanced message passing framework (see Section 4) map at least two distinct local 3D subgraphs with isometric local tree structures to different embeddings.

This theorem confirms that the enhanced 3D GNN (formula 4) is more expressive than the SchNet baseline, at least in testing local non-isomorphic geometric graphs. The existence of such local encoder ϕ is proved in two ways : 1. Equivariant construction by the Atomic Cluster Expansion (ACE) [29]; 2. Invariant construction under scalarization by edge-wise equivariant frames [30]. Note that there are other different perspectives on characterizing 3D structures, we will also briefly discuss them in Appendix D.

4 From local to global: the missing pieces

In the last section, we introduced a local 3D graph isomorphism hierarchy for testing the expressive power of 3D GNNs. Furthermore, we motivated adding a SE(3)invariant encoder to improve the expressive power of one-hop 3D GNNs by scalarizing not only pairwise distances but also their mutual 3D structures in Theorem 3.1. However, to build a powerful 3D GNN, it remains to be analyzed how a 3D GNN acquires higher order (beyond 1-hop neighbors) information by accumulating local messages. A natural question arises: Are local invariant messages enough for representing global geometric information?

To formally formulate this problem, we consider a two-hop aggregation case. From fig. 2, the central atom a is connected with atoms b and c . Except for the common neighbor a , other atoms that connect to b and c form two 3D clusters, denoted by B , C . Suppose the ground-truth interaction potential of B and C imposed on atom a is described by a tensor-valued function f a ( B , C ) . Since B and C are both beyond the 1-hop neighborhood of a , the information of f a ( B , C ) can only be acquired after two steps of message passing: 1. atoms b and c aggregate message separately from B and C ; 2. the central atom a receives the aggregated message (which contains information of B and C ) from its neighbors b and c .

b

c

cluster B

cluster C

a

R

CD

Frame Transition between B and C

h

B

h

C

f

a

=

h

B

·

h

C

x

b

x

c

y

c

y

b

z

b

z

c

Let S B ( S C ) denote the collection of all invariant scalars created by B ( C ) . For example, S B contains all relative distances and angles within the 3D structure B . Then, the following theorem holds:

̸

Theorem 4.1. Not all types of invariant interaction f a ( B , C ) can be expressed solely as functions of the union of two sets S B and S C . In other words, there exists E (3) invariant function f a ( B , C ) , such that it cannot be expressed as functions of S B and S C : f a ( B , C ) = ρ ( S B , S C ) for arbitrary invariant functions ρ .

Proof. The detailed proof is provided in Appendix E.

This theorem essentially demonstrates that simply aggregating "local" scalar information from different clusters is insufficient for approximating "global" interactions, even in the case of simple invariant potential learning tasks. In contrast to the previous section, where the local expressiveness was evaluated based on the ability to classify geometric shapes, in this theorem, we constructed counterexamples using continuous regression functions that depend strictly on information beyond the combination of local invariant scalars. Intuitively, the theorem is proved by introducing two local equivariant frames (three independent equivariant vectors, see Appendix B for the definition) determined by B ( C ) separately, through which all scalars in S B ( S C ) can be expressed. However, the transition matrix between these frames is not encoded in the aggregation, resulting in information loss when aggregating geometric features from two sub-clusters. In other words, local frames provide local observations of geometric quantities, while the transition matrix reveals the global changes between local geometries. Importantly, the proof also highlights that the missing information that causes the expressiveness gap is solely the Frame Transition (FT) information, which we will define immediately.

Frame Transition (FT). Let ( e i 1 , e i 2 , e i 3 ) and ( e j 1 , e j 2 , e j 3 ) be two orthonormal frames in R 3 . These frames are connected by an orthogonal matrix R ij ∈ SO (3) :

Furthermore, when ( e i 1 , e i 2 , e i 3 ) and ( e j 1 , e j 2 , e j 3 ) are equivariant frames, all elements of R ij are invariant scalars. For instance, in the case of a geometric graph where i and j represent indexes of two connected atoms, the fundamental torsion angle τ ij in ComeNet [23] corresponds to one element of R ij (see Appendix E).

To address this expressiveness gap, one approach is to directly incorporate all invariant edge-wise frame transition matrices ( FT ) into the model. This can be achieved using a geometric formulation based on neural sheaf diffusion , as described in Appendix H. However, it is worth noting that this method becomes computationally expensive when dealing with a large number of local clusters, as it requires O ( k 2 ) pairs of FT for each node. Instead, we propose a more efficient approach by introducing equivariant tensor features for each node i , denoted as m i . These tensor features allow us to maintain equivariant frames directly within each m i . We prove in Appendix E that FT can be easily derived through equivariant message passing and updating when utilizing these equivariant tensor features.

Equivariant Message Passing. Similarly with the standard one-hop message passing scheme 1, the aggregated tensor message m i from the l -1 layer to the l layer can be written as: m l -1 i = ∑ j ∈ N ( i ) m l -1 j . Since summation does not break the symmetry rule, it is obvious that m l -1 i are still equivariant tensors. However, the nontrivial part lies in the design of the equivariant update function ϕ :

To fully capture the information of FT , it is necessary for ϕ to possess sufficient expressive power while maintaining SE (3) equivariance. Here, we propose a novel way of updating scalar and tensor messages by performing node-wise scalarization and tensorization blocks (the FTE module of Figure 3). From the perspective of Eq. 2, m ( x u ) is transformed equivariantly as:

Here, m ( x u ) is decomposed to ( m 0 ( x u ) , . . . , m l ( x u )) according to different tensor types, and {M i ( g ) } l i =0 is a collection of different SE (3) tensor representations (see the precise definition in Appendix A).

To illustrate the benefit of aggregating equivariant messages from local patches, we study a simple case. Let f a ( B , C ) = h B · h C be an invariant function of B and C (see Fig. 2), then f a can be calculated by a direction composition of scalar messages and equivariant vector messages: f a ( B , C ) = 1 2 [ ∥ m a ∥ 2 -∥ h B ∥ 2 -∥ h C ∥ 2 ] , where m a = h B + h C is an equivariant vector. Note that m a follows the local equivariant aggregation formula 6, and the other vectors' norm ∥ h B ∥ and ∥ h C ∥ are obtained through local scalarization on atoms b and c . As a comparison, it's worth mentioning that f a ( B , C ) can also be expressed by local scalarization with the additional transition matrix data R BC defined by Eq. 5. Let ˜ h B and ˜ h C be the scalarized coordinates with respect to two local equivariant frames

adding the transition matrix R BC for each ( B , C ) pair into the aggregation procedure, which is computationally expensive compared to directly implementing equivariant tensor updates.

5 Building an efficient and expressive equivariant 3D GNN

We propose to leverage the full power of LSE and FTE modules (along with a necessary tensor update module) to push the limit of efficient and expressive 3D equivariant GNNs design. Specifically, we improve a recently proposed method based on constructing local frames but without the implementation of LSE and FTE [30]. We also analyze related works following this framework detailed in Appendix G.

Output

Subgraph

𝑺

!"#

Frame

ℱ

!#

, ℱ

!

Interaction

⋮

Interaction

Output block

Embedding

ℎ

!

, 𝐡

!

ℎ

!

$

, ℎ

#

$

LSE

to encode

local 3D

structures

Equivariant

message

passing

FTE

to encode

frame

transition

ℎ

!

$%&

Interaction block

𝑗 ∈ 𝒩

!

𝐡

!

$

, 𝐡

#

$

𝐡

!

$%&

𝑺

!"#

ℱ

!#

ℱ

!

𝐴

!#

𝐦

!

$

𝑚

!

$

Invariant

Equivariant

Scalarization

𝐦

!

$

FTE

Tensorization

ℱ

!

ℎ

!

$%&

ℎ

!

$%&

𝐡

!

$%&

MLP

𝑚

!

$

𝑺

!"#

ℱ

!#

𝐴

!#

Scalarization

𝒙

"

∈ 𝑺

!#$

& 𝑥

"

LSE

MLP

RBF

Average

'

∥

⨀

LSE Instantiation. We propose to apply edge-wise equivariant frames (following [30]) to encode the local 3D structures S i -j . By definition, S i -j contains edge e ij , nodes i and j , and their common neighbors. We use the equivariant frame F ij built on e ij (see the precise formula in Appendix F) to scalarize S i -j . After scalarization (8), the equivariant coordinates of all nodes in S i -j are transformed into invariant coordinates: { x k → ˜ x k for x k ∈ S i -j } . To encode these scalars sufficiently, we first weight each ˜ x k by the RBF distance embedding: ˜ x k → RBF ( ∥ x k ∥ ) ⊙ MLP (˜ x k ) for each x k ∈ S i -j . Note that to preserve the permutation symmetry, the MLP is shared among the nodes. Finally, the 3D structure weight A ij is obtained by the average pooling of all node features.

FTE Instantiation. We propose to introduce equivariant tensor message passing and update function for encoding local FT information. At initialization, let NF l ( x i , x j ) denote the embedded tensorvalued edge feature between i and j . We split it into two parts: 1. the scalar part SF l ( x i , x j ) for aggregating invariant messages; 2. the higher order tensor part TF l ( x i , x j ) for aggregating tensor messages. To transform TF l ( x i , x j ) , we turn to the equivariant frame F ij once again. After scalarization by F ij , TF l ( x i , x j ) becomes a tuple of scalars ˜ TF l ( x i , x j ) , which is then transformed by MLP. Finally, we output arbitrary tensor messages through equivariant tensorization 22:

Further details are provided in Appendix F. As we have discussed earlier, the node-wise update function ϕ in Eq. 6 is also one of the guarantees for a powerful FTE . As a comparison, ϕ is usually a standard MLP for updating node features in 2D GNNs, which is a universal approximator of invariant functions. Previous works [14, 31] updated equivariant features by taking linear combinations and calculating the invariant norm of tensors, which may suffer from information loss. Then a natural question arises: Can we design an equivariant universal approximator for tensor update? We answer this question by introducing a novel node-wise frame. Consider node i with its position x i , let ¯ x i := 1 N ∑ x j ∈ N ( x i ) x j be the center of mass around x i 's neighborhood. Then the orthonormal equivariant frame F i := ( e i 1 , e i 2 , e i 3 ) with respect to x i is defined similar to the edge-wise frame between x i and ¯ x i in [30], detailed in appendix B. Finally, we realize a powerful ϕ by the following theorem:

Theorem 5.1. Equipped with an equivariant frame F i for each node i , the equivariant function ϕ defined by the following composition is a universal approximator of tensor transformations: ϕ : Scalarization → MLP → Tensorization .

Proof.

LEFTNet. An overview of our { LSE , FTE } enhanced efficient graph neural network (LEFTNet) is depicted in Fig. 3. The detailed algorithm for LEFTNet is shown in Algorithm 1 of Appendix C. LEFTNet receives as input a collection of node embeddings { v 0 1 , . . . , v 0 N } , which contain the atom types and 3D positions for each node: v 0 i = ( z i , x i ) , where i ∈ { 1 , . . . , N } . For each edge e ij ∈ E ,

we denote the associated equivariant features consisting of tensors by e ij . During each messaging passing layer, the LSE module outputs the scalar weight coefficients A ij as enhanced invariant edge feature and feed into the interaction module. Moreover, scalarization and tensorization as two essential blocks are used in the equivariant update module that fulfills the function of FTE . The permutation equivariance of a geometric graph is automatically guaranteed for any message passing architecture, we provide a complete proof of SE(3)-equivariance for LEFTNet in Appendix F.

SE(3) vs E(3) Equivariance. Besides explicitly fitting the SE (3) invariant molecular geometry probability distribution, modeling the energy surface of a molecule system is also a crucial task for molecule property prediction. However, the Hamiltonian energy function E of a molecule is invariant under refection transformation: Energy ( X ) = Energy ( R X ) , for arbitrary reflection transformation R ∈ E (3) . In summary, there exist two different inductive biases for modeling 3D data: (1) SE(3) equivariance, e.g. chirality could turn a therapeutic drug to a killing toxin; (2) E(3) equivariance, e.g. energy remains the same under reflections.

Since we implement SE (3) equivariant frames in LEFTNet, our algorithm is naturally SE (3) equivariant. However, our method is flexible to implement E (3) equivariant tasks as well. For E (3) equivariance, we can either replace our frames to E (3) equivariant frames, or modify the scalarization block by taking the absolute value: x → ˜ x := ( x · e 1 , x · e 2 , x · e 3 ) ︸ ︷︷ ︸ SE (3) → ( x · e 1 , | x · e 2 | , x · e 3 ) ︸ ︷︷ ︸ E (3) .

Intuitively, since the second vector e 2 is a pseudo-vector, projections of any equivariant vectors along the e 2 direction are not E (3) invariant until taking the absolute value.

Efficiency. To analyze the efficiency of LEFTNet, suppose 3D graph G has n vertices, and its average node degree is k . Our algorithm consists of three phases: 1. Building equivariant frames and performing local scalarization; 2. Equivariant message passing; 3. Updating node-wise tensor features through scalarization and tensorization. Let l be the number of layers, then the computational complexity for each of our three phases are: 1. O ( nk ) for computing the frame and local (1-hop) 3D features; 2. O ( nkl ) for 1-hop neighborhood message aggregation; 3. O ( nl ) for node-wise tensorization and feature update.

6 Experiments

We evaluate our LEFTNet on both scalar value (e.g. energy) and vector value (e.g. forces) prediction tasks. The scalar value prediction experiment is conducted on the QM9 dataset [32] which includes 134 k small molecules with quantum property annotations; the vector value prediction experiment is conducted on the MD17 dataset [33] and the Revised MD17(rMD17) dataset [34] which includes the energies and forces of molecules. We compare LEFTNet with a list of state-of-the-art equivariant (invariant) GNNs including SphereNet [15], PaiNN [35], Equiformer [36], GemNet [37], etc [28, 38, 13, 39, 40, 14, 16, 41-48].The training details, results on rMD17, etc. are listed in Appendix I.

6.1 QM9 - scalar-valued property prediction

The QM9 dataset is a widely used dataset for predicting molecular properties. However, existing models are trained on different data splits. Specifically, Cormorant [40], EGNN [14], etc., use 100k, 18k, and 13k molecules for training, validation, and testing, while DimeNet [38], SphereNet [15], etc., split the data into 110k, 10k, and 11k. For a fair comparison with all methods, we conduct experiments using both data splits. Results are listed in Table 1. For the first data split, LEFTNet is the best on 7 out of the 12 properties and improves previous SOTA results by 20% on average. Consistently, LEFTNet is the best or second best on 10 out of the 12 properties for the second split. These results validate the effectiveness of LEFTNet. Detailed ablation in Appendix I shows that both LSE and FTE contribute to

4

6

8

10

12

14

100

150

200

250

300

350

SchNet

ClofNet

ComENet

DimeNet++

SphereNet

LEFTNet

MAE

Training time per epoch (s)

the final performance. In addition, we compare the forward time of existing methods to show the efficiency of LEFTNet. Specifically, we use the same batch size for all methods and report the one epoch training time. Fig. 4 shows that LEFTNet can achieve the best performance using similar forward time as SchNet, ClofNet, and ComENet and is much faster than DimeNet and SphereNet. This is consistent with our efficiency analysis in Sec. 5 and Table 4.

6.2 MD17 - vector-valued property prediction

We evaluate LEFTNet to predict forces on the MD17 dataset. Following existing studies [28, 38, 15], we train a separate model for each of the 8 molecules. Both training and validation sets contain 1000 samples, and the rest are used for testing. Note that all baseline methods are trained on a joint loss of energies and forces, but different methods use different weights of force over energy (WoFE). For example, SchNet [28] sets WoEF as 100, while GemNet [37] uses a weight of 1000. For a fair comparison with existing studies, we conduct experiments on two widely used weights of 100 and 1000 following Liu et al. [15]. The results are summarized in Table 2. Results show that when WoFE is 100, LEFTNet

Table 3: Ablation study on Aspirin of the MD17 dataset. Detailed ablations on other molecules are listed in Appendix I

outperforms all baseline methods on 5 of the 8 molecules. In addition, LEFTNet can outperform all baseline methods on 6 of the 8 molecules when WoFE is 1000. These experimental results on MD17 demonstrate the performance of LEFTNet on vector-valued property prediction tasks. The ablation results in Table 3 demonstrate that both LSE and FTE are important to the final results. Specifically, the algorithm for LSE only is in Algorithm 1 of Appendix C. Note that the original MD17 dataset we used has numerical noise [34], and a recomputed version of MD17, called Revised MD17 (rMD17) [34] is proposed to reduce the numerical noise. In addition to the original MD17, we also conduct experiments on the rMD17, and the results are listed in Appendix I.

7 Limitation and future work

In this paper, we seek a general recipe for building 3D geometric graph deep learning algorithms. Considering common prior of 2D graphs, such as permutation symmetry, has been incorporated in off-the-shelf graph neural networks, we mainly focus on the E (3) and SE (3) symmetry specific to 3D geometric graphs. Despite our framework being general for modeling geometric objects, we only conducted experiments on commonly used molecular datasets. It's worth exploring datasets in other domains in the future.

To elucidate the future design space of equivariant GNNs, we propose two directions that are worth exploring. Firstly, our current algorithms consider fixed equivariant frames for performing aggregation and node updates. Inspired by the high body-order ACE approach [49] (for modeling atom-centered potentials), it is worth investigating in the future if equivariant frames that relate to many body (e.g., the PCA frame in [50, 51]) can boost the performance of our algorithm. For example, to build the A-basis proposed in [49], we can replace our message aggregation Eq. 6 from summation to tensor product, which is also a valid pooling operation. Another direction is to explore geometric mesh graphs on manifolds M , where the local frame is defined on the tangent space of each point: F ( x ) ∈ T x M . Since our scalarization technique (crucial for realizing LSE in LEFTNet) originates from differential geometry on frame bundles [52], it is reasonable to expect that our framework also works for manifold data [53, 54].

Acknowledgments and Disclosure of Funding

Y.D., D.F. and C.P.G. are supported by the Eric and Wendy Schmidt AI in Science Postdoctoral Fellowship, a Schmidt Futures program; the National Science Foundation (NSF), the Air Force Office of Scientific Research (AFOSR). L.W. and S.J. are supported by NSF IIS-2243850.

References

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770-778, 2016.

Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv:1609.02907 , 2016.

Xuan Zhang, Limei Wang, Jacob Helwig, Youzhi Luo, Cong Fu, Yaochen Xie, Meng Liu, Yuchao Lin, Zhao Xu, Keqiang Yan, et al. Artificial intelligence for science in quantum, atomistic, and continuum systems. arXiv preprint arXiv:2307.08423 , 2023.

Anthony Simeonov, Yilun Du, Andrea Tagliasacchi, Joshua B Tenenbaum, Alberto Rodriguez, Pulkit Agrawal, and Vincent Sitzmann. Neural descriptor fields: Se (3)-equivariant object representations for manipulation. In 2022 International Conference on Robotics and Automation (ICRA) , pages 6394-6400. IEEE, 2022.

Frank Noé, Alexandre Tkatchenko, Klaus-Robert Müller, and Cecilia Clementi. Machine learning for molecular simulation. Annual review of physical chemistry , 71:361-390, 2020.

Lars Holdijk, Yuanqi Du, Priyank Jaini, Ferry Hooft, Bernd Ensing, and Max Welling. Path integral stochastic optimal control for sampling transition paths. In ICML 2022 2nd AI for Science Workshop .

Michael M Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geometric deep learning: going beyond euclidean data. IEEE Signal Processing Magazine , 34 (4):18-42, 2017.

Kenneth Atz, Francesca Grisoni, and Gisbert Schneider. Geometric deep learning on molecular representations. Nature Machine Intelligence , 3(12):1023-1032, 2021.

Zhengyang Wang, Meng Liu, Youzhi Luo, Zhao Xu, Yaochen Xie, Limei Wang, Lei Cai, Qi Qi, Zhuoning Yuan, Tianbao Yang, et al. Advanced graph and sequence neural networks for molecular property prediction and drug discovery. Bioinformatics , 38(9):2579-2586, 2022.

Arne Schneuing, Yuanqi Du, Charles Harris, Arian Jamasb, Ilia Igashov, Weitao Du, Tom Blundell, Pietro Lió, Carla Gomes, Max Welling, et al. Structure-based drug design with equivariant diffusion models. arXiv preprint arXiv:2210.13695 , 2022.

Yuanqi Du, Tianfan Fu, Jimeng Sun, and Shengchao Liu. Molgensurvey: A systematic survey in machine learning models for molecule design. arXiv preprint arXiv:2203.14500 , 2022.

Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veliˇ ckovi´ c. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478 , 2021.

Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equivariant neural networks for 3D point clouds. arXiv preprint arXiv:1802.08219 , 2018.

Vıctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International conference on machine learning , pages 9323-9332. PMLR, 2021.

Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3D molecular graphs. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=givsRXsOt9r .

Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning , pages 1263-1272. PMLR, 2017.

Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv:1806.01261 , 2018.

Kevin Yang, Kyle Swanson, Wengong Jin, Connor Coley, Philipp Eiden, Hua Gao, Angel Guzman-Perez, Timothy Hopper, Brian Kelley, Miriam Mathea, et al. Analyzing learned molecular representations for property prediction. Journal of chemical information and modeling , 59 (8):3370-3388, 2019.

Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations .

Lingxiao Zhao, Wei Jin, Leman Akoglu, and Neil Shah. From stars to subgraphs: Uplifting any gnn with local structure awareness. In International Conference on Learning Representations .

Fabrizio Frasca, Beatrice Bevilacqua, Michael Bronstein, and Haggai Maron. Understanding and extending subgraph gnns by rethinking their symmetries. Advances in Neural Information Processing Systems , 35:31376-31390, 2022.

Chaitanya K Joshi, Cristian Bodnar, Simon V Mathis, Taco Cohen, and Pietro Liò. On the expressive power of geometric graph neural networks. In The First Learning on Graphs Conference , 2022.

Limei Wang, Yi Liu, Yuchao Lin, Haoran Liu, and Shuiwang Ji. ComENet: Towards complete and efficient message passing for 3D molecular graphs. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. URL https://openreview.net/forum?id=mCzMqeWSFJ .

Ilyes Batatia, David Peter Kovacs, Gregor NC Simm, Christoph Ortner, and Gabor Csanyi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. In Advances in Neural Information Processing Systems .

Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint arXiv:1810.00826 , 2018.

Christopher Morris, Yaron Lipman, Haggai Maron, Bastian Rieck, Nils M Kriege, Martin Grohe, Matthias Fey, and Karsten Borgwardt. Weisfeiler and leman go machine learning: The story so far. arXiv preprint arXiv:2112.09992 , 2021.

Asiri Wijesinghe and Qing Wang. A new perspective on" how graph neural networks go beyond weisfeiler-lehman?". In International Conference on Learning Representations , 2021.

Kristof T Schütt, Huziel E Sauceda, P-J Kindermans, Alexandre Tkatchenko, and K-R Müller. Schnet-a deep learning architecture for molecules and materials. The Journal of Chemical Physics , 148(24):241722, 2018.

Ralf Drautz. Atomic cluster expansion for accurate and transferable interatomic potentials. Physical Review B , 99(1):014104, 2019.

Weitao Du, He Zhang, Yuanqi Du, Qi Meng, Wei Chen, Nanning Zheng, Bin Shao, and Tie-Yan Liu. Se (3) equivariant graph neural networks with complete local frames. In International Conference on Machine Learning , pages 5583-5608. PMLR, 2022.

Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations , 2020.

Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data , 1(1):1-7, 2014.

Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof T Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances , 3(5):e1603015, 2017.

Anders S Christensen and O Anatole Von Lilienfeld. On the role of gradients for machine learning of molecular energies and forces. Machine Learning: Science and Technology , 1(4): 045018, 2020.

Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning , pages 9377-9388. PMLR, 2021.

Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint arXiv:2206.11990 , 2022.

Johannes Gasteiger, Florian Becker, and Stephan Günnemann. GemNet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems , 34: 6790-6802, 2021.

Johannes Gasteiger, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03123 , 2020.

Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (3)-transformers: 3d roto-translation equivariant attention networks. Advances in Neural Information Processing Systems , 33:1970-1981, 2020.

Brandon Anderson, Truong Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. Advances in neural information processing systems , 32, 2019.

Marc Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. In International Conference on Machine Learning , pages 3165-3176. PMLR, 2020.

Tuan Le, Frank Noé, and Djork-Arné Clevert. Equivariant graph attention networks for molecular property prediction. arXiv preprint arXiv:2202.09891 , 2022.

Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications , 14(1):579, 2023.

Dávid Péter Kovács, Cas van der Oord, Jiri Kucera, Alice EA Allen, Daniel J Cole, Christoph Ortner, and Gábor Csányi. Linear atomic cluster expansion force fields for organic molecules: beyond rmse. Journal of chemical theory and computation , 17(12):7696-7711, 2021.

Felix A Faber, Anders S Christensen, Bing Huang, and O Anatole Von Lilienfeld. Alchemical and structural distribution based representation for universal quantum machine learning. The Journal of chemical physics , 148(24):241717, 2018.

Albert P Bartók, Mike C Payne, Risi Kondor, and Gábor Csányi. Gaussian approximation potentials: The accuracy of quantum mechanics, without the electrons. Physical review letters , 104(13):136403, 2010.

Xiang Gao, Farhad Ramezanghorbani, Olexandr Isayev, Justin S Smith, and Adrian E Roitberg. Torchani: a free and open source pytorch-based deep learning implementation of the ani neural network potentials. Journal of chemical information and modeling , 60(7):3408-3415, 2020.

Oliver T Unke, Stefan Chmiela, Michael Gastegger, Kristof T Schütt, Huziel E Sauceda, and Klaus-Robert Müller. Spookynet: Learning force fields with electronic degrees of freedom and nonlocal effects. Nature communications , 12(1):7273, 2021.

Ilyes Batatia, Simon Batzner, Dávid Péter Kovács, Albert Musaelian, Gregor NC Simm, Ralf Drautz, Christoph Ortner, Boris Kozinsky, and Gábor Csányi. The design space of e (3)equivariant atom-centered interatomic potentials. arXiv preprint arXiv:2205.06643 , 2022.

Omri Puny, Matan Atzmon, Edward J Smith, Ishan Misra, Aditya Grover, Heli Ben-Hamu, and Yaron Lipman. Frame averaging for invariant and equivariant network design. In International Conference on Learning Representations , 2021.

Alexandre Duval, Victor Schmidt, Alex Hernandez Garcia, Santiago Miret, Fragkiskos D Malliaros, Yoshua Bengio, and David Rolnick. Faenet: Frame averaging equivariant gnn for materials modeling. arXiv preprint arXiv:2305.05577 , 2023.

Elton P Hsu. Stochastic analysis on manifolds . Number 38. American Mathematical Soc., 2002.

Lingshen He, Yiming Dong, Yisen Wang, Dacheng Tao, and Zhouchen Lin. Gauge equivariant transformer. Advances in Neural Information Processing Systems , 34:27331-27343, 2021.

Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang. Equivariant graph mechanics networks with constraints. arXiv preprint arXiv:2203.06442 , 2022.

Nadav Dym and Haggai Maron. On the universality of rotation equivariant point cloud networks. arXiv preprint arXiv:2010.02449 , 2020.

Sheng Gong, Tian Xie, Yang Shao-Horn, Rafael Gomez-Bombarelli, and Jeffrey C Grossman. Examining graph neural networks for crystal structures: limitations and opportunities for capturing periodicity. arXiv preprint arXiv:2208.05039 , 2022.

Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems , 30, 2017.

Limei Wang, Haoran Liu, Yi Liu, Jerry Kurtin, and Shuiwang Ji. Learning hierarchical protein representations via complete 3d graph networks. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id= 9X-hgLDLYkQ .

Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learning , pages 5361-5370. PMLR, 2020.

Emiel Hoogeboom, Vıctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molecule generation in 3d. In International Conference on Machine Learning , pages 8867-8887. PMLR, 2022.

Xiyuan Wang and Muhan Zhang. Graph neural network with local frame for molecular potential energy surface. arXiv preprint arXiv:2208.00716 , 2022.

Cristian Bodnar, Francesco Di Giovanni, Benjamin Chamberlain, Pietro Liò, and Michael Bronstein. Neural sheaf diffusion: A topological perspective on heterophily and oversmoothing in gnns. Advances in Neural Information Processing Systems , 35:18527-18541, 2022.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An imperative style, high-performance deep learning library. Advances in Neural Information Processing Systems , 32, 2019.

Matthias Fey and Jan E. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds , 2019.

Meng Liu, Youzhi Luo, Limei Wang, Yaochen Xie, Hao Yuan, Shurui Gui, Haiyang Yu, Zhao Xu, Jingtun Zhang, Yi Liu, et al. Dig: A turnkey library for diving into graph deep learning research. The Journal of Machine Learning Research , 22(1):10873-10881, 2021.

Appendix for 'A new perspective on building efficient and expressive 3D equivariant graph neural networks'

Contents

A Supplementary background

In this section, we first detail the concept of scalarization and tensorization and then review the concept of (contravariant) tensor fields and their associated equivariant group representations.

Scalarization. Scalarization is a general technique that originated from differential geometry for realizing covariant operations on tensors [52]. Our method will apply a simple version of scalarization in R 3 to transform equivariant quantities. At the heart of its realization is the notion of equivariant orthonormal frames, which consist of three orthonormal equivariant vectors:

Based on F , we can build orthonormal equivariant frames for higher order tensors by taking tensor products ⊗ , see Eq. 21 in Appendix. By taking the inner product between F and a given equivariant vector (tensor) x , we get a tuple of invariant scalars (see [30] for a proof):

and ˜ x can be seen as the 'scalarized' coordinates of x .

Tensorization. Tensorization, on the other hand, is the 'reverse' process of scalarization. Given a tuple of scalars: ( x 1 , x 2 , x 3 ) , tensorization creates an equivariant vector (tensor) out of F :

The same procedure is extended to higher order cases, see Eq. 22 in Appendix.

A s order (contravariant) tensor T on a vector space V is a multilinear map:

where V ∗ denotes the dual space of V . In fact, there is a canonical 'multiplication' operation between two tensors. Define the tensor product S ⊗ T of two tensors S and T to be a tensor of order r + s :

where v i ∈ V ∗ .

From now on, we assume V = V ∗ = R 3 . Note that when s = 1 , T is exactly an equivariant vector. In practice, the tensor data in R 3 is usually given by its coefficients under a Cartesian coordinate system. Take a second-order tensor as an example, assume we are given an orthonormal frame (basis) ( e 1 , e 2 , e 3 ) and its dual frame ( e 1 , e 2 , e 3 ) , then the nine coefficients of T are given by

In other words, we say the collection { T ij } 1 ≤ i,j ≤ 3 is a faithful representation of T in a fixed coordinate system:

Once defined a tensor on R 3 , it's easy to extend it to a continuous manifold or a discrete graph. A tensor field of order s on a 3D graph G = ( V, E ) is a tensor-valued function f which assigns to each 3D node x i an order s tensor, denoted by f ( x i ) .

SE(3) Tensor Representations. Let V be a vector space, then the group SE(3) is said to act on V if there is a mapping ϕ : SE (3) × V → V satisfying the following two conditions:

if e ∈ SE (3) is the identity element, then

if g 1 , g 2 ∈ SE (3) , then

If we further require ϕ ( g, · ) is a linear map for all g ∈ SE (3) , then ϕ becomes a group representation of SE (3) . From now on, we only consider the rotation subgroup SO(3) and its group representations. When V = R 3 , there is a natural representation of SO (3) by rotating vectors in R 3 . In this way, an element g ∈ SO (3) is identified with a Rotation matrix, denoted by { g j i } 1 ≤ i,j ≤ 3 .

From the tensor definition (10), this natural representation on R 3 induces a tensor representation on T . Still take T = { T ij } 1 ≤ i,j ≤ 3 as an example, we have

for ∀ g ∈ SO (3) . It's easy to check that (12) is indeed a SO(3) representation on the vector space spanned by second-order tensors.

Relation with Spherical Harmonics. For the SO (3) group, all representations (including the tensor representations) can be decomposed as a direct sum of irreducible representations. For each type of irreducible representations, there is a subset of spherical harmonics formulating a basis for this specific representation. However, in terms of representing equivariant geometric quantities, the theorem in [55] claims that tensor representations and irreducible representations are equally powerful: They all form a complete basis in the space of continuous E (3) equivariant functions.

B Equivariant frame construction

Equivariant Frame Definition. An equivariant frame on a 3D graph G consists of three orthonormal vectors that transform equivariantly under SE (3) . In the following, we present specific constructions of equivariant frames. However, it is important to note that the proof of Theorem 4.1 does not depend on our particular choice of local equivariant frames.

It is worth noting that the number of vectors required in the equivariant frame is equal to the dimension of the space. For example, in the case of the four-dimensional Lorentz space, where the symmetry group is replaced by the Lorentz group, the equivariant frame consists of four independent vectors.

Edge-wise frame. From this point onward, we assume that the 3D graph's mass has been translated to zero. This translation ensures that the center of mass of the graph is located at the origin, ensuring the translation invariance of the system. Consider node i and one of its neighbors j with positions x i and x j , respectively. The orthonormal equivariant frame F ij := ( e ij 1 , e ij 2 , e ij 3 ) is defined with respect to x i and x j as follows:

Node-wise frame. Consider node i with 3D position x i , and let ¯ x i := 1 N ∑ x j ∈N ( x i ) x j be the center of mass around the 1-hop neighborhood of x i . The orthonormal equivariant frame F i := ( e i 1 , e i 2 , e i 3 ) is defined with respect to x i as follows:

Note that node frames can also be obtained by averaging the edge-wise frames and applying the GramSchmidt orthogonalization process. This approach provides an alternative method for constructing equivariant frames at the node level.

C Algorithm for the proposed LEFTNet

The detailed LEFTNet is shown in Algorithm 1. Note that for simplicity, the message m in this algorithm could include both invariant and equivariant terms in Fig. 3.

The design of LEFTNet ( LSE only) used in the ablation study is shown in Algorithm 2.

Algorithm 1 LEFTNet

1: Input: 3D graph with equivariant positions X = ( x 1 , . . . , x n ) ∈ R n × 3 , invariant node features h i ∈ R d , invariant relative distances d ij ∈ R 1 , equivariant edge features e ij ∈ R c .

3: for ( i = 1; i < = n ; i ++ ) do

2: Centralize the positions: X ← X -CoM ( X ) .

4: Compute node-wise equivariant frames F i via Eq. 14 .

6: Compute edge-wise equivariant frames F ij via Eq. 13

5: for j ∈ N ( i ) do

7: Get the mutual 3D structure S i -j , perform local scalarization through F ij :

8: Calculate the SE(3)-invariant structural coefficients: A ij = g ( t ij , d ij )

9: Perform equivariant message passing as in Eq. 4:

10: end for

11: Equivariant message aggregation: m i = ∑ j ∈N ( i ) m ij ;

12: Transform equivariant node features through F i :

14: Equivariant Output: Perform tensorization through F i :

13: Update invariant node features:

15: end for

16: Output: for invariant properties (e.g. energy): AvgPooling ( h 1 , . . . , h n ) , for equivariant properties (e.g. force): ( h 1 , . . . , h n ) .

Algorithm 2 LEFTNet ( LSE only)

1: Input: 3D graph with equivariant positions X = ( x 1 , . . . , x n ) ∈ R n × 3 , invariant node features h i ∈ R d , invariant relative distances d ij ∈ R 1 .

3: for ( i = 1; i < = n ; i ++ ) do

2: Centralize the positions: X ← X -CoM ( X ) .

4: for j ∈ N ( i ) do

5: Compute edge-wise equivariant frames F ij via Eq. 13:

6: Get the mutual 3D structure S i -j , perform local scalarization:

7: Calculate the SE(3)-invariant structural coefficients: A ij = g ( t ij , d ij )

8: Perform invariant message passing:

9: end for

10: Update invariant node features:

11: end for

12: Output: AvgPooling ( h 1 , . . . , h n ) .

D Related proofs and discussions of Section 3

In section 3, we proposed a novel hierarchy of local geometric isomorphisms that further motivates the design of incorporating the mutual 3D substructure's information into equivariant GNNs. Different from our fine-grained local characterization, a cocurrent work GWL [22]) proposes to measure geometric isomorphism from local to global by the k-hop partition.

From another point of view, we essentially demonstrated that encoding mutual 3D substructures expands the capacity of the transformation function class with respect to an equivariant GNN. [23] put forward the Completeness concept for characterizing these transformation functions. However, it mainly concentrates on testing whether a function can discriminate global geometric isomorphism (in the sense of Eq. 3).

Discussion on the Completeness Concept. Following our terminology in the preliminary section, completeness of a transformation f can be translated into claiming that f is invariant among 3D graphs if and only if they are globally isomorphic (see definition (3)). Therefore, it's easy to refine the notion of completeness that adapts to our local version by replacing the global isomorphism to local isomorphism:

if and only if X and Y are local {tree-, triangular-, subgraph-} isomorphic. Then, in terms of function class capacities, the following relation holds:

Global complete ⊂ Subgraph complete ⊂ Triangular complete ⊂ Tree complete .

Note that our equivalent description of complete transformation reveals the fact that the completeness concept in [23] is defined from the global 3D isomorphism point of view. Therefore, we shall claim that the above series of completeness notions belong to the structure completeness. Indeed, the theory developed in section 3 indicates that a GNN which can express structure complete functions may not be sufficient in expressing general tensor potential functions on a 3D graph.

On the other hand, a non-negligible proportion of 3D graph tasks may not be sensitive to the global 3D non-isomorphism. For example, some chemical properties (formulated as a function defined on molecular graphs) are characterized by local substructures [56]. In these scenarios, we are looking for a geometric transformation f that is global non-complete, but (tree-, triangular-, subgraph-) local complete.

Proof of Theorem 3.1.

Remark.

When compared to the construction of topological structural coefficients that take the discrete adjacency matrix of the graph as input, as proposed in [27], the input 3D structural coefficients we seek are of a continuous nature. For instance, as illustrated in Fig. 1, the key features that differentiate these 3D local shapes are the dihedral angle and relative distance, both of which are continuous functions of the 3D positions. Therefore, we require the existence of universal approximators for continuous functions to ensure that the parametrization of A ij has sufficient expressiveness.

The reason we desire the LSE module to be invariant is because it ensures the local injectivity of message passing, following the standard argument proposed in Xu et al. [25]. This local injectivity is crucial for proving the second part of Theorem 3.1.

From a mathematical perspective, the common approach to proving the universal approximator property is by demonstrating that a neural network can express a complete basis in the function space, typically leveraging the Stone-Weierstrass theorem. Two successful examples are: a. Piecewise linear function bases achieved by wide Multi-layer perceptrons ( MLP ); b. Polynomial bases. In the 3D GNN setting, equivariant polynomials that are also permutation-invariant can be rigorously constructed using ACE [29]. Additionally, through local scalarization and the Kolmogorov representation theorem [57], the universal approximator property of 3D GNNs reduces to the universal approximator property of MLP , which has been proven for a broad range of non-linear activation functions.

Proof. The first part of the theorem is proved by providing an explicit example. Consider the first pair of 3D shapes in Figure 1, where the mutual 3D substructures S i -q and S j -q ′ both contain two triangles: ( x i , x p , x q ) , ( x i , x m , x q ) and ( x j , x p ′ , x q ′ ) , ( x j , x m ′ , x q ′ ) respectively.

The difference between these two triangular non-isomorphic shapes can be captured by the distance function:

and similarly, d ( x p ′ , x m ′ ) for the second shape. Note that this function cannot be expressed solely by tree-level features since there is no edge connecting x p ( x p ′ ) and x m ( x m ′ ). The fact that this function produces different output values for the two tree isometric but triangular non-isometric 3D shapes implies that ϕ is capable of distinguishing 3D shapes beyond tree isomorphism.

Therefore, if the structural coefficients A ip ( A jp ′ ) take the value of d ( x p , x m ) ( d ( x p ′ , x m ′ ) ), then we know that it can differentiate these two triangular non-isomorphic shapes. Thus, it suffices to prove that the encoder ϕ is a universal approximator of continuous functions (which obviously contain the distance functions) that takes the mutual 3D substructures S i -q ( S j -q ′ ) as input.

According to the Stone-Weierstrass theorem, it is sufficient to demonstrate that ϕ can express a complete basis within the continuous function space. In the 3D GNN setting, this basis should satisfy two additional requirements: 1) permutation equivariance, and 2) SE (3) equivariance. An example of such a basis is the set of E (3) equivariant polynomials that are also permutation equivariant, as demonstrated in Drautz [29]. To obtain invariant coefficients from the combination of equivariant basis functions, we can employ a projection layer, similar to the one used in Dym and Maron [55].

Alternatively, a simpler approach to achieve the universal approximator property is by leveraging local scalarization. Specifically, we use an equivariant frame on the edge e iq (as described in the previous section) to scalarize the mutual 3D structure S i -q ( S j -q ′ ), which includes x p ( x p ′ ) and x m ( x m ′ ). Let ˜ S i -q denote the scalarized 3D positions of S i -q . By quoting the information lossless theorem and the Kolmogorov representation theorem for permutation-invariant functions from Zaheer et al. [57], Du et al. [30], we know that any φ that is a SE (3) and permutation invariant function of S i -q can be expressed as:

where f and g are certain continuous functions that can be approximated by wide MLP , following the universal approximator property of MLP [25].

The second part of the proof demonstrates that there exists an enhanced message passing framework (under injectivity assumptions, which can be realized similarly to [25, 27]) that incorporates the structural coefficients {{ A ij := ϕ ( S i -j ) } e ij ∈ E } to map two distinct local 3D subgraphs with isometric local tree structures in Figure 1 to different embeddings. Since we specifically construct ϕ as an invariant encoder, the collection {{ A ij := ϕ ( S i -j ) } e ij ∈ E } consists of invariant numbers, and the construction of the message passing framework follows a similar approach to [27]. For completeness, we provide a brief review of the injectivity condition and the message passing construction.

To establish the injectivity condition and prove the second part, we introduce the multi-set notation { {·} } , following Xu et al. [25]. A basic equivariant GNN within our enhanced framework consists of at least two steps: 1. Message passing, defined by (4); 2. Node-wise update:

For simplicity, we denote the composition of these two steps as Ψ . Then, the additional injectivity condition requires that Ψ satisfies the following:

is injective for each layer t and each node i . It is worth noting that this condition can be trivially achieved by incorporating weighted residue terms, similar to formula (5) in Wijesinghe and Wang [27]. Consequently, based on the injectivity condition, it is evident that two non-identical collections of {{ A ij } e ij ∈ E } would yield two different embedded feature vectors.

Furthermore, in the first part, we have proven the existence of at least two distinct local 3D subgraphs with isometric local tree structures such that the corresponding geometric weights {{ A ij } e ij ∈ E } generated by the encoder ϕ are different. Combining both parts, we have successfully demonstrated the theorem.

E Related proofs and discussions of Section 4

Torsion Angle is Secretly Hidden in FT

Recall the edge-wise (signed) torsion angle τ ij [58] involves the 1-hop atom pairs i and j and two 2-hop atoms k and l , then τ ij is defined to be the dihedral angle between plane k -i -j and plane l -j -i . Although exhausting all torsion angles requires O ( k 2 ) complexity, Wang et al. [23] reduces the computation to O ( k ) order by selecting a canonical 2-hop atom k and l , which is enough for detecting the relative orientations between atoms (insufficient for general tasks like many body interactions).

Now we show how τ ij naturally appears as one of the derivatives from frame transition functions. For node i , define the equivariant frame F i by

z

z

x

x

y

i

y

i

i

i

x

j

x

j

y

j

y

j

z

j

z

j

F i is normalized through the Gram-Schmidt algorithm. For node j, F j is defined similarly by

Then following the transition formula (5),

Note that for an orthonormal matrix, its inverse is equal to its transpose. Therefore, based on the standard definition of a dihedral angle, we have:

In conclusion, τ ij represents only one component of the transition matrix R ij . To fully determine R ij , we still require two additional angles, as a transition matrix is uniquely determined by three Euler angles.

Proof of Theorem 4.1.

Remark.

The crucial property we leverage to prove this theorem is that all local scalars can be expressed through local scalarization using one local frame and a nonlinear transformation. However, it is important to note that the transition between local frames reveals the relative change of local geometries, which is inherently non-local.

Interestingly, stemming from the fact that invariant scalars can be expressed using scalarization and a universal approximator, we can deduce that by incorporating the frame transition matrix into the union set of S B and S C , all (continuous) invariant interactions f a ( B , C ) can be expressed.

Proof. This theorem is proved in two steps:

The first step characterizes all scalars S B ( S C ) determined by isolated local clusters B and C through equivariant frames and local scalarization;

The second step constructs a specific invariant function f a ( B , C ) that cannot be expressed by taking the union of S B and S C : { S B ∥ S C } .

Let G denote a 3D point cloud. As proven in Du et al. [30], when equipped with an equivariant frame F G , all equivariant features of G can be transformed into scalar features through scalarization without information loss. More precisely, follow the convention in the main text, let ˜ G be the output of performing scalarization (using any equivariant frames purely determined by G , see Section B for examples) on the 3D point cloud G . For any invariant function f ( G ) , there exists a corresponding function ˜ f such that:

Since S G represents the collection of all invariant scalars produced by G , which is equivalent to the collection of all invariant functions that depend on the 3D point cloud G , Eq. 17 implies that S G can be generated by a finite set of invariant scalars ˜ G . To apply this insight to our current theorem, we have two different 3D clouds B and C . Therefore, we need to build two local equivariant frames F B = ( e B 1 , e B 2 , e B 3 ) and F C = ( e C 1 , e C 2 , e C 3 ) . Importantly, the frame F B itself doesn't depend on C , and scalarization through F B is only performed on the local 3D point cloud B . Therefore, performing operations like scalarizing equivariant information of C through F B would violate the assumptions of the theorem.

We are now ready to construct an explicit counterexample f a ( B , C ) using F B and F C :

Since f a is the inner product of two equivariant vectors, it automatically becomes an invariant function. Next, we check whether f a ( B , C ) can be expressed as a function of f a ( { S B ∥ S C } ) . The equivariant component of f a related to B is precisely e B 1 . Applying the above local scalarization principle, we scalarize e B 1 through F B and obtain:

Similarly, e C 1 is also transformed into a constant scalar tuple ˜ e C 1 = (1 , 0 , 0) through F C . As constant inputs generate constant outputs, we conclude that the derived local scalars can only approximate constant functions. However, since the local frames change as we vary the 3D structure of B and C , it is evident that e B 1 · e C 1 is not a constant function of ( B , C ) . Thus, we complete the proof by contradiction.

Similarly, e C 1 is also transformed to a constant scalar tuple ˜ e C 1 = (1 , 0 , 0) through F C . As constant inputs generate constant outputs, we conclude that the deduced local scalars can only approximate constant functions. However, since the local frames are changing as we vary the 3D structure of B and C , it's obvious that e B 1 · e C 1 is not a constant function of ( B , C ) . Therefore, we finish the proof by contradiction.

Further Comments.

It is important to note that the counterexample we constructed is not the only valid counterexample. For instance, the norm of any tensors that depend on both B and C (e.g., ∥ ∥ e B 1 + e C 1 ∥ ∥ , ∥ ∥ e B 1 ⊗ e B 2 ⊗ e C 2 ⊗ e C 3 ∥ ∥ ) can also serve as valid counterexamples.

Let R BC denote the frame transition function between F B and F C , defined in the same way as in Equation 5. The proof of Theorem 4.1 implies that all types of invariant interactions f a ( B , C ) can be expressed solely as functions of the union of three sets S B, S C, and R BC : { S B ∥ S C ∥ R BC } . This is because, through R BC , the 3D point cloud C can be scalarized by the local frame F B . Specifically, let ˜ S B C , ˜ S C C denote the scalarized 3D point cloud of S C with respect to F B and F C , respectively. We have the following diagram:

Utilizing formula 15 on the union of the two 3D point clouds, ( B , C ) , leads us to this conclusion.

Realizing FT by Equivariant Messages: From the FT definition 5, each element of the 3 × 3 matrix R ij is calculated by

Now we show how to reproduce R ij ( k, l ) through equivariant messages. Let the equivariant message m i be the following:

It's easy to check that m i ∈ R 3 × 9 consists of 9 equivariant vectors ( multi-channels ). For atom j , m j is defined symmetrically. For each node, we also store the scalar messages, e.g., ∥ ∥ e i k ∥ ∥ for 1 ≤ k ≤ 3 . Flattening the whole matrix R ij into a R 1 × 9 array, then R ij is obtained by simple summation and taking the vector norm:

where the norm is taken for each column of m i + m j , such that ∥ m i + m j ∥ ∈ R 1 × 9 . Then,

Our illustration also demonstrates the importance of keeping multi-channel tensor messages.

Relation with Previous Equivariant Update Methods. Following the efficiency principle established in section 4, we don't encode the data of the transition matrices explicitly. Instead, we implement tensor messages to fill in the expressiveness gap. Among the tremendously different designs of equivariant graph neural networks, Schütt et al. [35] is closely related to our equivariant updating method. By the above argument, the inner product operation for node i (see (9) of Schütt et al. [35])

can also be reinterpreted as a realization of the (aggregated) frame transition matrix (5).

Moreover, since the equivariant vectors Uv i and Vv i are both aggregated vector features that belong to the same node i and the inner product operation between them is performed in the node-wise updating phase, Schütt et al. [35] actually avoids the 2-hop O ( k 2 ) complexity of computing R xy for all neighborhood node pairs ( x , y ) (while able to express the torsion angle implicitly). For our algorithm, we utilize the scalarization and tensorization in the node-wise updating phase. By the universal approximation theorem 5.1, our method can approximate any inner product operations.

F Related proofs and discussions of Section 5

Equivariant Frames and Higher Order Scalarization and Tensorization. Given an edge e ij with two atom positions ( x i , x j ) , we define our edge-wise SE (3) equivariant frames F ij as follows:

To ensure frame translation invariance, we adopt the approach followed by previous works [59, 60] by restricting the entire 3D conformer space to a linear subspace where the center of mass (CoM) of the system (either the entire system or the sub-cluster to which i and j belong) is set to zero. Alternatively, constructing an E (3) frame is also possible but requires an additional atom position x k , which can be selected using the K-Nearest Neighbor algorithm. If ( x i , x j , x k ) spans the 3D space, we obtain an E (3) equivariant frame by performing Gram-Schmidt orthogonalization. For different constructions of E (3) frames, readers can refer to Wang and Zhang [61]. However, in this work, we focus on SE (3) frames since molecular 3D conformers exhibit SE (3) symmetry but do not possess reflection symmetry.

Once we have an equivariant frame, every vector is a linear combination of the three orthogonal vectors in the frame. Moreover, the unique combination coefficients are exactly the 'scalarized' coordinates in (8). A similar procedure also applies to higher order tensors. Indeed, the vector frame F 1 extends to a tensor frame F r of arbitrary order r > 1 :

Since the orthonormal frame F r is complete in the sense that it spans the whole tensor space of order r , every r-th order tensor admits a unique decomposition:

It's easy to prove that the collection { T i 1 ,...,i r } 1 ≤ i 1 ,...,i r ≤ 3 consists of invariant scalars. We call the process from T to { T i 1 ,...,i r } 1 ≤ i 1 ,...,i r ≤ 3 scalaraization .

Tensorization is the inverse of scalarization, in the sense that it sends scalars { T i 1 ,...,i r } 1 ≤ i 1 ,...,i r ≤ 3 to tensor T . Under the same frames we use during scalarization, the following diagram demonstrates the pipeline of producing L second-order tensors out of { T i 1 i 2 j } 1 ≤ i 1 ,i 2 ≤ 3 :

(22)

where ⊙ denotes the element-wise product.

Proof of Theorem 5.1

Proof. The proof relies on the invertibility of Scalarization and Tensorization operations, as demonstrated in Appendix A.5 of Du et al. [30]. This invertibility allows us to establish a commutative diagram as follows:

This diagram illustrates that for each mapping ρ , there exists a corresponding "scalarized" mapping ˜ ρ given by:

By applying Tensorize , followed by ρ , and then Scalarize , we obtain an invariant representation ˜ ρ . Since MLP serves as a universal approximator of invariant functions, we can always find an MLP that expresses ˜ ρ . By reversing the arrows, we conclude the proof.

Proof of Equivariance for LEFTNet LEFTNet consists of multiple layers of LSE and FTE . LSE is realized by scalarization, and FTE is realized by scalarization and tensorization. Since the invariance of scalarization and the equivariance of tensorization have been proved, we finish the proof.

G Extended Related Work

Based on the discussions in Section 3 and 4, we identify two essential components for constructing expressive equivariant 3D Graph Neural Networks (GNNs): (1) Local 3D Substructure Encodings ( LSE ), which enable the local message passing to capture diverse local 3D structures; and (2) Frame Transition Encodings ( FTE ), which incorporate equivariant coordinate transformations between different local patches into the 3D GNN.

G.1 Modular overview of 3D GNN

Based on the general notion of Local 3D Substructure Encodings ( LSE ) and Frame Transition Encodings ( FTE ), we have provided concrete constructions to realize them. However, it is important to note that there are also implicit methods to encode the information of LSE and FTE . In order to provide a comprehensive overview, we review previous 3D Graph Neural Networks (GNNs) that follow this framework and summarize the findings in Table 4. To ensure a fair comparison, we include the computational complexity as it is often a trade-off with expressiveness. A detailed analysis of this trade-off is provided at the end of Section G.2.

Regarding Local 3D Substructure Encodings ( LSE ), SphereNet [15] and GemNet [37] (implicitly) encode local 3D substructures through a computation-intensive edge-based update. We indicate these two architectures with a ∗ in Table 4 due to the following reasons: In the case of GemNet [37], the message passing framework considers the 1-hop neighborhood with respect to both end points of an edge e ij , which inherently contains all the nodes of the mutual 3D subgraph S i -j . However, when compared to SphereNet [15], which encodes all node positions under an angular coordinate system, GemNet [37] only captures partial geometric information. On the other hand, SphereNet [15] utilizes an edge-based update that aggregates 1-hop neighbors with respect to the source node. While this approach considers the immediate neighbors, it may not include all the nodes inside the subgraph S i -j (e.g., the neibors of the target node).

Regarding Frame Transition Encodings ( FTE ), we have demonstrated (see the proof of Theorem 4.1) that the essential information of FTE does not rely on the specific choices of local equivariant frames. Therefore, most 3D GNNs with equivariant vector updates that can incorporate at least three ( multi-channel ) independent equivariant vectors as edge features are capable of expressing local frame transitions ( FT ). However, EGNN [14] is an exception as it only updates the position vector (i.e., one channel), which is insufficient for capturing the entire FT . In other words, the power of the update function ϕ in Equation (6) also influences the encoding of FT . Additionally, models that encode torsion angle information partially express FTE , as illustrated in Appendix E. While there are multiple ways to realize LSE and FTE , there is a trade-off between efficiency and expressiveness in

terms of the number of hops considered for message passing, as indicated in the last column of Table 4.

In contrast to our invariant realization of LSE , Batatia et al. [24] constructs their framework by building a complete ( E (3) + permutation) equivariant polynomial basis using spherical harmonics and tensor product, where the monomial variables are a combination of 3D features from different nodes (bodies). On the other hand, we achieve the function of LSE and FTE through edgewise scalarization A ij and equivariant message passing (see Fig. 3). While Batatia et al. [24] utilize the atomic cluster expansion (ACE) mechanism, we provide an illustration of how to equivariantly realize LSE based on local ACE in the proof of Theorem 3.1. Another work Wang and Zhang [61] uses an ensemble of frames to model the local environment defined through a distance cutoff, which is suitable for pure 3D geometry. However, our local 3D hierarchy incorporates well with the 2D topology. Therefore, whether [61] encodes LSE depends on the radius of the local environment and whether the ensemble of 'frames' purely built as a combination of radical directions (see Equation (6) in [61]) spans R 3 . In comparison to geometric GNNs like SphereNet [15] and GemNet [37], which encode reflection-antisymmetric torsion angles, and LeftNet, which implements SE (3) equivariant, reflectionantisymmetric frames, the ensemble of frames in [61] is reflection symmetric and therefore cannot differentiate local chemical isomers. Additionally, Wang and Zhang [61] defines the '(ensembled) frame to frame' projection, which is somehow equivalent to our FTE , and utilizes it to detect the global isomorphism from local observations. In contrast, our FTE is employed to bridge the expressiveness gap of global continuous functions, such as regression tasks on graphs (see the comments in Section E), rather than focusing on detecting global isomorphisms, as in classification tasks on graphs. Moreover, we have demonstrated that using one SE (3) equivariant frame for each local patch is sufficient for achieving local expressiveness, considering efficiency as well. Finally, we provide another invariant realization of FTE , inspired by [62], which is presented at the end of Section H.

G.2 Relationship to Geometric WL test

Recently, Joshi et al. [22] propose a geometric k-WL test (GWL) to assess the expressiveness of geometric GNN algorithms. In essence, our tree isomorphism aligns with the 1-hop geometric isomorphism introduced in GWL, while the fine-grained triangular isomorphism falls between the 1-hop and 2-hop geometric isomorphism described in GWL. From a model design perspective, we achieve the realization of LSE through local scalarization, which guarantees expressiveness through the Kolmogorov representation theorem ( Zaheer et al. [57]) and the universal approximator property of MLP. Different from our geometry perspective, the fundamental concepts of body order and tensor order, which stem from classical inter-atomic potential theories and exhibit equivariance, play a crucial role in measuring expressive power in Joshi et al. [22]. In addition to the local geometric isomorphism hierarchy, we discover the significance of FTE as the connecting bridge between local invariant scalars and global geometric expressiveness. This realization, along with LSE on mutual 3D substructures, sheds light on the insufficiency of the 1-hop local scalarization implemented in ClofNet ( Du et al. [30]). We further explore the connection between FTE and the neural sheaf Laplacian in Section H.

Computational efficiency of LEFTNet In the analysis of a graph with n nodes and an average of k edges per node, the computational efficiency of a message passing-based graph neural network (MPNN) is typically described by the following form:

where f represents the computational cost of the node-level updating function (usually an MLP), and g represents the computational cost of calculating the message for each edge. The power c is a crucial factor that determines the computational cost of performing each message passing operation. For instance, SphereNet (Liu et al., 2021) aggregates all 1-hop nodes with respect to the source node of an edge to construct a message, resulting in a complexity of O ( nk 2 ) for message passing.

To analyze the computational complexity of LEFTNet, we can divide it into the following computational steps:

Building the S i -j and scalarizing with respect to edge frame F ij . The 3D subgraph S i -j is constructed by intersecting the 1-hop neighbors N i and N j , which requires O ( nk )

computations. Additionally, collecting the 1-hop neighbors for each node has a complexity of O ( nk ) . Finally, scalarizing each S i -j with the edge-wise frame F ij (which involves taking inner products) and extending the result to m channels has a complexity of O ( mnk ) . Therefore, the total complexity for this step remains at O ( mnk ) .

Equivariant message passing. From the first step, we get the scalarized ˜ S i -j from the 3D subgraph S i -j . Then, the edge-wise message is obtained by transforming the scalarized ˜ S i -j with a MLP , and the output invariant scalars is multiplied with both the invariant and equivariant node features (let d denotes the dimension of the node feature, then the multiplication's complexity is O ( dnk ) . Let l denotes the complexity of MLP (depends on the depth and hidden dimension of MLP ), then the message passing takes O (( d + l ) nk ) computation.

Node level updating. This step involves nodewise scalarization, applying an MLP , and performing tensorization. Suppose the input and output channel numbers for equivariant tensor features are both d , then the scalarization using the node frame (similar to step 2) has a complexity of O ( dn ) computations. Similarly, tensorization, which is the inverse operation of scalarization, also has a complexity of O ( dn ) . Therefore, the total complexity for this step is O (( l + d ) n ) .

In conclusion, we have analyzed the three steps required for performing one layer of LEFTNet. It is worth noting that the efficiency of LEFTNet is comparable to other 1-hop based GNN algorithms like SchNet [28] ( c = 1 in Eq. 23). The final complexity is O ( ldn + ( d + m + l ) nk ) , which is summarized in Table 4.

H Neural sheaf interpretation

Since our invariant scalarization, utilized in the LSE module, shares similarities with the scalarization technique on vector sheaves (see Hsu [52]), it is natural to explore a potential interpretation of LEFTNET within the framework of neural sheaf diffusion proposed by Bodnar et al. [62]. In this section, we aim to modify the neural sheaf diffusion architecture to preserve SE (3) equivariance by incorporating our node and edge frames.

We begin by revisiting the concept of a neural sheaf and its associated sheaf Laplacian operator. A cellular sheaf over a discrete graph is a mathematical object that assigns a vector space to each node and edge in the graph and specifies a linear map between these spaces for each incident node-edge pair:

Definition H.1. A cellular sheaf ( G, F ) on a graph G = ( V, E ) consists of:

A vector space F ( v ) for each v ∈ V .

A vector space F ( e ) for each e ∈ E .

A linear map F v ⊴ e : F ( v ) →F ( e ) for each incident v ⊴ e node-edge pair.

In the case of a 3D graph neural network (GNN), the node-wise features typically consist of tensors (as defined in 10), which naturally form vector spaces. The key aspect of defining a sheaf lies in determining the specific assignments of the linear maps F v ⊴ e for each node v and each edge e . For non-geometric GNNs, F v ⊴ e captures how the opinions (represented as vectors in F ( v ) ) manifest in a "discourse space" formed by F ( e ) . In our approach, we equip each node with an equivariant frame denoted by F v and assign an equivariant frame to each edge denoted by F e . By definition, F v and F e capture the local geometry around the node and the edge, respectively. Thus, a natural choice for F v ⊴ e is given by:

To ensure SE (3) equivariance in the sheaf structure, we define F ( v ) as the set of scalarized vectors obtained by applying scalarization using F v . It is important to note that the linear map F v ⊴ e between vectors can be extended to a linear map between higher-order tensors through tensor product operations (see Section F). This construction guarantees the entire structure to be SE (3) -invariant. With a well-defined sheaf structure, we can introduce the sheaf Laplacian operator that defines a

transformation between node features:

In our case, the sheaf Laplacian can be expressed as:

using the properties of orthogonal transformations. Thus, this operator quantifies the collective "disagreement of local geometries" at each node. Neural sheaf diffusion propagates information through nodes using the following partial differential equation (PDE):

Here, X represents the node feature matrix. In accordance with the nonlinear parametrization proposed in [62], we have the following model:

where any applicable nonlinear activation function σ can be utilized since our cellular sheaf and its Laplacian are invariant. It is important to note that the key component of L F (denoted by F T v F u in Wq. 24) corresponds to the frame transition ( FTE ) module between node frames. Hence, Equation (26) provides an invariant realization of the FTE module based on equivariant node frames. For equivariant outputs, tensorization can be applied after the final layer of Equation (26). To incorporate the LSE module, we can include the LSE information as a part of invariant node features to F v prior to the message aggregation step in Equation (26).

I Additional experiments

Experiment Detail for Table 1 and 2. For QM9, baseline results are taken from Liao and Smidt [36]. For MD17, baseline results are taken from the original papers (with unit conversions if needed). All models are trained on energies and forces, and WoFE is the weight of force over energy in loss functions. All experiments are conducted on a single NVIDIA GeForce RTX 2080 Ti 11GB GPU. Our implementation is based on the libraries including PyTorch [63], PyG [64], and DIG [65].

Ablation Study. As discussed in Section 5, there are two main modules in LEFTNet, namely LSE and FTE . We conduct experiments on QM9 and MD17 to show the importance of each component. Experimental results are summarized in Table 5 and Table 6. The results show that using LSE can outperform the model without both LSE and FTE on all tasks. Adding FTE can further improve the performance. The results demonstrate the importance of LSE and FTE modules.

Results on rMD17. Following Batatia et al. [24], we conduct experiments on rMD17 to compare with recent studies. Results show that our LEFTNet can achieve comparable performance to stateof-the-art methods such as MACE and NequIP, while outperforming other baseline methods like GemNet and PaiNN.

Model and training hyperparameters. Model and training hyperparameters for our method on different datasets are listed in Table 8.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents the expressive power of equivariant graph neural networks in the context of 3D local isomorphism from a 2D local isomorphism perspective. Similar to the subgraph isomorphism of GNNs, the paper proposes three types of isomorphism in 3D space: Tree Isometric, Triangular Isometric, Subgraph Isometric. Based on the above definitions, the paper defines geometric Weisfeiler-Lehman (WL) tests. Based on the conclusions in paper [1], this paper enhances the expressive power of equivariant GNN by incorporating mutual 3D substructures. On the other hand, to further enhance the expressive power of the equivariant graph neural network, this paper introduces an SE(3)-invariant encoder to exploit mutual 3D structures. Through a theoretical analysis (theoreom 4.1), the paper demonstrates that simple aggregation of local messages cannot approximate global interactions, thus the authors further enhance the model's expressive power through the frame transition matrix. The effectiveness of the proposed model is demonstrated by experimental results on scalar and vector properties of molecules on QM9.\n\n[1] A New Perspective on \"How Graph Neural Networks Go Beyond Weisfeiler-Lehman?\"\n\nSTRENGTHS:\n1. To the best of my knowledge, this is the first paper to analyze the expressive power of equivariant neural networks from the 3D isomorphism perspective. This could inspire future work to analyze the expressive power of current equivariant neural networks from this perspective.\n\n2. Based on the analysis proposed in the paper on the expressivity of equivariant neural networks, the paper incorporates mutual 3D substructures into 3D GNNs (LSE in this paper) to further enhance their expressive power. The effectiveness of the proposed analysis and methods has been demonstrated experimentally. \n\n3. To reduce the computational cost of neural sheaf diffusion, the paper proposes a novel frame transition matrix to express global geometric information. The computational complexity analysis in line 304 validates the effectiveness of the proposed method.\n\nWEAKNESSES:\n1. I have to say that the notation in this paper is confusing, which imposes a burden on me. For instance, h is defined as a node feature on line 81, and f is a tensor-valued function on line 70. However, f is then used as a bijection on line 112. Does this mean the definition of h on line 112 is different from that on line 81? \n\n2. Although the analysis of the expressivity of 3D GNNs is the first from the 3D isomorphism perspective, this paper seems to be based on the analysis in [1]. Moreover, the method proposed by the analysis in [1] can incorporate structural properties into many GNNs, such as GCN, GIN, and GAT. Can this paper similarly incorporate mutual 3D substructures into other 3D GNNs?\n\n3. Recently, some work [2] has analyzed the expressive power of GNNs from the perspective of the attention mechanism, so how can we analyze the expressive power of 3D GNNs from the attention mechanism? \n\n\n\n[1] A New Perspective on \"How Graph Neural Networks Go Beyond Weisfeiler-Lehman?\"  \n\n[2] Interpretable and Generalizable Graph Learning via Stochastic Attention Mechanism\n\nQUESTIONS:\nSee my comments on weakness.\n\nLIMITATIONS:\nIn my personal opinion, one of the limitations of this paper is that the notation is a bit confusing, and the paper's layout is somewhat disorganized, making it hard to follow. I would suggest that the authors provide a guideline at the end of the introduction to help readers understand the role of each section.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces local structure encoding and frame transition encoding for more expressive representation learning in 3D graph neural networks. The local structure encoding is inspired by observations in the proposed local hierarchy of 3D graph isomorphisms; specifically the observation that subgraph isomorphism has greater discriminative power than triangle or tree isomorphism. The frame transition encoding is inspired by the observation that not all types of invariant interactions between disjoint clusters can be expressed as functions of the union of those clusters. The authors show that the proposed 3D GNN exceeds the performance of or performs on par with a multitude of existing models.\n\nSTRENGTHS:\n* Originality: The proposed work appears original. The authors introduce a classification of local and global structure, and use it to classify existing 3D GNNs. They also show how to incorporate the most expressive types of local and global structural elements into a new 3D GNN and show that it outperforms/performs on par with existing methods.\n* Quality/Clarity: The proposed work is of good quality. The paper is well written, preliminary ideas are presented in an accessible way; the local hierarchy of 3D graph isomorphisms is paired with a very nice figure; the work is well motivated; the proposed model outperforms/performs on par with existing methods.\n* Significance: The problem of building expressive representations in GNNs is challenging and of significance to the community.\n\nWEAKNESSES:\n* Quality: see questions\n\nQUESTIONS:\n* SpookyNet is not cited or discussed although it shows very strong performance in the analysis. This method appears to have local and global features, it would be interesting to see where it lands in the proposed classification.\n\nLIMITATIONS:\nThe authors discuss some limitations of the method", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

comments: SUMMARY:
Authors analize the expressive power of 3D equivariant GNNs and introduce a new expressive equivariant GNN architecture based on local node and edge wise frames.

STRENGTHS:
The proposed way to achieve equivariance is well grounded and conceptually very interesting and novel.
The construction is also based on solid theoretical motivation and achieved good results in practical experiments.
The paper is well written

WEAKNESSES:
The theory about the local hierarchy of 3D Graph Isomorphism (section 3) relatively straightforwardly follows from existing works on equivariant GNN expressive power and subgraph GNN expressive power. That being said, the overall package is still very solid.

QUESTIONS:
The local frame construction reminded me a bit of Frame Averaging by Puny et al. (ICLR 2022). It would be interesting to see how they stack up.

LIMITATIONS:
Authors have addressed the limitations
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates expressive message passing architectures for processing 3D geometric graphs with permutation and Euclidean symmetries. The authors first provide an investigation of a local hierarchy of isomorphism separability. Then, the authors show cases where powerful local invariant models may fall short of encoding global structures due to information loss during propagation, and aim to reduce the information loss by turning to equivariant message passing incorporating frame transition matrices that are related to neural sheaf diffusion. Based on the theoretical framework, the authors propose a message passing architecture that is expressive in terms of the local isomorphism separability as well as incorporating frame transition. Experimental results on invariant molecular property prediction on QM9 and equivariant force prediction on MD17 shows that the proposed approach outperforms previous equivariant neural networks.\n\nSTRENGTHS:\nS1. This paper addresses an important problem of understanding the theoretical expressiveness of geometric graph neural networks and developing architectures that are provably expressive.\n\nS2. The local isomorphism hierarchy proposed in Section 3 is novel as far as I can tell (in context of 3D geometric graphs) and correct.\n\nS3. The visual example shown in Figure 1 was helpful in understanding the proposed hierarchy of local isomorphism.\n\nS4. The connection of the proposed approach to sheaf diffusion mentioned in Appendix I was interesting.\n\nS5. The experimental results and ablation study on QM9 and MD17 seems to support the main arguments.\n\nWEAKNESSES:\nW1. The equality in Eq. 3 seems not correct as it neglects permutation symmetry. If we consider positional features represented as matrices $\\mathbf{X}\\in\\mathbb{R}^{n\\times 3}$ (Line 77), the precise symmetry under consideration is a direct product of permutation group and Euclidean group $S_n \\times SE(3)$ (Dym et al., 2020). I think it should be clarified that the equality in Eq. 3 is up to permutation of nodes.\n\nW2. In Line 111, I think  {-tree, -triangular, -subgraph} should be {tree-, triangular-, subgraph-}. Same applies to Line 603-604 and Line 612.\n\nW3. In the definition of tree and triangular isometries (Line 114-118), is there any reason to use the term tree and triangular? I think the definitions are a straightforward extension of the subtree and overlap isomorphism proposed in Wijesinghe et al., 2022, respectively, and I see no reason to use different name. Especially, for the triangular isometry, I found the name misleading since the isometry condition is imposed on the entire set of triangles that share $e_{ui}$, rather than individual triangle.\n\nW4. In the definition of triangular isometry (Line 116), the condition is specified for each edge in the intersection of the local subgraphs $e_{iu}\\in\\mathbf{S}\\_{i-j}$. Please correct me if I am wrong, but to make sense, I think the condition should be for each edge in one of the local subgraphs, $e_{iu}\\in\\mathbf{S}\\_{i}$.\n\nW5. In Line 606, the term structure completeness is used without proper definition (although it is emphasized by boldface). Without clear definition, I was not able to properly understand and correctness of the argument in Line 608.\n\nW6. The argument in Line 162-164 is ambiguous; do you mean a composition of Atoministc Cluster Expansion and invariant map constructed using scalarization by edge-wise equivariant frames is a universal approximator of SE(3) invariant functions, so it can serve as $\\phi$ in Theorem 3.1? Is the proof of universality available in literature?\n\nW7. For LSE, based on Appendix H, it seems that the hierarchy of local isomorphism proposed in this paper is a subcase of the geometric WL test proposed in Joshi et al., 2023. Since this paper provides 3-way classification while GWL provides a more sensitive separation of expressive power of geometric GNNs based on the orders of node tuple and feature tensors, one can argue that the local isomorphism hierarchy proposed in this paper is less of practical interest compared to GWL.\n\nW8. For FTE, in Section 4, frame transition is introduced to solve the limitation of SE(3) invariant messages as outlined in Theorem 4.1. But as far as I know, a majority of practical equivariant geometric GNNs in literature are already using equivariant (tensor) messages (Thomas et al., 2018; Fuchs et al., 2020; Satorras et al., 2022; Geiger et al., 2022; Kohler et al., 2019), and the expressive power of some of them (TFN (Thomas et al., 2018) and SE(3)-Transformers (Fuchs et al., 2020)) are shown to be universal i.e., maximally expressive (Dym et al., 2020). Is the motivation of frame transition also relevant for these architectures? Also, is there specific reasons to favor the proposed frame based propagation over these already equivariant message passing algorithms, despite the added complexity of incorporating node- and edge-wise frames?\n\nW9. In Line 214, does the frame here refer to frames in Riemannian manifold theory? Since there is no proper definition (Eq. 5 provides a description of key property but is not a definition), I was confused on what exactly are the frames mentioned here and was unable to verify whether Line 216-221 is correct. It seems the precise definition and descriptions of equivariant frame are completely deferred to Appendix B and F although it serves as a main component of the proposed architecture, which makes the paper hard to read and understand; I think key parts of the proposed algorithm should be described in the main text in a self-contained manner.\n\nW10. Overall, I find that the writing of the paper has room for improvement in terms of organization and readability. There are many cross-references back and forth across the paper, so I couldn't serially read through each section without going through a large portion of the entire text multiple text. For example, Line 236-237 (Section 4) refers to Figure 3 (Section 5), Line 555 (Appendix A) refers to Eq. 21 (Appendix F), Line 592 (Appendix C) refers to Figure 3 (Section 5), Algorithm 1 (Appendix C) refers to Eq. 4 (Section 3), and so on.\n\nW11. In Eq. 7, I don't see how the equation describes equivariance. I think it should be something like $\\mathbf{m}(g\\mathbf{x}\\_u) = \\sum\\_{i=0}\\^l\\mathcal{M}\\^i(g)\\mathbf{m}\\_i(\\mathbf{x}\\_u)$, is this a typo?\n\nW12. Typo in Line 1 of Algorithm 1, gragh -> graph\n\nW13. Specification of the output is missing in Algorithm 1, and the output in Algorithm 2 is not clear since normal and boldfaced h are mixed in the algorithm. Furthermore, In Line 7 of Algorithm 1, how do you scalarize the mutual 3D structure, which is a subgraph, based on edge frame $F_{ij}$? I don't see how this can be done in Eq. 8 which defines scalarization. Similarly, in Line 6 of Algorithm 2, how do you scalarize the mutual 3D structure without frame? Also, if the algorithm 2 is LSE only, why should one compute edge-wise frames in Line 5?\n\nW14. In the proof of Theorem 5.1, I can see that the proposed parameterization is universal, but does it correctly guarantee equivariance as well (as specified in Theorem 5.1)? The current proof seems to clearly show universality, but I am not sure about equivariance.\n\nDym et al., On the universality of rotation equivariant point cloud networks (2020)\n\nWijesinghe et al., A new perspective on \"how graph neural networks go beyond weisfeiler-lehman\" (2022)\n\nJoshi et al., On the Expressive Power of Geometric Graph Neural Networks (2023)\n\nThomas et al., Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds (2018)\n\nFuchs et al., SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks (2020)\n\nSatorras et al., E(n) Equivariant Graph Neural Networks (2022)\n\nGeiger et al., e3nn: Euclidean Neural Networks (2022)\n\nKohler et al., Equivariant Flows: sampling configurations for multi-body systems with symmetric energies (2019)\n\nQUESTIONS:\nThe questions are included in the weakness section.\n\nLIMITATIONS:\nThe authors have addressed the limitations of their work in Section 7.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-14"}