SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Weitao Du*  
Chinese Academy of Sciences  
duweitao@mass.ac.cn
Yuanqi Du* Cornell University yd392@cornell. edu
Limei Wang*  
Texas A&M University  
limei@tamu.edu
Dieqiao Feng Cornell University
Guifeng Wang  
Zhejiang University
Shuiwang Ji  
Texas A&M University
Carla P Gomes Cornell University
Zhi-Ming Ma Chinese Academy of Sciences
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.
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,
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.
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 permutation-invariant graph neural networks. The formal way of describing Euclidean symmetry is the group  $E(3) = O(3) \rtimes 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) \rtimes 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.
Equivalence. A tensor-valued function  $\varphi(\mathbf{x})$  is said to be equivariant with respect to  $SE(3)$  if for any translation or rotation  $g \in SE(3)$  acting on  $\mathbf{x} \in \mathbf{R}^3$ , we have
$$
\varphi (g \mathbf {x}) = \mathcal {M} (g) \varphi (\mathbf {x}),
$$
where  $\mathcal{M}(\cdot)$  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.,  $\mathbf{x}$  as a position vector. It is worth noting that when  $\varphi(\mathbf{x}) \in \mathbf{R}^1$  and  $\mathcal{M}(g) \equiv 1$  (the constant group representation), the equivariant function  $\varphi(\mathbf{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} \in V$  denotes the set of nodes (vertices, atoms), and  $e_{ij} \in E$  denotes the set of edges. For brevity, the edge feature attached on  $e_{ij}$  is also denoted by  $e_{ij}$ .  $\mathbf{X} = (\mathbf{x}_1, \dots, \mathbf{x}_n) \in \mathbf{R}^{n \times 3}$  denotes the node positions which determine the geometric structure of  $G$ .
(b)
Figure 1: (a)  $\mathbf{S}_i$  and  $\mathbf{S}_j$  share the same tree structure (edge lengths are identical), but they are not triangular isomorphic (different dihedral angles); (b)  $\mathbf{S}_i$  and  $\mathbf{S}_j$  are triangular isomorphic but not subgraph isomorphic (the relative distance between the two triangles is different).
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_i^l, h_j^l$  denote the  $l$ -th layer's node features of source  $i$  and target  $j$  that also depend on the 3D positions  $(\mathbf{x}_i, \mathbf{x}_j)$ , then the aggregated message is
$$
m _ {i} ^ {l} = \bigoplus_ {j \in \mathcal {N} (i)} m _ {i j} \left(h ^ {l} \left(\mathbf {x} _ {i}\right), h ^ {l} \left(\mathbf {x} _ {j}\right), e _ {i j} ^ {l}\right), \tag {1}
$$
and  $\bigoplus_{j\in \mathcal{N}(i)}$  is any permutation-invariant pooling operation between the 1-hop neighbors of  $i$ . We also include the edge features  $e_{ij}^{l}$  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}$ :
$$
m _ {i j} \left(h ^ {l} \left(g \mathbf {x} _ {i}\right), h ^ {l} \left(g \mathbf {x} _ {j}\right)\right) = \mathcal {M} (g) m _ {i j} \left(h ^ {l} \left(\mathbf {x} _ {i}\right), h ^ {l} \left(\mathbf {x} _ {j}\right)\right), \tag {2}
$$
where  $g \in SE(3)$  is acting on the whole geometric graph simultaneously:  $(\mathbf{x}_1, \ldots, \mathbf{x}_n) \to (g\mathbf{x}_1, \ldots, g\mathbf{x}_n)$ . For example, the invariant model ComENet [23] satisfies Eq. 2 by setting  $\mathcal{M}(g) \equiv 1$  and MACE [24] realized Eq. 2 for nonconstant irreducible group representations  $\mathcal{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  $\mathbf{X},\mathbf{Y}$  are defined to be globally isomorphic, if there exists  $g\in SE(3)$  such that
$$
\mathbf {Y} = g \mathbf {X}. \tag {3}
$$
In other words,  $\mathbf{X}$  and  $\mathbf{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  $\mathbf{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} \in E$ , the mutual 3D substructure  $\mathbf{S}_{i-j}$  is defined by the intersection of  $\mathbf{S}_i$  and  $\mathbf{S}_j$ :  $\mathbf{S}_{i-j} = \mathbf{S}_i \cap \mathbf{S}_j$ .
Given two local subgraphs  $\mathbf{S}_i$  and  $\mathbf{S}_j$  that correspond to two nodes  $i$  and  $j$  (not necessarily adjacent), we say  $\mathbf{S}_i$  is {tree-, triangular-, subgraph-} isometric to  $\mathbf{S}_j$ , if there exists a bijective function
$f:\mathbf{S}_i\to \mathbf{S}_j$  such that  $h_{f(u)} = h_u$  for every node  $u\in \mathbf{S}_i$ , and the following conditions hold respectively:
- Tree Isometric: If there exists a collection of group elements  $g_{iu} \in SE(3)$ , such that  $(\mathbf{x}_{f(u)}, \mathbf{x}_{f(i)}) = (g_{iu}\mathbf{x}_u, g_{iu}\mathbf{x}_i)$  for each edge  $e_{iu} \in \mathbf{S}_i$ ;
- Triangular Isometric: If there exists a collection of group elements  $g_{iu} \in SE(3)$ , such that the corresponding mutual 3D substructures satisfy:  $\mathbf{S}_{f(u) - f(i)} = g_{iu}\mathbf{S}_{u - i}$  for each edge  $e_{iu} \in \mathbf{S}_i$ ;
- Subgraph Isometric: for any two adjacent nodes  $u, v \in \mathbf{S}_i$ ,  $f(u)$  and  $f(v)$  are also adjacent in  $\mathbf{S}_j$ , and there exist a single group element  $g_i \in SE(3)$  such that  $g_i \mathbf{S}_i = \mathbf{S}_j$ .
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  $\Rightarrow$  Triangular Isometric  $\Rightarrow$  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  $\nRightarrow$  Triangular Isometric  $\nRightarrow$  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_{ij}^{t} = \mathbf{RBF}(d(\mathbf{x}_{i},\mathbf{x}_{j}))$ , where  $\mathbf{RBF}(\cdot)$  is a set of radial basis functions. Although it is powerful enough for testing tree nonisomorphism (assuming that  $\mathbf{RBF}(\cdot)$  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  $\phi$ , define the 3D structure weights  $A_{ij} \coloneqq \phi(\mathbf{S}_{i-j})$  for each edge  $e_{ij} \in E$ . Then, the message passing framework (1) is generalized to:
$$
m _ {i} ^ {l} = \bigoplus_ {j \in \mathcal {N} (i)} m _ {i j} \left(h ^ {l} \left(\mathbf {x} _ {i}\right), h ^ {l} \left(\mathbf {x} _ {j}\right), A _ {i j} h ^ {l} \left(\mathbf {x} _ {j}\right), e _ {i j} ^ {l}\right). \tag {4}
$$
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  $\phi$  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} := \phi(S_{i-j})\}_{e_{ij} \in 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  $\phi$  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  $\mathbf{B}$ ,  $\mathbf{C}$ . Suppose the ground-truth interaction potential of  $\mathbf{B}$  and  $\mathbf{C}$  imposed on atom  $a$  is described by a tensor-valued function  $f_{a}(\mathbf{B},\mathbf{C})$ . Since  $\mathbf{B}$  and  $\mathbf{C}$  are both beyond the 1-hop neighborhood of  $a$ , the information of  $f_{a}(\mathbf{B},\mathbf{C})$  can only be acquired after two steps of message passing: 1. atoms  $b$  and  $c$  aggregate message separately from  $\mathbf{B}$  and  $\mathbf{C}$ ; 2. the central atom  $a$  receives the aggregated message (which contains information of  $\mathbf{B}$  and  $\mathbf{C}$ ) from its neighbors  $b$  and  $c$ .
Figure 2: Illustrations of different local frames and their transition.
Let  $S_{\mathbf{B}}(S_{\mathbf{C}})$  denote the collection of all invariant scalars created by  $\mathbf{B}(\mathbf{C})$ . For example,  $S_{\mathbf{B}}$  contains all relative distances and angles within the 3D structure  $\mathbf{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) \neq \rho(S_{B}, S_{C})$  for arbitrary invariant functions  $\rho$ .
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  $\mathbf{B}(\mathbf{C})$  separately, through which all scalars in  $S_{\mathbf{B}}(S_{\mathbf{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.
Proof. The detailed proof is provided in Appendix D.
Frame Transition (FT). Let  $(e_1^i,e_2^i,e_3^i)$  and  $(e_1^j,e_2^j,e_3^j)$  be two orthonormal frames in  $\mathbf{R}^3$ . These frames are connected by an orthogonal matrix  $R_{ij}\in SO(3)$ :
$$
\left(\mathbf {e} _ {1} ^ {i}, \mathbf {e} _ {2} ^ {i}, \mathbf {e} _ {3} ^ {i}\right) = R _ {i j} \left(\mathbf {e} _ {1} ^ {j}, \mathbf {e} _ {2} ^ {j}, \mathbf {e} _ {3} ^ {j}\right). \tag {5}
$$
Furthermore, when  $(e_1^i,e_2^i,e_3^i)$  and  $(e_1^j,e_2^j,e_3^j)$  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  $\tau_{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  $\mathbf{m}_i$ . These tensor features allow us to maintain equivariant frames directly within each  $\mathbf{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  $\mathbf{m}_i$  from the  $l - 1$  layer to the  $l$  layer can be written as:  $\mathbf{m}_i^{l - 1} = \sum_{j\in N(i)}\mathbf{m}_j^{l - 1}$ . Since summation does not break the symmetry rule, it is obvious that  $\mathbf{m}_i^{l - 1}$  are still equivariant tensors. However, the nontrivial part lies in the design of the equivariant update function  $\phi$ :
$$
\mathbf {m} _ {i} ^ {l} = \phi \left(\mathbf {m} _ {i} ^ {l - 1}\right). \tag {6}
$$
To fully capture the information of FT, it is necessary for  $\phi$  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,  $\mathbf{m}(\mathbf{x}_u)$  is transformed equivariantly as:
$$
\mathbf {m} \left(g \mathbf {x} _ {u}\right) = \sum_ {i = 0} ^ {l} \mathcal {M} ^ {i} (g) \mathbf {m} _ {i} \left(\mathbf {x} _ {u}\right), g \in S E (3). \tag {7}
$$
Here,  $\mathbf{m}(\mathbf{x}_u)$  is decomposed to  $(\mathbf{m}_0(\mathbf{x}_u),\dots ,\mathbf{m}_l(\mathbf{x}_u))$  according to different tensor types, and  $\{\mathcal{M}^i (g)\}_{i = 0}^l$  is a collection of different  $SE(3)$  tensor representations (see the precise definition in Appendix A).