SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zian Li $^{1,2}$ , Xiyuan Wang $^{1,2}$ , Yinan Huang $^{1}$ , Muhan Zhang $^{1,*}$
Graph Neural Networks (GNNs) are often used for tasks involving the 3D geometry of a given graph, such as molecular dynamics simulation. While incorporating Euclidean distance into Message Passing Neural Networks (referred to as Vanilla DisGNN) is a straightforward way to learn the geometry, it has been demonstrated that Vanilla DisGNN is geometrically incomplete. In this work, we first construct families of novel and symmetric geometric graphs that Vanilla DisGNN cannot distinguish even when considering all-pair distances, which greatly expands the existing counterexample families. Our counterexamples show the inherent limitation of Vanilla DisGNN to capture symmetric geometric structures. We then propose  $k$ -DisGNNs, which can effectively exploit the rich geometry contained in the distance matrix. We demonstrate the high expressive power of  $k$ -DisGNNs from three perspectives: 1. They can learn high-order geometric information that cannot be captured by Vanilla DisGNN. 2. They can unify some existing well-designed geometric models. 3. They are universal function approximators from geometric graphs to scalars (when  $k \geq 2$ ) and vectors (when  $k \geq 3$ ). Most importantly, we establish a connection between geometric deep learning (GDL) and traditional graph representation learning (GRL), showing that those highly expressive GNN models originally designed for GRL can also be applied to GDL with impressive performance, and that existing complicated, equivariant models are not the only solution. Experiments verify our theory. Our  $k$ -DisGNNs achieve many new state-of-the-art results on MD17.
Many real-world tasks are relevant to learning the geometric structure of a given graph, such as molecular dynamics simulation, physical simulation, drug designing and protein structure prediction [Schmitz et al., 2019, Sanchez-Gonzalez et al., 2020, Jumper et al., 2021, Guo et al., 2020]. Usually in these tasks, the coordinates of nodes and their individual properties, such as atomic numbers, are given. The goal is to accurately predict both invariant properties, such as the energy of the molecule, and equivariant properties, such as the force acting on each atom. This kind of graphs is also referred to as geometric graphs by researchers [Bronstein et al., 2021, Han et al., 2022].
In recent years, Graph Neural Networks (GNNs) have achieved outstanding performance on such tasks, as they can learn a representation for each graph or node in an end-to-end fashion, rather than relying on handcrafted features. In such setting, a straightforward idea is to incorporate Euclidean distance, the most basic geometric feature, into Message Passing Neural Network [Gilmer et al., 2017] (we call these models Vanilla DisGNN) [Schütt et al., 2018, Kearnes et al., 2016]. While Vanilla DisGNN is simple yet powerful when considering all-pair distances (we assume all-pair
*Corresponding author: Muhan Zhang (muhan@pku.edu.cn).
distance is used by default to research a model's maximal expressive power), it has been proved to be geometrically incomplete [Zhang et al., 2021, Garg et al., 2020, Schütt et al., 2021, Pozdnyakov and Ceriotti, 2022], i.e., there exist pairs of geometric graphs which Vanilla DisGNN cannot distinguish. It has led to the development of various GNNs which go beyond simply using distance as input. Instead, these models use complex group irreducible representations, first-order equivariant representations or manually-designed complex invariant geometric features such as angles or dihedral angles to learn better representations from geometric graphs. It is generally believed that pure distance information is insufficient for complex GDL tasks [Klicpera et al., 2020a, Schütt et al., 2021].
On the other hand, it is well known that the distance matrix, which contains the distances between all pairs of nodes in a geometric graph, holds all of the geometric structure information [Satorras et al., 2021]. This suggests that it is possible to obtain all of the desired geometric structure information from distance graphs, i.e., complete weighted graphs with Euclidean distance as edge weight. Therefore, the complex GDL task with node coordinates as input (i.e., 3D graph) may be transformed into an equivalent graph representation learning task with distance graphs as input (i.e., 2D graph).
Thus, a desired question is: Can we design theoretically and experimentally powerful geometric models, which can learn the rich geometric patterns purely from distance matrix?
In this work, we first revisit the counterexamples in previous work used to show the incompleteness of Vanilla DisGNN. These existing counterexamples either consider only finite pairs of distance (a cutoff distance is used to remove long edges) [Schütt et al., 2021, Zhang et al., 2021, Garg et al., 2020], which can limit the representation power of Vanilla DisGNN, or lack diversity thus may not effectively reveal the inherent limitation of Vanilla DisGNN [Pozdnyakov and Ceriotti, 2022]. In this regard, we further constructed plenty of novel and symmetric counterexamples, as well as a novel method to construct families of counterexamples based on several basic units. These counterexamples significantly enrich the current set of counterexample families, and can reveal the inherent limitations of MPNNs in capturing symmetric configurations, which can explain the reason why they perform badly in real-world tasks where lots of symmetric (sub-)structures are included.
Given the limitations of Vanilla DisGNN, we propose  $k$ -DisGNNs, models that take pair-wise distance as input and aim for learning the rich geometric information contained in geometric graphs.  $k$ -DisGNNs are mainly based on the well-known  $k$ -(F)WL test [Cai et al., 1992], and include three versions:  $k$ -DisGNN,  $k$ -F-DisGNN and  $k$ -E-DisGNN. We demonstrate the superior geometric learning ability of  $k$ -DisGNNs from three perspectives. We first show that  $k$ -DisGNNs can capture arbitrary  $k$ - or  $(k + 1)$ -order geometric information (multi-node features), which cannot be achieved by Vanilla DisGNN. Then, we demonstrate the high generality of our framework by showing that  $k$ -DisGNNs can implement DimeNet [Klicpera et al., 2020a] and GemNet [Gasteiger et al., 2021], two well-designed state-of-the-art models. A key insight is that these two models are both augmented with manually-designed high-order geometric features, including angles (three-node features) and dihedral angles (four-node features), which correspond to the  $k$ -tuples in  $k$ -DisGNNs. However,  $k$ -DisGNNs can learn more than these handcrafted features in an end-to-end fashion. Finally, we demonstrate that  $k$ -DisGNNs can act as universal function approximators from geometric graphs to scalars (i.e., E(3) invariant properties) when  $k \geq 2$  and vectors (i.e., first-order O(3) equivariant and translation invariant properties) when  $k \geq 3$ . This essentially answers the question posed in our title: distance matrix is sufficient for GDL. We conduct experiments on benchmark datasets where our models achieve state-of-the-art results on a wide range of the targets in the MD17 dataset.
Our method reveals the high potential of the most basic geometric feature, distance. Highly expressive GNNs originally designed for traditional GRL can naturally leverage such information as edge weight, and achieve high theoretical and experimental performance in geometric settings. This opens up a new door for GDL research by transferring knowledge from traditional GRL, and suggests that existing complex, equivariant models may not be the only solution.
Equivariant neural networks. Symmetry is a rather important design principle for GDL [Bronstein et al., 2021]. In the context of geometric graphs, it is desirable for models to be equivariant or invariant under the group (such as E(3) and SO(3)) actions of the input graphs. These models include those using group irreducible representations [Thomas et al., 2018, Batzner et al., 2022, Anderson et al., 2019, Fuchs et al., 2020], complex invariant geometric features [Schütt et al., 2018, Klicpera
et al., 2020a, Gasteiger et al., 2021] and first-order equivariant representations [Satorras et al., 2021, Schutt et al., 2021, Tholke and De Fabritiis, 2021]. Particularly, Dym and Maron [2020] proved that both Thomas et al. [2018] and Fuchs et al. [2020] are universal for SO(3)-equivariant functions. Besides, Villar et al. [2021] showed that one could construct powerful (first-order) equivariant outputs by leveraging invariances, highlighting the potential of invariant models to learn equivariant targets.
GNNs with distance. In geometric settings, incorporating 3D distance between nodes as edge weight is a simple yet efficient way to improve geometric learning. Previous work [Maron et al., 2019, Morris et al., 2019, Zhang and Li, 2021, Zhao et al., 2022] mostly treat distance as an auxiliary edge feature for better experimental performance and do not explore the expressiveness or performance of using pure distance for geometric learning. Schutt et al. [2018] proposes to expand distance using a radial basis function as a continuous-filter and perform convolutions on the geometric graph, which can be essentially unified into the Vanilla DisGNN category. Zhang et al. [2021], Garg et al. [2020], Schutt et al. [2021], Pozdnyakov and Ceriotti [2022] demonstrated the limitation of Vanilla DisGNN by constructing pairs of non-congruent geometric graphs which cannot be distinguished by it, thus explaining the poor performance of such models [Schutt et al., 2018, Kearnes et al., 2016]. However, none of these studies proposed a complete purely distance-based geometric model. Recent work [Hordan et al., 2023] proposed theoretically complete GNNs for distinguishing geometric graphs, but they go beyond pair-wise distance and instead utilize gram matrices or coordinate projection.
Expressive GNNs. In traditional GRL (no 3D information), it has been proven that the expressiveness of MPNNs is limited by the Weisfeiler-Leman test [Weisfeiler and Leman, 1968, Xu et al., 2018, Morris et al., 2019], a classical algorithm for graph isomorphism test. While MPNNs can distinguish most graphs [Babai and Kucera, 1979], they are unable to count rings, triangles, or distinguish regular graphs, which are common in real-world data such as molecules [Huang et al., 2023]. To increase the expressiveness and design space of GNNs, Morris et al. [2019, 2020] proposed high-order GNNs based on  $k$ -WL. Maron et al. [2019] designed GNNs based on the folklore WL (FWL) test [Cai et al., 1992] and Azizian and Lelarge [2020] showed that these GNNs are the most powerful GNNs for a given tensor order. Beyond these, there are subgraph GNNs [Zhang and Li, 2021, Bevilacqua et al., 2021, Frasca et al., 2022, Zhao et al., 2021], substructure-based GNNs [Bouritsas et al., 2022, Horn et al., 2021, Bodnar et al., 2021] and so on, which are also strictly more expressive than MPNNs. For a comprehensive review on expressive GNNs, we refer the readers to Zhang et al. [2023].
More related work can be referred to Appendix E.
In this paper, we denote multiset with  $\{\{\}\}$ . We use  $[n]$  to represent the set  $\{1,2,\dots,n\}$ . A complete weighted graph with  $n$  nodes is denoted by  $G = (V,E)$ , where  $V = [n]$  and  $\pmb{E} = [e_{ij}]_{n\times n}\in \mathbb{R}^{n\times n}$ . The neighborhoods of node  $i$  are denoted by  $N(i)$ .
Distance graph vs. geometric graph. In many tasks, we need to deal with geometric graphs, where each node  $i$  is attached with its 3D coordinates  $\mathbf{x}_i\in \mathbb{R}^3$  in addition to other invariant features. Geometric graphs contain rich geometric information useful for learning chemical or physical properties. However, due to the significant variation of coordinates under E(3) transformations, they can be redundant when it comes to capturing geometric structure information.
Corresponding to geometric graphs are distance graphs, i.e., complete weighted graph with Euclidean distance as edge weight. Unlike geometric graphs, distance graphs do not have explicit coordinates attached to each node, but instead, they possess distance features that are naturally invariant under E(3) transformation and can be readily utilized by most GNNs originally designed for traditional GRL. Distance also provides an inherent inductive bias for effectively modeling the interaction/relationship between nodes. Moreover, a distance graph maintains all the essential geometric structure information, as stated in the following theorem:
Theorem 3.1. [Satorras et al., 2021] Two geometric graphs are congruent (i.e., they are equivalent by permutation of nodes and  $E(3)$  transformation of coordinates)  $\iff$  their corresponding distance graphs are isomorphic.
By additionally incorporating orientation information, we can learn equivariant features as well [Villar et al., 2021]. Hence, distance graphs provide a way to represent geometric graphs without referring to a canonical coordinate system. This urges us to study their full potential for GDL.
Weisfeiler-Leman Algorithms. Weisfeiler-Lehman test (also called as 1-WL) [Weisfeiler and Leman, 1968] is a well-known efficient algorithm for graph isomorphism test (traditional setting, no 3D information). It iteratively updates the labels of nodes according to nodes' own labels and their neighbors' labels, and compares the histograms of the labels to distinguish two graphs. Specifically, we use  $l_i^t$  to denote the label of node  $i$  at iteration  $t$ , then 1-WL updates the node label by
$$
l _ {i} ^ {t + 1} = \operatorname {H A S H} \left(l _ {i} ^ {t}, \left\{\left. l _ {j} ^ {t} \mid j \in N (i) \right\} \right\}\right), \tag {1}
$$
where HASH is an injective function that maps different inputs to different labels. However, 1-WL cannot distinguish all the graphs [Zhang and Li, 2021], and thus  $k$ -dimensional WL ( $k$ -WL) and  $k$ -dimensional Folklore WL ( $k$ -FWL),  $k \geq 2$ , are proposed to boost the expressiveness of WL test.
Instead of updating the label of nodes,  $k$ -WL and  $k$ -FWL update the label of  $k$ -tuples  $\mathbf{v} \coloneqq (v_{1}, v_{2}, \dots, v_{k}) \in V^{k}$ , denoted by  $l_{\mathbf{v}}$ . Both methods initialize  $k$ -tuples' labels according to their isomorphic types and update the labels iteratively according to their  $j$ -neighbors  $N_{j}(\mathbf{v})$ . The difference between  $k$ -WL and  $k$ -FWL mainly lies in the definition of  $N_{j}(\mathbf{v})$ . To be specific, tuple  $\mathbf{v}$ 's  $j$ -neighbors in  $k$ -WL and  $k$ -FWL are defined as follows respectively
$$
N _ {j} (\boldsymbol {v}) = \{\{(v _ {1}, \dots , v _ {j - 1}, w, v _ {j + 1}, \dots , v _ {k}) \mid w \in V \} \}, \tag {2}
$$
$$
N _ {j} ^ {F} (\boldsymbol {v}) = \big ((j, v _ {2}, \dots , v _ {k}), (v _ {1}, j, \dots , v _ {k}), \dots , (v _ {1}, v _ {2}, \dots , j) \big). \tag {3}
$$
To update the label of each tuple,  $k$ -WL and  $k$ -FWL iterates as follows
$$
k \text {- W L}: l _ {\boldsymbol {v}} ^ {t + 1} = \operatorname {H A S H} \left(l _ {\boldsymbol {v}} ^ {t}, \left(\left\{\left. l _ {\boldsymbol {w}} ^ {t} \mid \boldsymbol {w} \in N _ {j} (\boldsymbol {v}) \right\} \right\rvert   j \in [ k ]\right)\right), \tag {4}
$$
$$
k \text {- F W L}: l _ {\boldsymbol {v}} ^ {t + 1} = \operatorname {H A S H} \left(l _ {\boldsymbol {v}} ^ {t}, \{\{l _ {\boldsymbol {w}} ^ {t} \mid \boldsymbol {w} \in N _ {j} ^ {F} (\boldsymbol {v}) \} \mid j \in [ | V | ] \}\}\right). \tag {5}
$$
According to Cai et al. [1992], there always exists a pair of non-isomorphic graphs that cannot be distinguished by  $k$ -(F)WL but can be distinguished by  $(k + 1)$ -(F)WL. This means that  $k$ -(F)WL forms a strict hierarchy, which, however, still cannot solve the graph isomorphism problem with a finite  $k$ .
4 Revisiting the Incompleteness of Vanilla DisGNN
As mentioned in previous sections, Vanilla DisGNN is the edge-enhanced version of MPNN, which can unify many model frameworks [Schütt et al., 2018, Kearnes et al., 2016]. Its message passing formula can be generalized from its discrete version, which we call 1-WL-E [Pozdnyakov and Ceriotti, 2022]:
$$
l _ {i} ^ {t + 1} = \operatorname {H A S H} \left(l _ {i} ^ {t}, \{(l _ {j} ^ {t}, e _ {i j}) \mid j \in | V | \}\right). \tag {6}
$$
To provide valid counterexamples that Vanilla DisGNN cannot distinguish, Pozdnyakov and Ceriotti [2022] proposed both finite-size and periodic counterexamples and showed that they included real chemical structures. This work demonstrated for the first time the inherent limitations of Vanilla DisGNN. However, it should be noted that the essential change between these counterexamples is merely the distortion of  $C^{\pm}$  atoms, which lacks diversity despite having high manifold dimensions. In this regard, constructing new families of counterexamples cannot only enrich the diversity of existing families, but also demonstrate Vanilla DisGNN's limitations from different angles.
In this paper, we give a simple valid counterexample which Vanilla DisGNN cannot distinguish even with an infinite cutoff, as shown in Figure 1. In both geometric graphs, all nodes have exactly the same unordered list of distances (infinite cutoff considered), which means that Vanilla DisGNN will always label them identically. Nevertheless, the two geometric graphs are obviously non-congruent, since there are two small equilateral
Figure 1: A pair of geometric graphs that are non-congruent but cannot be distinguished by Vanilla DisGNN. Only red nodes belong to the two geometric graphs. The grey nodes and the "edges" are for visualization purpose only. Note that the nodes of the geometric graphs are sampled from regular icosahedrons.
triangles on the right, and zero on the left. Beyond this, we construct three kinds of counterexamples in Appendix A, which can significantly enrich the counterexamples found by Pozdnyakov and Ceriotti [2022], including:
1. Individual counterexamples sampled from regular polyhedrons that can be directly verified.
2. Small families of counterexamples that can be transformed in one dimension.
3. Families of counterexamples constructed by arbitrary combinations of basic symmetric units.
These counterexamples further demonstrate the incompleteness of Vanilla DisGNN in learning the geometry: even quite simple geometric graphs as shown in Figure 1 cannot be distinguished by it. While the counterexamples we constructed may not correspond to real molecules, such symmetric structures are commonly seen in other geometry-relevant tasks, such as physical simulations and point clouds. The inability to distinguish these symmetric structures can have a significant impact on the geometric learning performance of models, even for non-degenerated configurations, as demonstrated by Pozdnyakov and Ceriotti [2022]. This perspective can provide an additional explanation for the poor performance of models such as SchNet [Schütt et al., 2018] and inspire the design of more powerful and efficient geometric models.
5  $k$ -DisGNNs:  $k$ -order Distance Graph Neural Networks
In this section, we propose the framework of  $k$ -DisGNNs, complete and universal geometric models learning from pure distance features.  $k$ -DisGNNs consist of three versions:  $k$ -DisGNN,  $k$ -F-DisGNN, and  $k$ -E-DisGNN. Detailed implementation is referred to Appendix C.
Initialization Block. Given a geometric  $k$ -tuple  $\pmb{v}$ , high-order DisGNNs initialize it with an injective function by  $h_{\pmb{v}}^{0} = f_{\mathrm{init}}\Big((z_{v_{i}}\mid i\in [k]),(e_{v_{i}v_{j}}\mid i,j\in [k],i < j)\Big)\in \mathbb{R}^{K}$ , where  $K$  is the hidden dimension. This function can injectively embed the ordered distance matrix along with the  $z$  type of each node within the  $k$ -tuple, thus preserving all the geometric information within it.
Message Passing Block. The key difference between the three versions of high-order DisGNNs lies in their message passing blocks. The message passing blocks of  $k$ -DisGNN and  $k$ -F-DisGNN are based on the paradigms of  $k$ -WL and  $k$ -FWL, respectively. Their core message passing function  $f_{\mathrm{MP}}^t$  and  $f_{\mathrm{MP}}^{\mathrm{F},t}$  are formulated simply by replacing the discrete tuple labels  $l_{\mathbf{v}}^t$  in  $k$ -(F)WL (Equation (4, 5)) with continuous geometric tuple representation  $h_{\mathbf{v}}^t \in \mathbb{R}^K$ , see Equation (7, 8). Tuples containing local geometric information interact with each other in message passing blocks, thus allowing models to learn considerable global geometric information (as explained in Section 6.1).
$$
k - \operatorname {D i s G N N}: h _ {\boldsymbol {v}} ^ {t + 1} = f _ {\mathrm {M P}} ^ {t} \left(h _ {\boldsymbol {v}} ^ {t}, \left(\left\{\left\{h _ {\boldsymbol {w}} ^ {t} \mid \boldsymbol {w} \in N _ {j} (\boldsymbol {v}) \right\} \right\} \mid j \in | k |\right)\right), \tag {7}
$$
$$
k - \mathrm {F} - \operatorname {D i s G N N}: h _ {\boldsymbol {v}} ^ {t + 1} = f _ {\mathrm {M P}} ^ {\mathrm {F}, t} \left(h _ {\boldsymbol {v}} ^ {t}, \{\left\{\left(h _ {\boldsymbol {w}} ^ {t} \mid \boldsymbol {w} \in N _ {j} ^ {F} (\boldsymbol {v})\right) \mid j \in [ | V | ] \right\} \}\right). \tag {8}
$$
However, during message passing in  $k$ -DisGNN, the information about distance is not explicitly used but is embedded implicitly in the initialization block when each geometric tuple is embedded according to its distance matrix and node types. This means that  $k$ -DisGNN is unable to capture the relationship between a  $k$ -tuple  $v$  and its neighbor  $w$  during message passing, which could be very helpful for learning the geometric structural information of the graph. For example, in a physical system, a tuple  $w$  far from  $v$  should have much less influence on  $v$  than another tuple  $w'$  near  $v$ .
Based on this observation, we propose  $k$ -E-DisGNN, which maintains a representation  $e_{ij}^{t}$  for each edge  $e_{ij}$  at every time step and explicitly incorporates it into the message passing procedure. The edge representation  $e_{ij}^{t}$  is defined as
$$
e _ {i j} ^ {t} = f _ {\mathrm {e}} ^ {t} \left(e _ {i j}, \left(\left\{\left\{h _ {\boldsymbol {w}} ^ {t} \mid \boldsymbol {w} \in V ^ {k}, \boldsymbol {w} _ {u} = i, \boldsymbol {w} _ {v} = j \right\} \right\} \mid u, v \in [ k ], u <   v\right)\right). \tag {9}
$$
Note that  $e_{ij}^{t}$  not only contains the distance between nodes  $i$  and  $j$ , but also pools all the tuples related to edge  $ij$ , making it informative and general. For example, in the special case  $k = 2$ , Equation (9) is equivalent to  $e_{ij}^{t} = f_{\mathrm{e}}^{t}(e_{ij}, h_{ij}^{t})$ .
We realize the message passing function of  $k$ -E-DisGNN,  $f_{\mathrm{MP}}^{\mathrm{E},t}$ , by replacing the neighbor representation,  $h_w^t$  in  $f_{\mathrm{MP}}^t$ , with  $(h_w^t, e_{\mathbf{v} \setminus \mathbf{w}, \mathbf{w} \setminus \mathbf{v}}^t)$  where  $\mathbf{v} \setminus \mathbf{w}$  gives the only element in  $\mathbf{v}$  but not in  $\mathbf{w}$ , see Equation 10. In other words,  $e_{\mathbf{v} \setminus \mathbf{w}, \mathbf{w} \setminus \mathbf{v}}^t$  gives the representation of the edge connecting  $\mathbf{v}$ ,  $\mathbf{w}$ . By this means, a tuple can be aware of how far it is to its neighbors and by what kind of edges each neighbor
(A2)
Figure 2: A: High-order geometric information contained in the distance matrix of  $k$ -tuples. We mark different orders with different colors, with brown, green, blue for 2-,3-,4-order respectively. (A1) High-order geometric information contained in 4-tuples, including distances, angles and dihedral angles. (A2) More 3-order geometric features, such as vertical lines, middle lines and the area of triangles. B: Examples explaining that neighboring 3-tuples can form a 4-tuple. Blue represents the center tuple, the other colors represent neighbor tuples, and the red node is the one node of that neighbor tuple which is not in the center tuple. (B1) Example for 3-E-DisGNN. With the green edge, two 3-tuples can form a 4-tuple. (B2) Example for 3-F-DisGNN. Four 3-tuples form a 4-tuple.