Title: Is Distance Matrix Enough for Geometric Deep Learning?
Venue: NeurIPS 2023
Accepted: yes
Keywords: geometric deep learning

expressiveness

equivariant neural networks

universality
Abstract: 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.

Is Distance Matrix Enough for Geometric Deep Learning?

Abstract

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 ≥ 2 ) and vectors (when k ≥ 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.

1 Introduction

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

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 ≥ 2 and vectors (i.e., first-order O(3) equivariant and translation invariant properties) when k ≥ 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.

2 Related Work

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, Schütt et al., 2021, Thölke 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. Schütt 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], Schütt 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 [Schütt 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.

3 Preliminaries

In this paper, we denote multiset with { {} } . We use [ n ] to represent the set { 1 , 2 , ..., n } . A complete weighted graph with n nodes is denoted by G = ( V, E ) , where V = [ n ] and E = [ e ij ] n × n ∈ R n × 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 x i ∈ 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) ⇐⇒ 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 t i to denote the label of node i at iteration t , then 1-WL updates the node label by

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 ≥ 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 v := ( v 1 , v 2 , ..., v k ) ∈ V k , denoted by l v . Both methods initialize k -tuples' labels according to their isomorphic types and update the labels iteratively according to their j -neighbors N j ( v ) . The difference between k -WL and k -FWL mainly lies in the definition of N j ( v ) . To be specific, tuple v 's j -neighbors in k -WL and k -FWL are defined as follows respectively

To update the label of each tuple, k -WL and k -FWL iterates as follows

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]:

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 ± 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

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:

Individual counterexamples sampled from regular polyhedrons that can be directly verified.

Small families of counterexamples that can be transformed in one dimension.

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 v , high-order DisGNNs initialize it with an injective function by h 0 v = f init ( ( z v i | i ∈ [ k ] ) , ( e v i v j | i, j ∈ [ k ] , i < j ) ) ∈ 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 t MP and f F ,t MP are formulated simply by replacing the discrete tuple labels l t v in k -(F)WL (Equation (4, 5)) with continuous geometric tuple representation h t v ∈ 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).

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 t ij for each edge e ij at every time step and explicitly incorporates it into the message passing procedure. The edge representation e t ij is defined as

Note that e t ij 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 t ij = f t e ( e ij , h t ij ) .

We realize the message passing function of k -E-DisGNN, f E ,t MP , by replacing the neighbor representation, h t w in f t MP , with ( h t w , e t v \ w , w \ v ) where v \ w gives the only element in v but not in w , see Equation 10. In other words, e t v \ w , w \ v gives the representation of the edge connecting v , 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

(A1)

(A2)

(B1)

(B2)

is connected. This can boost the ability of geometric structure learning (as explained in Section 6.1).

Output Block. The output function t = f out ( { { h T v | v ∈ V k } } ) , where T is the final iteration, injectively pools all the tuple representations, and generates the E(3) and permutation invariant geometric target t ∈ R .

Like the conclusion about k -(F)WL for unweighted graphs, the expressiveness (in terms of approximating functions) of DisGNNs does not decrease as k increases. This is simply because all k -tuples are contained within some ( k +1) -tuples, and by designing message passing that ignores the last index, we can implement k -DisGNNs with ( k +1) -DisGNNs.

6 Rich Geometric Information Learned by k -DisGNNs

In this section, we aim to delve deeper into the geometry learning capability of k -DisGNNs from various perspectives. In Subsection 6.1, we will examine the high-order geometric features that the models can extract from geometric graphs. This analysis is more intuitive and aids in comprehending the high geometric expressiveness of k -DisGNNs. In Subsection 6.2, we will show that DimeNet and GemNet, two classical and widely-used GNNs for GDL employing invariant geometric representations, are special cases of k -DisGNNs, highlighting the generality of k -DisGNNs. In the final subsection, we will show the completeness (distinguishing geometric graphs) and universality (approximating functions over geometric graphs) of k -DisGNNs, thus answering the question posed in the title: distance matrix is enough for GDL.

6.1 Ability of Learning High-Order Geometry

We first give the concept of high-order geometric information for better understanding.

Definition 6.1. k -order geometric information is the E(3)-invariant features calculable from k nodes' 3D coordinates.

For example, in a 4-tuple, one can find various high-order geometric information, including distance (2-order), angles (3-order) and dihedral angles (4-order), as shown in Figure 2(A1). Note that highorder information is not limited to the common features listed above: we show some of other possible 3-order geometric features in Figure 2(A2).

We first note that k -DisGNNs can learn and process at least k -order geometric information. Theorem 3.1 states that we can reconstruct the whole geometry of the k -tuple from the embedding of its distance matrix. In other words, we can extract all the desired k -order geometric information solely from the embedding of the k -tuple's distance matrix, which is calculated at the initialization step of k -DisGNNs, with a learnable function.

Furthermore, both k -E-DisGNN and k -F-DisGNN can actually learn ( k +1) -order geometric information in their message passing layers. For k -E-DisGNN, information about ( h v , h w , e v \ w , w \ v ) , where w is some j -neighbor of v , is included in the input of its update function. Since the distance matrices of tuples v and w can be reconstructed from h v and h w , the all-pair distances of ( v 1 , v 2 , ..., v k , w \ v ) can be reconstructed from ( h v , h w , e v \ w , w \ v ) , as shown in Figure 2(B1). Similarly, the distance matrix of ( v 1 , v 2 , ..., v k , j ) can also be reconstructed from ( h v , ( h t w | w ∈ N F j ( v )) ) in update function of k -F-DisGNN as shown in Figure 2(B2). These enable k -E-DisGNN and k -F-DisGNN to reconstruct all ( k +1) -tuples' distance matrices during message passing and thus learn all the ( k +1) -order geometric information contained in the graph.

With the ability to learn high-order geometric information, k -DisGNNs can learn geometric structures that cannot be captured by Vanilla DisGNNs. For example, consider the counterexample shown in Figure 1. The two geometric graphs have a different number of small equilateral triangles, which cannot be distinguished by Vanilla DisGNNs even with infinite message passing layers. However, 3 -DisGNN and 2 -E/F-DisGNN can easily distinguish the graphs by counting the number of small equilateral triangles, which is actually a kind of 3-order geometric information. This example also illustrates that high-order DisGNNs are strictly more powerful than Vanilla DisGNNs.

6.2 Unifying Existing Geometric Models with DisGNNs

There have been many attempts to improve GDL models by manually designing and incorporating high-order geometric features such as angles (3-order) and dihedral angles (4-order). These features are all invariant geometric features that can be learned by some k -DisGNNs. It is therefore natural to ask whether these models can be implemented by k -DisGNNs. In this subsection, we show that two classical models, DimeNet [Klicpera et al., 2020a] and GemNet [Gasteiger et al., 2021], are just special cases of k -DisGNNs, thus unifying existing methods based on hand-crafted features with a learning paradigm that learns arbitrary high-order features from distance matrix.

DimeNet embeds atom i with a set of incoming messages m ji , i.e., h i = ∑ j ∈ N i m ji , and updates the message m ji by

where d ji is the distance between node j and node i , and ϕ kji is the angle kji . We simplify the subscript of ∑ , same for that in Equation (12). The detailed ones are referred to Appendix B.2, B.1.

DimeNet is one early approach that uses geometric features to improve the geometry learning ability of GNNs, especially useful for learning the angle-relevant energy or other chemical properties. Note that the angle information that DimeNet explicitly incorporates into its message passing function is actually a kind of 3-order geometric information, which can be exploited from distance matrices by our 2-E/F-DisGNN with a learning paradigm. This gives the key insight for the following proposition.

Proposition 6.2. 2-E-DisGNN and 2-F-DisGNN can implement DimeNet.

GemNet is also a graph neural network designed to process graphs embedded in Euclidean space. Unlike DimeNet, GemNet incorporates both angle and dihedral angle information into the message passing functions, with a 2-step message passing scheme. The core procedure is as follows:

where θ cabd is the dihedral angle of planes cab and abd . The use of dihedral angle information allows GemNet to learn at least 4-order geometric information, making it much more powerful than models that only consider angle information. We now prove that GemNet is just a special case of 3-E/F-DisGNN, which can also learn 4-order geometric information during its message passing stage.

Proposition 6.3. 3-E-DisGNN and 3-F-DisGNN can implement GemNet.

It is worth noting that while both k -DisGNNs and existing models can learn some k -order geometric information, k -DisGNNs have the advantage of learning arbitrary k -order geometric information . This means that we can learn different high-order geometric features according to specific tasks in a data-driven manner, including but not limited to angles and dihedral angles.

6.3 Completeness and Universality of k -DisGNNs

We have provided an illustrative example in Section 6.1 that demonstrates the ability of k -DisGNNs to distinguish certain geometric graphs, which is not achievable by Vanilla DisGNNs. Actually, all counterexamples presented in Appendix A can be distinguished by 3-DisGNNs or 2-E/F-DisGNNs. This leads to a natural question: Can k -DisGNNs distinguish all geometric graphs with a finite and small value of k ? Furthermore, can they approximate arbitrary functions over geometric graphs?

Wenow present our main findings that elaborate on the completeness and universality of k -DisGNNs with finite and small k , which essentially highlight the high theoretical expressiveness of k -DisGNNs.

Theorem 6.4. (informal) Let M ( θ ) ∈ MinMods = { 1-round 4-DisGNN, 2-round 3-E/F-DisGNN } with parameters θ . Denote h m = f node ( { { h T v | v 0 = m } } ) ∈ R K ′ as node m 's representation produced by M where f node is an injective multiset function, and x c m as node m 's coordinates w.r.t. the center. Denote MLPs : R K ′ → R as a multi-layer perceptron. Then we have:

( Completeness ) There exists θ 0 such that M ( θ 0 ) can distinguish all pairs of non-congruent geometric graphs.

( Universality for scalars ) M ( θ ) is a universal function approximator for continuous, E (3) and permutation invariant functions f : R 3 × n → R over geometric graphs.

( Universality for vectors ) f equiv out = ∑ | V | m =1 MLPs ( h m ) x c m is a universal function approximator for continuous, O (3) equivariant and translation and permutation invariant functions f : R 3 × n → R 3 over geometric graphs.

In fact, completeness directly implies universality for scalars [Hordan et al., 2023]. Since we can actually recover the whole geometry from arbitrary node representation h m , universality for vectors can be proved with conclusions from Villar et al. [2021]. Also note that as order k and round number go higher, the completeness and universality still hold. Formal theorems and detailed proof are provided in Appendix B.3.

There is a concurrent work by Delle Rose et al. [2023], which also investigates the completeness of k -WL on distance graphs. In contrast to our approach of using only one tuple's final representation to reconstruct the geometry, they leverage all tuples' final representations. They prove that with the distance matrix, 3-round geometric 2-FWL and 1-round geometric 3-FWL are complete for geometry. This conclusion can also be extended to k -DisGNNs, leading to the following theorem:

Theorem 6.5. The completeness and universality for scalars stated in Theorem 6.4 hold for M ( θ ) ∈ MinMods ext = { 3-round 2-E/F-DisGNN, 1-round 3-E/F-DisGNN } with parameters θ .

This essentially lowers the required order k for achieving universality on scalars to 2. However, due to the lower order and fewer rounds, the expressiveness of a single node's representation may be limited, thus the universality for vectors stated in Theorems 6.4 may not be guaranteed.

7 Experiments

In this section, we evaluate the experimental performance of k -DisGNNs. Our main objectives are to answer the following questions:

Q1 Does 2/3-E/F-DisGNN outperform their counterparts in experiments (corresponding to Section 6.2)?

Q2 As universal models for scalars (when k ≥ 2 ), do k -DisGNNs also have good experimental performance (corresponding to Section 6.3)?

Q3 Does incorporating well-designed edge representations in the k -E-DisGNN result in improved performance?

The best and the second best results are shown in bold and underline respectively in tables. Detailed experiment configuration and supplementary experiment information can be found in Appendix D. Our code is available at https://github.com/GraphPKU/DisGNN .

MD17. MD17 [Chmiela et al., 2017] is a dataset commonly used to evaluate the performance of machine learning models in the field of molecular dynamics. It contains a collection of molecular

dynamics simulations of small organic molecules such as aspirin and ethanol. Given the atomic numbers and coordinates, the task is to predict the energy of the molecule and the atomic forces. We mainly focus on the comparison between 2-F-DisGNN/3-E-DisGNN and DimeNet/GemNet. At the same time, we also compare k -DisGNNs with the state-of-the-art models: FCHL [Christensen et al., 2020], PaiNN [Schütt et al., 2021], NequIP [Batzner et al., 2022], TorchMD [Thölke and De Fabritiis, 2021], GNN-LF [Wang and Zhang, 2022]. The results are shown in Table 1. 2-F-DisGNN and 3-E-DisGNN outperform their counterparts on 16/16 and 6/8 targets, respectively, with an average improvement of 43.9% and 12.23%, suggesting that data-driven models can subsume carefully designed manual features given high expressiveness. In addition, k -DisGNNs also achieve the best performance on 8/16 targets, and achieve the second-best performance on all the other targets , outperforming the best baselines GNN-LF and TorchMD by 10.19% and 12.32%, respectively. Note that GNN-LF and TorchMD are all complex equivariant models. Beyond these, we perform experiments on revised MD17 , which has better data quality, and k -DisGNNs outperforms the SOTA models [Batatia et al., 2022, Musaelian et al., 2023] on a wide range of targets. The results are referred to Appendix D.2. The results firmly answer Q1 and Q2 posed at the beginning of this section and demonstrate the potential of pure distance-based methods for graph deep learning.

QM9. QM9 [Ramakrishnan et al., 2014, Wu et al., 2018] consists of 134k stable small organic molecules with 19 regression targets. The task is like that in MD17, but this time we want to predict the molecule's properties such as the dipole moment. We mainly compare 2-F-DisGNN with DimeNet on this dataset and the results are shown in Table 2. 2-F-DisGNN outperforms DimeNet on 11/12 targets, by 14.27% on average, especially on targets µ (65.0%) and ⟨ R 2 ⟩ (90.4%), which also answers Q1 well. A full comparison to other state-of-the-art models is included in Appendix D.2.

Effectiveness of edge representations. To answer Q3 , we split the edge representation e t ij (Equation (9)) into two parts, namely the pure distance (the first element) and the tuple representations (the second element), and explore whether incorporating the two elements is beneficial. We conduct experiments on MD17 with three versions of 2-DisGNNs, including 2-DisGNN (no edge representation), 2-e-DisGNN (add only edge weight) and 2-E-DisGNN (add full edge representation). The results are shown in Table 3. Both 2-e-DisGNN and 2-E-DisGNN exhibit significant performance improvements over 2-DisGNN, highlighting the significance of edge representation. Furthermore, 2-E-DisGNN outperforms 2-DisGNN and 2-e-DisGNN on 16/16 and 12/16 targets, respectively, with average improvements of 39.8% and 3.8%. This verifies our theory that capturing the full edge representation connecting two tuples can boost the model's representation power.

8 Conclusions and Limitations

Conclusions. In this work we have thoroughly studied the ability of GNNs to learn the geometry of a graph solely from its distance matrix. We expand on the families of counterexamples that Vanilla DisGNNs are unable to distinguish from their distance matrices by constructing families of symmetric and diverse geometric graphs, revealing the inherent limitation of Vanilla DisGNNs in capturing symmetric configurations. To better leverage the geometric structure information contained in distance graphs, we proposed k -DisGNNs, geometric models with high generality (ability to unify two classical and widely-used models, DimeNet and GemNet), provable completeness (ability to distinguish all pairs of non-congruent geometric graphs) and universality (ability to universally approximate scalar functions when k ≥ 2 and vector functions when k ≥ 3 ). In experiments, k -DisGNNs outperformed previous state-of-the-art models on a wide range of targets of the MD17 dataset. Our work reveals the potential of using expressive GNN models, which were originally designed for traditional GRL, for the GDL tasks, and opens up new opportunities for this domain.

Limitations. Although DisGNNs can achieve universality (for scalars) with a small order of 2, the results are still generated under the assumption that the distance graph is complete. In practice, for large geometric graphs such as proteins, this assumption may lead to intractable computation due to the O ( n 3 ) time complexity and O ( n 2 ) space complexity. However, we can still balance theoretical universality and experimental tractability by using either an appropriate cutoff, which can furthest preserve the completeness of the distance matrix in local clusters while cutting long dependencies to improve efficiency (and generalization), or using sparse but expressive GNNs such as subgraph GNNs [Zhang and Li, 2021]. We leave it for future work.

References

Gunnar Schmitz, Ian Heide Godtliebsen, and Ove Christiansen. Machine learning for potential energy surfaces: An extensive database and assessment of methods. The Journal of chemical physics , 150 (24):244113, 2019.

Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia. Learning to simulate complex physics with graph networks. In International Conference on Machine Learning , pages 8459-8468. PMLR, 2020.

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature , 596(7873):583-589, 2021.

Yulan Guo, Hanyun Wang, Qingyong Hu, Hao Liu, Li Liu, and Mohammed Bennamoun. Deep learning for 3d point clouds: A survey. IEEE transactions on pattern analysis and machine intelligence , 43(12):4338-4364, 2020.

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.

Han, Y. Rong, T. Xu, and W. Huang. Geometrically equivariant graph neural networks: A survey. 2022.

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.

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.

Steven Kearnes, Kevin McCloskey, Marc Berndl, Vijay Pande, and Patrick Riley. Molecular graph convolutions: moving beyond fingerprints. Journal of computer-aided molecular design , 30: 595-608, 2016.

Yaolong Zhang, Junfan Xia, and Bin Jiang. Physically motivated recursively embedded atom neural networks: incorporating local completeness and nonlocality. Physical Review Letters , 127(15): 156002, 2021.

Vikas Garg, Stefanie Jegelka, and Tommi Jaakkola. Generalization and representational limits of graph neural networks. In International Conference on Machine Learning , pages 3419-3430. PMLR, 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.

Sergey N Pozdnyakov and Michele Ceriotti. Incompleteness of graph convolutional neural networks for points clouds in three dimensions. arXiv preprint arXiv:2201.07136 , 2022.

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

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.

Jin-Yi Cai, Martin Fürer, and Neil Immerman. An optimal lower bound on the number of variables for graph identification. Combinatorica , 12(4):389-410, 1992.

Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 6790-6802. Curran Associates, Inc., 2021. URL https://proceedings.neurips.cc/ paper/2021/file/35cf8659cfcb13224cbd47863a34fc58-Paper.pdf .

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.

Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications , 13(1):1-11, 2022.

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

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

Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations , 2021.

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

Soledad Villar, David W Hogg, Kate Storey-Fisher, Weichi Yao, and Ben Blum-Smith. Scalars are universal: Equivariant machine learning, structured like classical physics. Advances in Neural Information Processing Systems , 34:28848-28863, 2021.

Haggai Maron, Heli Ben-Hamu, Hadar Serviansky, and Yaron Lipman. Provably powerful graph networks. Advances in neural information processing systems , 32, 2019.

Christopher Morris, Martin Ritzert, Matthias Fey, William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. In Proceedings of the AAAI conference on artificial intelligence , volume 33, pages 4602-4609, 2019.

Muhan Zhang and Pan Li. Nested graph neural networks. Advances in Neural Information Processing Systems , 34:15734-15747, 2021.

Lingxiao Zhao, Louis Härtel, Neil Shah, and Leman Akoglu. A practical, progressively-expressive gnn. arXiv preprint arXiv:2210.09521 , 2022.

Snir Hordan, Tal Amir, Steven J Gortler, and Nadav Dym. Complete neural networks for euclidean graphs. arXiv preprint arXiv:2301.13821 , 2023.

Boris Weisfeiler and Andrei Leman. The reduction of a graph to canonical form and the algebra which appears therein. NTI, Series , 2(9):12-16, 1968.

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

László Babai and Ludik Kucera. Canonical labelling of graphs in linear average time. In 20th Annual Symposium on Foundations of Computer Science (sfcs 1979) , pages 39-46. IEEE, 1979.

Yinan Huang, Xingang Peng, Jianzhu Ma, and Muhan Zhang. Boosting the cycle counting power of graph neural networks with i 2 -gnns. 2023.

Christopher Morris, Gaurav Rattan, and Petra Mutzel. Weisfeiler and leman go sparse: Towards scalable higher-order graph embeddings. Advances in Neural Information Processing Systems , 33: 21824-21840, 2020.

Waiss Azizian and Marc Lelarge. Expressive power of invariant and equivariant graph neural networks. arXiv preprint arXiv:2006.15646 , 2020.

Beatrice Bevilacqua, Fabrizio Frasca, Derek Lim, Balasubramaniam Srinivasan, Chen Cai, Gopinath Balamurugan, Michael M Bronstein, and Haggai Maron. Equivariant subgraph aggregation networks. arXiv preprint arXiv:2110.02910 , 2021.

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.

Lingxiao Zhao, Wei Jin, Leman Akoglu, and Neil Shah. From stars to subgraphs: Uplifting any gnn with local structure awareness. arXiv preprint arXiv:2110.03753 , 2021.

Giorgos Bouritsas, Fabrizio Frasca, Stefanos Zafeiriou, and Michael M Bronstein. Improving graph neural network expressivity via subgraph isomorphism counting. IEEE Transactions on Pattern Analysis and Machine Intelligence , 45(1):657-668, 2022.

Max Horn, Edward De Brouwer, Michael Moor, Yves Moreau, Bastian Rieck, and Karsten Borgwardt. Topological graph neural networks. arXiv preprint arXiv:2102.07835 , 2021.

Cristian Bodnar, Fabrizio Frasca, Yuguang Wang, Nina Otter, Guido F Montufar, Pietro Lio, and Michael Bronstein. Weisfeiler and lehman go topological: Message passing simplicial networks. In International Conference on Machine Learning , pages 1026-1037. PMLR, 2021.

Bohang Zhang, Shengjie Luo, Liwei Wang, and Di He. Rethinking the expressive power of gnns via graph biconnectivity. arXiv preprint arXiv:2301.09505 , 2023.

Valentino Delle Rose, Alexander Kozachinskiy, Cristóbal Rojas, Mircea Petrache, and Pablo Barceló. Three iterations of ( d -1) -wl test distinguish non isometric clouds of d -dimensional points. arXiv e-prints , pages arXiv-2303, 2023.

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, Lars A Bratholm, Felix A Faber, and O Anatole von Lilienfeld. Fchl revisited: Faster and more accurate quantum machine learning. The Journal of chemical physics , 152(4): 044107, 2020.

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

Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gábor Csányi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems , 35:11423-11436, 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.

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.

Zhenqin Wu, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learning. Chemical science , 9(2):513-530, 2018.

Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. arXiv preprint arXiv:1812.09902 , 2018.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

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.

Johannes Klicpera, Shankari Giri, Johannes T Margraf, and Stephan Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. arXiv preprint arXiv:2011.14115 , 2020b.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. Mlp-mixer: An all-mlp architecture for vision. Advances in neural information processing systems , 34:2426124272, 2021.

Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision , pages 10012-10022, 2021.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

Jun Xia, Chengshuai Zhao, Bozhen Hu, Zhangyang Gao, Cheng Tan, Yue Liu, Siyuan Li, and Stan Z Li. Mole-bert: Rethinking pre-training graph neural networks for molecules. 2023.

Shuqi Lu, Zhifeng Gao, Di He, Linfeng Zhang, and Guolin Ke. Highly accurate quantum chemical property prediction with uni-mol+. arXiv preprint arXiv:2303.16982 , 2023.

Christopher Morris, Gaurav Rattan, Sandra Kiefer, and Siamak Ravanbakhsh. Speqnets: Sparsityaware permutation-equivariant graph networks. arXiv preprint arXiv:2203.13913 , 2022.

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.

Albert P Bartók, Risi Kondor, and Gábor Csányi. On representing chemical environments. Physical Review B , 87(18):184115, 2013.

Alexander V Shapeev. Moment tensor potentials: A class of systematically improvable interatomic potentials. Multiscale Modeling & Simulation , 14(3):1153-1173, 2016.

Jörg Behler and Michele Parrinello. Generalized neural-network representation of high-dimensional potential-energy surfaces. Physical review letters , 98(14):146401, 2007.

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

Chaitanya K Joshi, Cristian Bodnar, Simon V Mathis, Taco Cohen, and Pietro Liò. On the expressive power of geometric graph neural networks. arXiv preprint arXiv:2301.09308 , 2023.

A Counterexamples for Vanilla DisGNNs

In this section, we will give some counterexamples, i.e., pairs of geometric graphs that are not congruent but cannot be distinguished by Vanilla DisGNNs. We will organize the section as follows: First, we will give some isolated counterexamples that can be directly verified; Then, we will give some counterexamples that can hold the property (i.e., to be counterexamples) after some simple continuous transformation; Finally, we will give a family of counterexamples that can be obtained by the combination of some basic units. Since the latter two cases cannot be directly verified, we will also give the detailed proof.

We say there are M kinds of nodes if there are M different labels after infinite iterations of Vanilla DisGNNs. Note that in this section, all the grey nodes and the 'edges' are just for visualization purpose . Only colored nodes belong to the geometric graph.

Note that we provide verification programs in our code for all the counterexamples. The first type of counterexamples (Appendix A.1) is limited in quantity and can be directly verified. The remaining two types (Appendix A.2 and Appendix A.3) are families of counterexamples and have an infinite number of cases. Our code can verify them to some extent by randomly selecting some parameters. Theoretical proofs for these two types are also provided in Appendix A.2 and Appendix A.3, respectively.

A.1 Counterexamples That Can Be Directly Verified

Since the distance graph, which DisGNNs take as input, is a kind of complete graph and contains lots of geometric constraints, if we want to get a pair of non-congruent geometric graphs that cannot be distinguished by DisGNNs, they must exhibit great symmetry and have just minor difference . Inspired by this, we can try to sample nodes from regular polyhedrons, which themselves exhibit high symmetry. The first pair of geometric graphs is sampled from regular icosahedrons, which is referred to the main body of our paper, see Figure 1. Since all the nodes from both graphs have the same unordered distance list, there is only 1 kind of nodes for both graphs. We note that it is the only counterexample we can sample from just regular icosahedron.

The following pairs of geometric graphs are all sampled from regular dodecahedrons. We will distinguish them by the graph size.

6-size and 14-size counterexamples.

A pair of 6-size geometric graphs (red nodes) and a pair of 14-size geometric graphs (blue nodes) that are counterexamples are shown in the figure above. They are actually complementary in terms of forming all vertices of a regular dodecahedron.

For the 6-size counterexample, the label distribution after infinite rounds of Vanilla DisGNN is shown in the picture above, and the nodes are colored differently to represent different labels. It can be directly verified that there are 2 kinds of nodes. Note that the two geometric graphs are non-congruent: The isosceles triangle of the two geometric graphs formed by one green point and two pink nodes have the same leg length but different base lengths.

For the 14-size counterexample, we give the conclusion that there are 4 kinds of nodes and the number of nodes of each kind is 2, 4, 4, 4 respectively. Since the counterexamples here and in the following are complex and not so intuitionistic, we recommend readers to directly verify them with our programs.

8-size and 12-size counterexamples.

A pair of 8-size geometric graphs (red nodes) and a pair of 12-size geometric graphs (blue nodes) that are counterexamples are shown in the figure above.

For the 8-size counterexample, note that it is actually obtained by carefully inserting two nodes into the 6-size counterexample mentioned just before. We note that there are 2 kinds of nodes in this counterexample, and the numbers of nodes of each kind are 4, 4 respectively.

For the 12-size counterexample, it is actually obtained by carefully deleting two nodes from the 14-size counterexample (corresponds to the way to obtain the 8-size counterexample) mentioned just before. There are 4 kinds of nodes in this counterexample, and the numbers of nodes of each kind are 4, 4, 2, 2 respectively.

Two pairs of 10-size counterexamples.

The above figure shows the first pair of 10-size counterexamples. There are 3 kinds of nodes and the numbers of nodes of each kind are 4, 4, 2 respectively.

The above figure shows the second pair of 10-size counterexamples. There is actually only 1 kind of nodes, all the nodes of both graphs have the same unordered distance list. It is interesting because

there are two regular pentagons in the left graph and a ring in the right graph, which means that the two graphs are non-congruent, and it is quite similar to the case shown in Figure 1, where there are two equilateral triangles in the left graph and also a ring in the right graph.

In the end, we note that the above counterexamples are probably all the counterexamples one can sample from a single regular polyhedron.

A.2 Counterexamples That Hold After Some Continuous Transformation

Inspired by the intuition mentioned in Appendix A.1, we can also sample nodes from the combination of regular polyhedrons, which also have great symmetry but with more nodes.

Cube + regular octahedron.

scale cube

scale cube

The first two counterexamples are sampled from the combination of a cube and a regular octahedron. We combine them in the following way: Bring the center of both polyhedrons together, and align the vertex of the regular octahedron, the center of the cube's face, and the center of the polyhedrons in a straight line. We do not limit the relative size of the two polyhedrons: one can scale the cube to any size as long as the constraints are not broken. In this way, a small family of counterexamples is given (we say small here because the dimension of the transformation space is small).

Proof of the red case. Let the side length of the cube be 2 a , the diagonal of the regular octahedron be 2 b . Initially, all nodes are labeled l 0 . We now simulate Vanilla DisGNNs on both graphs.

At the first iteration, there are actually two kinds of unordered distance lists for both of the graphs: For the vertexes on the regular octahedron, the list is { ( (2 b, l 0 ) , 1 ) , ( ( √ 3 a 2 + b 2 +2 ab, l 0 ) , 2 ) , ( ( √ 3 a 2 + b 2 -2 ab, l 0 ) , 2 ) } . For the vertexes on the cube, the list is { ( (2 a, l 0 ) , 1 ) , ( (2 √ 2 a, l 0 ) , 1 ) , ( (2 √ 3 a, l 0 ) , 1 ) , ( ( √ 3 a 2 + b 2 +2 ab, l 0 ) , 1 ) , ( ( √ 3 a 2 + b 2 -2 ab, l 0 ) , 1 ) } . Thus they will be labeled differently at this step, let the labels be l 1 and l 2 respectively.

At the second iteration, the lists of all the vertexes on the regular octahedron from both graphs are still the same, i.e. { ( (2 b, l 1 ) , 1 ) , ( ( √ 3 a 2 + b 2 +2 ab, l 2 ) , 2 ) , ( ( √ 3 a 2 + b 2 -2 ab, l 2 ) , 2 ) } , as well as the lists of all the vertexes on the cube from both graphs, i.e. { ( (2 a, l 2 ) , 1 ) , ( (2 √ 2 a, l 2 ) , 1 ) , ( (2 √ 3 a, l 2 ) , 1 ) , ( ( √ 3 a 2 + b 2 +2 ab, l 1 ) , 1 ) , ( ( √ 3 a 2 + b 2 -2 ab, l 1 ) , 1 ) } .

Since Vanilla DisGNNs cannot subdivide the vertexes at the second step, they can still not at the following steps. So it cannot distinguish the two geometric graphs. But they are non-congruent: on the left graph, the nodes on the regular octahedron can only form isosceles triangles with the nodes on the face diagonal of the cube. On the right graph, they can only form isosceles triangles with nodes on the cubes that are on the same side. We note that the counterexample in Figure 1 is not a special case of this, because on the left graph of this case all the nodes are on the same surface, but none of graphs in Figure 1 have all nodes on the same surface.

Proof of the blue case. The proof of the blue case is quite similar to that of the red case; thus, we omit the proof here. Note that the labels will stabilize after the first iteration. At the end of Vanilla DisGNN, there are 2 kinds of nodes, each kind has 4 nodes. And the two geometric graphs are

non-congruent because the plane formed by four nodes on the regular octahedron is perpendicular to the plane formed by four nodes on the cube in the left graph, and is not in the right graph.

2 cubes.

scale cube

scale cube

The second case is the combination of 2 cubes. We combine them in the following way: Let the center of the two cubes coincide as well as the axis formed by the centers of two opposite faces. Also, we do not limit the relative size of the two cubes: one can scale the cube to any size as long as the constraints are not broken. The proof of this case is quite similar to that in Cube + regular octahedron ; thus, we omit the proof here. Note that the labels will stabilize after the first iteration. If the sizes of the two cubes are the same, then at the end of Vanilla DisGNN, there is only 1 kind of node. If not, there are 2 kinds of nodes, each kind has 4 nodes.

A.3 Families of Counterexamples

In this section, we will prove that any one of the counterexamples given in Appendix A.1 can be augmented to a family. To facilitate the discussion, we will introduce some notations.

Notation. Consider an arbitrary counterexample ( G L ori , G R ori ) given in Appendix A.1. It is worth noting that for each graph, the nodes are sampled from a regular polyhedron and are located on the same spherical surface. We can denote the geometric graph on this sphere as G ori,r , where r represents the radius of the sphere. Since the nodes are sampled from some regular polyhedron (denoted by G all,r ), there must be nodes that aren't sampled, which we call as complementary nodes. We denote the corresponding geometric graph as G com,r . Given two graphs G Q 1 ,r 1 , G Q 2 ,r 2 where Q i ∈ { ori, com, all } for i ∈ [2] , we say that we align them if we make sure that the spherical centers of the two geometric graphs coincide and one can get G Q 2 ,r 2 by just scaling all the nodes of G Q 2 ,r 1 along the spherical center. We use L = { ( v, l v ) | v ∈ V } to represent a label state for some graph G , and call it a stable state if Vanilla DisGNN cannot further refine the labels of the graph G with initial label L . We say that L is the final state of G if it is the label state obtained by performing Vanilla DisGNN on G with no initial labels for infinite rounds. Note that final state is a special stable state. We denote the final state of G ori and G com as L ori and L com respectively.

̸

Now, let us introduce an augmentation method for arbitrary counterexample ( G L ori , G R ori ) shown in Appendix A.1. For a given set QR k = { ( Q 1 , r 1 ) , ( Q 2 , r 2 ) , ..., ( Q k , r k ) } where r i > 0 , Q i ∈ { ori, com, all } and r i = r j for i = j , we get the pair of augmented geometric graphs ( G L , G R ) = AUG QR k ( G L ori , G R ori ) by the following steps (take G L for example):

̸

Generate k geometric graphs G L Q 1 ,r 1 , i ∈ [ k ] according to QR k .

Align the k geometric graphs.

Combine the k geometric graphs G L Q i ,r i , i ∈ [ k ] to obtain G L .

We give an example of the augmentation in Figure 3. Now we give our theorem.

Theorem A.1. Consider an arbitrary counterexample ( G L ori , G R ori ) given in Appendix A.1. Let QR k = { ( Q 1 , r 1 ) , ( Q 2 , r 2 ) , . . . , ( Q k , r k ) } be an arbitrary set, where Q i ∈ { ori, com, all } and

̸

r i = r j for i = j . If ∃ i ∈ [ k ] such that Q i = all , then the augmented geometric graphs ( G L , G R ) = AUG QR k ( G L ori , G R ori ) form a counterexample.

̸

̸

Before presenting the proof of the theorem, we will introduce several lemmas that will be used in the proof.

Lemma A.2. If Vanilla DisGNN cannot distinguish ( G L , G R ) with some initial label ( L L , L R ) , then Vanilla DisGNN cannot distinguish ( G L , G R ) without initial label (with identical labels).

̸

̸

Proof of Lemma A.2. We use l t u to represent the label of u after t iterations of Vanilla DisGNN without initial labels, and use l L ,t u to represent the label of u after t iterations of Vanilla DisGNN with initial label L . We use l ∞ u and l L , ∞ u to represent the corresponding final label. In order to prove this lemma, we first prove a conclusion: for arbitrary graphs G = ( V, E ) and arbitrary initial label L , let u, v be two arbitrary nodes in V , then we have ∀ t ∈ Z + , l t u = l t v → l L ,t u = l L ,t v . We'll prove the conclusion by induction.

· t = 1 holds.

̸

If ( l 0 u , { { ( d us , l 0 s ) | s ∈ N ( u ) } } ) = ( l 0 v , { { ( d vs , l 0 s ) | s ∈ N ( v ) } } ) , then we have { { d us | s ∈ N ( u ) } } = { { d vs | s ∈ N ( v ) } } , since l 0 s are all identical for s ∈ V . Then it is obvious that ( l L , 0 u , { { ( d us , l L , 0 s ) | s ∈ N ( u ) } } ) = ( l L , 0 v , { { ( d vs , l L , 0 s ) | s ∈ N ( v ) } } ) .

̸

For ∀ k ∈ Z + , t = k holds ⇒ t = k +1 holds.

Assume t = k +1 does not hold, i.e., ∃ u, v ∈ V , l ( k +1) u = l ( k +1) v but l L , ( k +1) u = l L , ( k +1) v . We can derive from l L , ( k +1) u = l L , ( k +1) v that

̸

Since t = k holds, for ∀ s 1 , s 2 , we have

̸

Together with Equation (13), we have l k u = l k v and { { ( d us , l k s ) | s ∈ N ( u ) } } = { { ( d vs , l k s ) | s ∈ N ( v ) } } . This means l ( k +1) u = l ( k +1) v , which is contradictory to assumptions.

This means that the conclusion holds for ∀ t ∈ Z + , as well as t = ∞ . In other words, specifying an initial label L will only result in the subdivision of the final labels.

Back to Lemma A.2, if Vanilla DisGNN cannot distinguish ( G L , G R ) even when the final labels are furthermore subdivided, it can still not when the final labels are not subdivided, i.e., cannot distinguish ( G L , G R ) without initial labels.

Lemma A.3. Consider an arbitrary counterexample ( G L ori , G R ori ) given in Appendix A.1. We have that ( L L ori ∪ L L com , L R ori ∪ L R com ) is a pair of stable states of ( G L all , G R all ) . Denote the labels as L L all and L R all respectively. Moreover, we have that Vanilla DisGNN cannot distinguish ( G L all , G R all ) with initial labels L L all and L R all .

Proof of Lemma A.3. Since the number of situations is finite, this lemma can be directly verified using a computer program. It is worth noting that we have also included the verification program in our code.

This lemma allows us to draw several useful prior conclusions. In the graph G all , the node set V all can be split into two subsets by definition, namely V ori and V com . Now for arbitrary node u from G L all with initial label L L all and v from G R all with initial label L R all , if the initial labels of node u and node v are the same, they must be in the same subset, i.e., respectively in V L ori and V R ori or V L com and V R com . Without loss of generality, let us assume that both u and v are in V ori . Since ( L L all , L R all ) is a stable state of ( G L all , G R all ) , we can obtain the following equation:

Since ( L L ori , L R ori ) is a stable state of the induced subgraphs ( G L ori , G R ori ) where the node sets are ( V L ori , V R ori ) , we can obtain the following equation:

Notice that in Equation (14, 15), we can obtain a new equation by replacing V L ori or V L all with ( V L ori -u ) or ( V L all -u ) and doing the same for the v side, while keeping the equation valid, since l u = l v . By subtracting Equation (15) from Equation (14), we obtain the following equation:

It is important to note that the same conclusion holds if both u and v are from type com . These equations provide valuable prior knowledge that we can use to prove Theorem A.1.

Now, let us prove Theorem A.1. Our proof is divided into four steps:

Construct an initial state ( L L , L R ) for the augmented geometric graphs ( G L , G R ) .

Prove that ( L L , L R ) is a pair of stable states for ( G L , G R ) .

Explain that Vanilla DisGNN cannot distinguish ( G L , G R ) with initial labels ( L L , L R ) .

Explain that ( G L , G R ) are non-congruent.

By applying Lemma A.2, we can get the conclusion from step 2 and 3 that Vanilla DisGNN cannot distinguish ( G L , G R ) without initial labels. Since ( G L , G R ) are non-congruent (as established in step 4), it follows that ( G L , G R ) is a counterexample.

Proof of Theorem A.1.

Step 1. We first construct an initial state for the augmented graphs ( G L , G R ) .

For simplicity, we omit the superscripts, as the rules are the same for both graphs. For each layer (i.e., the sphere of some radius) of the graph G , we label the nodes with L tp if the layer is of type tp , for tp ∈ { ori, com, all } . Note that we also ensure that the labels in different layers are distinct.

Step 2. Now we prove that ( L L , L R ) is a stable state for ( G L , G R ) . This is equivalent to proving that for arbitrary node u in ( G L , L L ) and node v in ( G R , L R ) , if l 0 u = l 0 v , then l 1 u = l 1 v .

Since G L is considered as a complete distance graph by DisGNNs, the neighbors of node u are all the nodes in G L except itself. We denote the neighbor nodes from the layer with radius r i by N r i ( u ) .

This is similar for node v . Now we need to prove that ( l 0 u , { { ( d us , l 0 s ) | s ∈ ⋃ i ∈ [ k ] N r i ( u ) } } ) = ( l 0 v , { { ( d vs , l 0 s ) | s ∈ ⋃ i ∈ [ k ] N r i ( v ) } } ) . Since l 0 u = l 0 v , we only need to prove that { { ( d us , l 0 s ) | s ∈ ⋃ i ∈ [ k ] N r i ( u ) } } = { { ( d vs , l 0 s ) | s ∈ ⋃ i ∈ [ k ] N r i ( v ) } } .

We split the multiset { { ( d us , l 0 s ) | s ∈ ⋃ i ∈ [ k ] N r i ( u ) } } into k multisets, namely { { ( d us , l 0 s ) | s ∈ N r i ( u ) } } , i ∈ [ k ] , and do the same for the multiset of node v . Our goal is to prove that { { ( d us , l 0 s ) | s ∈ N r i ( u ) } } = { { ( d vs , l 0 s ) | s ∈ N r i ( v ) } } for each i ∈ [ k ] .

Let the coordinates of node s in a spherical coordinate system be ( r s , θ s , ϕ s ) . Since nodes u and v have the same initial label, they must be from the same layer, meaning that r u = r v . Additionally, we always realign the coordinate systems of G L and G R to ensure that the direction of the polyhedra is the same. The distance between node u and node s in a spherical coordinate system is given by √ r 2 u + r 2 s -2 r u r s ( sin θ u sin θ s cos( ϕ u -ϕ s ) + cos θ u cos θ s ) . We denote the angle term ( sin θ u sin θ s cos( ϕ u -ϕ s ) + cos θ u cos θ s ) by Θ us for simplicty.

For each value of r i , it can be one of three types: ori , com , or all . Similarly, nodes u and v can belong to one of three categories: ori , com , all . Regardless of the combination, these situations can always be found in the Equations (14, 15, 16) derived from Lemma A.3 and can produce the following conclusion:

Then we have:

since r u = r i . This means that for all values of r i , { { ( d us , l 0 s ) | s ∈ N r i ( u ) } } = { { ( d vs , l 0 s ) | s ∈ N r i ( v ) } } . By merging all the multisets of radius r i , we can get { { ( d us , l 0 s ) | s ∈ ⋃ i ∈ [ k ] N r i ( u ) } } = { { ( d vs , l 0 s ) | s ∈ ⋃ i ∈ [ k ] N r i ( v ) } } , which concludes the proof.

Step 3. Since the stable states ( L L , L R ) are obtained by assigning the stable states of each layer, i.e. L com , L ori or L all , that cannot be distinguished by Vanilla DisGNN, the histogram of both graphs are exactly the same, which means that Vanilla DisGNN cannot distinguish the two graphs.

Step 4. Notice that the construction of an isomorphic mapping requires that the radius of the layer of node u in G L and node v in G R must be the same. If there exists a pair of layers in the two geometric graphs that are non-congruent, then the two geometric graphs are non-congruent. According to the definition of ( G L , G R ) , such layers always exist, therefore the two geometric graphs are non-congruent.

B Proofs

B.1 Proof of Proposition 6.3

The main proof of our proposition is to construct the key procedures in GemNet [Gasteiger et al., 2021] using basic message passing layers of 3-E/F-DisGNN, and the whole model can be constructed by stacking these message passing layers up. Since the k -E-DisGNN and k -F-DisGNN share the same initialization and output blocks and have similar update procedure (they can both learn 4-order geometric information), we mainly focus on the proof of E version, and one can check easily that the derivation for F version also holds.

Basic methods. Assume that we want to learn O tar = f tar ( I tar ) with our function O fit = f fit ( I fit ) . In our proof, the form of O tar and O fit , as well as the form of I tar and I fit , are quite different. For example, consider the case where O fit is an embedding h abc for a 3-tuple abc and I fit = I ( abc ) fit contains the information of all the neighbors of abc , while O tar is an embedding m ab for a 2-tuple ab and I tar = I ( ab ) tar contains the information of all the neighbors of ab . Therefore, directly learning functions by f fit that produces exactly the same output as f tar is inappropriate. Instead, we will learn functions that can calculate several different outputs in a way f tar does and appropriately embeds them into the output of f fit . For example, we still consider the case mentioned before, and we want to learn a function f fit that can extract I ( ab ) tar , I ( ac ) tar , I ( bc ) tar , I ( ba ) tar , I ( ca ) tar , I ( cb ) tar from I ( abc ) fit respectively, and calculates m ab , m ac , m bc , m ba , m ca , m cb with these information like f tar , and embed them into the output h abc in an injective way.

Since we realize f fit as a universal function approximator (such as MLPs and deep multisets), f fit can always learn the function we want. So the only concern is that whether we can extract exact I tar from I fit , i.e., whether there exists an injective function ( I (1) tar , I (2) tar , ..., I ( n ) tar ) = f ext ( I fit ) . We will mainly discuss about this in our proof.

Notations. We use the superscript G to represent functions in GemNet and 3E to represent functions in 3-E-DisGNN, and use I with the same superscript and subscript to represent the input of a function. We use the superscript ( a 1 a 2 ..a k ) to represent some input I if it 's for tuple ( a 1 a 2 ..a k ) . If there exists an injective function I tar = f ext ( I fit ) , then we denote it by I fit →I tar , meaning that I tar can be derived from I fit . If some geometric information I geo (such as distance and angles) is contained in the distance matrix of a tuple a 1 a 2 ..a k , then we denote it by I geo ∈ a 1 a 2 ..a k . For simplicity, We omit all the time superscript t if the context is clear.

B.1.1 Construction of Embedding Block

Initialization of directional embeddings. GemNet initialize all the two-tuples m (also called directional embeddings) at the embedding block by the following paradigm

At the initial step, 3-E-DisGNN follows the paradigm outlined below:

Then we have

meaning that we can extract all the information to calculate m ab , m ac , m bc , m ba , m ca , m cb from the input of f 3E init . And thanks to the universal property of f 3E init , we can approximate a function that accurately calculates these variables and injectively embeds them into h abc , such that h abc → ( m ab , m ac , m bc , m ba , m ca , m cb ) .

Initialization of atom embeddings. GemNet initializes all the one-tuples u i (also called atom embeddings) at the embedding block simply by passing atomic number z i through an embedding layer.

Note that since we can learn all the m ij and embed them into h abc by f 3E init , it is also possible to learn all the u i , i ∈ { a, b, c } and embed them into h abc at the same time with some function, since I 3E , ( abc ) init contains z i , i ∈ { a, b, c } . Therefore, h abc → ( u a , u b , u c ) holds.

Initialization of geometric information. It is an important observation that since f 3E init take all the pair-wise distance within 3-tuple abc as input, all the geometric information within the tuple can be included in h abc . This makes geometric information rich in the 3-tuple embedding h abc .

̸

To remind the readers, now we prove that there exists a function f 3E init which can correctly calculate m ij ( i = j, i, j ∈ { a, b, c } ) , u i ( i ∈ { a, b, c } ) and all the geometric information I geo within the triangle abc , and injectively embed them into h abc :

̸

B.1.2 Construction of Atom Embedding Block

Atom embedding block. GemNet updates the 1-tuple embeddings u by summing up all the relevant 2-tuple embeddings m in the atom emb block. The process can be formulated as

Since this function involves the process of pooling, it cannot be learned by f 3E init . However, it can be easily learned by the basic message passing layers of 3-E-DisGNN f 3E MP , which is formulated as

We now want to learn a function f 3E MP that updates the u a , u b , u c embedded in h abc like f G atom and keep the other variables and information unchanged. Note that h abc is the first input of f 3E MP , so all the old information is maintained. As what we talked about earlier, the main focus is to check whether the information to update u is contained in f 3E MP 's input. In fact, the following derivation holds

Note that d ka ∈ abc , and e ( ka ) RBF can be calculated from d ka . Similarly, we can derive that I 3E , ( abc ) MP → I G , ( b ) atom and I 3E , ( abc ) MP →I G , ( c ) atom . This means we can update u in h using a basic message passing layer of 3-E-DisGNN.

B.1.3 Construction of Interaction Block

Message passing. There are two key procedures in GemNet's message passing block, namely two-hop geometric message passing (Q-MP) and one-hop geometric message passing (T-MP), which can be abstracted as follows

̸

̸

̸

̸

Note that what we need to construct is a function f 3E MP that can update the information about m ij , i, j ∈ { a, b, c } , i = j embedded in h abc just like what f G TMP and f G QMP do, and keep the other variables and information unchanged. Since it is quite similar among different m ij , we will just take the update process of m ab for example.

First, T-MP. For this procedure, the following derivation holds

̸

Note that in the derivation above, there is an important conclusion implicitly used: the tuple ( h abc , h kbc , e ak ) actually contains all the geometric information in the 4-tuple abck , because the

distance matrix of the four nodes can be obtained from it. Thus d kb , d ab , ϕ abk can be obtained from it, and since e ( kb ) RBF and e ( abk ) CBF are just calculated from these geometric variables, the derivation holds.

And we can exclude the element in the multiset where the index k = a , simply because these tuples have different patterns from others.

Second, Q-MP. In Q-MP, the pooling objects consist of two indices (which we call two-order pooling), namely k 1 and k 2 in Equation (23). One alternative way to do this is two-step pooling, i.e. pool two times and once for one index. For example we can pool index k 1 before we pool index k 2 as follows:

̸

̸

Note that the expressiveness of two-step pooling is not less than two-order pooling, i.e.

̸

̸

̸

̸

̸

Thus, if we use two-step pooling to implement Q-MP, it does not reduce expressiveness. Inspired by this, in order to update m ab in h abc like f G QMP , we first learn a function by f 3E MP that calculates an intermediate variable w c = f 3E inter1 ( { { ( m ck , e ( ck ) RBF , e ( bkc ) CBF , e ( abkc ) SBF ) | k ∈ [ N ] , k = a } } ) by pooling all the m ck at index k , which is feasible because the following derivation holds

̸

Note that in the derivation process above, m ck is directly derived from e ck because it contains the information by definition 21. Then we apply another message passing layer f 3E MP but this time we learn a function that just pools all the w c in h abc and finally updates m ab :

̸

This means we can realize f G QMP ′ by stacking two message passing layers up.

Atom self-interaction. This sub-block actually involves two procedures: First, update atom embeddings u according to the updated directional embeddings m . Second, update the directional embeddings m according to the updated atom embeddings u . The first step is actually an atom embedding block, which is already realized by our f 3E MP in Appendix B.1.1. The second procedure can be formulated as

It is obvious that we can update m ab in h abc by f 3E MP because the following derivation holds

B.1.4 Construction of Output Block

In GemNet, the final output t is obtained by summing up all the sub-outputs from each interaction block. While it is possible to add additional sub-output blocks to our model, in our proof we only consider the case where the output is obtained solely from the final interaction block.

The following function produces the output of GemNet

where W out is a learnable matrix.

And our output function is

We can realize Equation (26) by stacking a message passing layer and an output block. First, the message passing layer updates all the atom embeddings u in h abc , and then the output block extracts u from h as follows, and calculates t like GemNet.

B.2 Proof of Proposition 6.2

In this section, we will follow the basic method and notations in Appendix B.1. We will mainly prove for 2-F-DisGNN, and since when k = 2 , k -E-DisGNN's update function can implement k -F-DisGNN's (See Appendix B.4), the derivations also holds for 2-E-DisGNN. We will use the superscript D to represent functions in DimeNet and 2F to represent functions in 2-F-DisGNN.

B.2.1 Construction of Embedding Block

DimeNet initializes all the two tuples in the embedding block just like GemNet, which can be formulated as

What 2-F-DisGNN does at the initialization step is

Now assume we want to learn such a function that can learn both m ab and m ba , then embed it into h ab . As what we talked about in Appendix B.1, it is possible because the following derivation holds

Note that h ab also contains the geometric information, but in this case, just the distance. Different from GemNet, DimeNet does not "track" the 1-tuple embeddings: it does not initialize and update the atom embeddings in the model, and only pool the 2-tuples into 1-tuples in the output block. So there is no need to embed the embeddings of atom a and atom b into h ab .

B.2.2 Construction of Interaction Block

The message passing framework of DimeNet (so-called directional message passing) can be formulated as

̸

And the message passing framework of 2-F-DisGNN can be formulated as

̸

Now we want to learn a function f 2F MP that can updates the m ab and m ba embedded in h ab like what f D MP does. We need to check if f 2F MP have sufficient information to update m ab and m ba embedded in its output h ab . The derivation holds:

̸

̸

Note that we again used the observation that tuple ( h ab , h kb , h ak ) contains all the geometric information of 3-tuple abk . Similarly we can get I 2F , ( ab ) MP →I D , ( ba ) MP .

B.2.3 Construction of Output Block

The output block of DimeNet is quite similar to that of GemNet (Appendix B.1.4), which can be formulated as

This can be realized by stacking a message passing layer of 2-F-DisGNN and its output block up.

The message passing layer of 2-F-DisGNN can learn u a and u b and embed it into the output h ab because the following derivation holds

Similarly we can get I 2F , ( ab ) MP →I D , ( b ) atom .

And the output block of 2-F-DisGNN is formulated as follows

Note that the following derivation holds

This means that the output block of 2-F-DisGNN can implement the sum operation in Equation (33).

B.3 Proof of Theorem 6.4

We first restate the theorem formally.

Theorem B.1. Let M ( θ ) ∈ MinMods = { 1-round 4-DisGNN, 2-round 3-E-DisGNN, 2-round 3-F-DisGNN } with parameters θ . Denote h m = f node ( { { h T v | v ∈ V k , v 0 = m } } ) ∈ R K ′ as node m 's representation produced by M where f node is an injective multiset function, and x c m as node m 's coordinates w.r.t. the center. Denote MLPs( θ ′ ) : R K ′ → R as a multi-layer perceptron with parameters θ ′ .

( Completeness ) Given arbitrary two geometric graphs X 1 , X 2 ∈ R 3 × n , then there exists θ 0 such that M ( X 1 ; θ 0 ) = M ( X 2 ; θ 0 ) ⇐⇒ X 1 and X 1 are congruent.

( Universality for scalars ) Let f : R 3 × n → R be an arbitrary continuous, E (3) invariant and permutation invariant function over geometric graphs, then for any compact set M ⊂ R 3 × n and ϵ > 0 , there exists θ 0 such that ∀ X ∈ M, | f ( X ) -M ( X ; θ 0 ) | ≤ ϵ .

( Universality for vectors ) Let f : R 3 × n → R 3 be an arbitrary continuous, O (3) equivariant, permutation and translation invariant function over geometric graphs, let f equiv out = ∑ | V | m =1 MLPs ( h m ) x c m , then for any compact set M ⊂ R 3 × n and ϵ > 0 , there exists θ 0 for M and θ ′ 0 for MLPs such that ∀ X ∈ M, | f ( X ) -f equiv out ( X ; θ 0 , θ ′ 0 ) | ≤ ϵ .

In the forthcoming proof for completeness part, we examine the optimal scenario for k -(E/F)DisGNNs, namely the geometric version k -(E/F)WL, where the hyper-parameters such as hidden dimensions are set and parameters θ of M ( θ ) are learned to ensure that all the relevant functions are hash functions. To simplify the analysis for k -EWL, we substitute the edge representation in k -E-DisGNN with the edge weight (distance) alone, which is already enough.

Basic idea for proof of completeness. We prove the completeness part of the theorem, which demonstrates that the optimal DisGNNs have the ability to distinguish all geometric graphs, through a reconstruction method. Specifically, if we can obtain all points' coordinates up to permutation and E(3) translation from the output of DisGNNs, it implies that DisGNNs can differentiate between all geometric graphs.

B.3.1 Proof of Completeness for One-Round 4-WL

Proof. We begin by assuming the existence of four affinely independent nodes, meaning that there exists a 4-tuple of nodes that can form a tetrahedron. If this condition is not met, and all the nodes lie on the same plane, the problem is simplified, as the cloud degenerates into 2D space. We will address this case later in the discussion. We assume the point number is N .

Our first step is to prove that we can recover the 4-tuple formed by the four affinely independent nodes from the output of one-round 4-WL, denoted by { { c 1 i | i ∈ [ N ] 4 } } . Since all the tuple colors are initialized injectively according to the tuple's distance matrices, there exists a function f 0 such that f 0 ( c 0 i ) = D i , where D i is the distance matrix of tuple i . Note that since the update procedures of the colors are HASH functions, c t i = HASH( c t -1 i , C t -1 0 , i , C t -1 1 , i , C t -1 2 , i , C t -1 3 , i ) where C t -1 m, i denotes the color multiset of tuple i 's m -th neighbors, there also exists a series functions f t such that f t ( c t i ) = c t -1 i . This simply means that the latter colors of a tuple will contain all the information of its former colors. Given the fact, we can use the function f 1 ◦ f 0 to reconstruct the distance matrices of all the tuples, and find the one that represents a tetrahedron geometry (Theorem 3.1). We mark the found tuple with k .

Now we prove that one can reconstruct the whole geometry from just c 1 k .

First, we can reconstruct the 4 points' 3D coordinates (given any arbitrary center and orientation) from k 's distance matrix (Theorem 3.1). Formally, there exists a function f D ( D k ) = X k , where X k ∈ R 4 ∗ 3 represents the coordinates. And since f 1 ◦ f 0 ( c 1 k ) = D k , we have that f 1 ◦ f 0 ◦ f D ( c 1 k ) = X k .

The update function of c k is c 1 k = HASH( c 0 k , C 0 0 , k , C 0 1 , k , C 0 2 , k , C 0 3 , k ) , where C 0 m, k = { { c 0 Φ m ( k ,j ) | j ∈ [ N ] } } , m ∈ [4] and Φ m ( k , j ) replaces the m -th element in tuple k with j . Since the function is HASH function, we can also reconstruct each of C 0 m, k from c 1 k . In C 0 m, k , there exists N 4-tuples' initial color, from which we can reconstruct the 4-tuple's distance matrix. Note that given a color c 0 Φ m ( k ,j ) in the multiset, we can reconstruct 3 distances, namely d ( x j , x k ( m +1)%4 ) , d ( x j , x k ( m +2)%4 ) , d ( x j , x k ( m +3)%4 ) , where x j representes the 3D coordinates of point j and d ( x, y ) calculates the l 2 norm of x, y .

There is a strong geometric constrain in 3D space: Given the distances between a point and 3 affinely independent points (whose postions are known), one can calculate at most 2 possible positions of the point, and the 2 possible postions are mirror-symmetric relative to the plane formed by the 3 affinely independent points. Moreover, the point is on the plane formed by the 3 affinely independent points iff there is only one solution of the distance equations.

Now since we have already reconstructed X k , and k ( m +1)%4 , k ( m +2)%4 , k ( m +3)%4 are affinely independent, we can calculate 2 possible positions of point j . And for the whole multiset C 0 m, k , we can calculate a N -size multiset where each element is a pair of possible positions of some point, denoted as P m, k = { {{ x (1) j , x (2) j } | j ∈ [ N ] } } . Note that now the possible geometry of the N points form a 0-dimension manifold, with no more than 2 N elements.

There are two important properties of P m, k :

Given arbitrary two elements (pairs) of P m, k , denoted as p j 1 and p j 2 , we have p j 1 = p j 2 or p j 1 ∩ p j 2 = ∅ . These correspond to two possible situations: j 1 and j 2 are either mirror-symmetric relative to the plane formed by k ( m +1)%4 , k ( m +2)%4 , k ( m +3)%4 , or not.

Arbitrary element p j of P m, k has multiplicity value at most 2. This corresponds two possible situations: j either has a mirror-symmetric point relative to the plane formed by k ( m +1)%4 , k ( m +2)%4 , k ( m +3)%4 , or does not.

The follows proves that the four 0-dimension manifolds determined by P m, k , m ∈ [4] can intersect at a unique solution, which is the real geometry.

We first 'refine' the multiset P m, k . There are 3 kinds of pairs in P m, k :

The pair { x (1) j , x (2) j } that have multiplicity value 2. This means that there are two mirrorsymmetric points relative to the plane formed by k ( m +1)%4 , k ( m +2)%4 , k ( m +3)%4 . So we can ensure that the two points with coordinates x (1) j , x (2) j both exist.

The pair where x (1) j = x (2) j = x j . This means that by solving the 3 distance equations, there is only one solution. So we can uniquely determine the coordinates of point j (and the point is on the plane).

The other pairs.

We can determine the real coordinates of the points from the first two kinds of pairs, so we record the real coordinates and delete them from the multiset P m, k (for the first kind, delete both two pairs). We denote the real coordinates determined by P m, k as A m . Now we have 4 preliminarily refined sets, denoted as P ′ m, k . Note that now in each P ′ m, k :

Arbitrary two elements intersect at ∅ .

Arbitrary element has two distinct possible coordinates, one is real and the other is fake.

The two properties above also mean that elements in multiset ⋃ p ∈ P ′ m, k p has only multiplicity 1.

Since the real coordinates A m determined by different P m, k may be not the same, we further refine the P ′ m, k by removing the pairs that contain coordinates in ( A 0 ∪ A 1 ∪ A 2 ∪ A 3 ) -A m . Note that according to the property of P ′ m, k mentioned above, each coordinates in ( A 0 ∪ A 1 ∪ A 2 ∪ A 3 ) -A m will and will only corresponds to one pair in P ′ m, k . Denote the further refined sets as P ′′ m, k . After the refinement, each P ′′ m, k has equally M = N -| A 0 ∪ A 1 ∪ A 2 ∪ A 3 | elements (pairs), corresponding to M undetermined real coordinates.

̸

We use p ∩ P to denote the multiset { { p ′ | p ′ ∈ P, p ′ ∩ p = ∅} } . Now we prove that ∃ m ∈ [4] , ∃ p ∈ P ′′ m, k , ∃ m ′ ∈ [4] -m, s.t. | p ∩ P ′′ m ′ , k | = 1 by contradiction.

Note that ∀ m ∈ [4] , ∀ p ∈ P ′′ m, k , ∀ m ′ ∈ [4] -m, | p ∩ P ′′ m ′ , k | ∈ { 1 , 2 } . That's because each element in P ′′ m, k contains one real coordinates and one fake coordinates, each real/fake coordinates can at most intersect with one pair in P ′′ m ′ , k , and there must exist one element in P ′′ m ′ , k that contains the same real coordinates.

So opposite to the conclusion is: ∀ m ∈ [4] , ∀ p ∈ P ′′ m, k , ∀ m ′ ∈ [4] -m, | p ∩ P ′′ m ′ , k | = 2 , we assume this holds. Now there are also two properties:

For all m ∈ [4] , elements in multiset ⋃ p ∈ P ′′ m, k p has multiplicity 1.

For all m ∈ [4] , ⋃ p ∈ P ′′ m, k p with different m shares the same M real coordinates.

So we have that: ⋃ p ∈ P ′′ m, k p is the same across all m ∈ [4] . Note that for each p ∈ P ′′ m, k , since the two coordinates in p are mirror-symmetric relative to plane k ( m +1)%4 , k ( m +2)%4 , k ( m +3)%4 , the midpoint of the two points is on the plane. This means that by averaging all the coordinates in ⋃ p ∈ P ′′ m, k p , we have a center point on the plane k ( m +1)%4 , k ( m +2)%4 , k ( m +3)%4 . However, since ⋃ p ∈ P ′′ m, k p for all m ∈ [4] are the same, we have one same center point on each four planes k ( m +1)%4 , k ( m +2)%4 , k ( m +3)%4 , m ∈ [4] . This contradicts to the assumption that k 1 , k 2 , k 3 , k 4 are affinely independent points. So the conclusion that ∃ m ∈ [4] , ∃ p ∈ P ′′ m, k , ∃ m ′ ∈ [4] -m, s.t. | p ∩ P ′′ m ′ , k | = 1 holds.

Since the conclusion holds, we can obtain the p and the unique element p ′ ∈ p ∩ P ′′ m ′ , k . Moreover, we have that | p ′ ∩ p | = 1 . This is because, the two pairs share one real coordinates, and the fake coordinates must be different (The two planes formed by k ( m +1)%4 , k ( m +2)%4 , k ( m +3)%4 and k ( m ′ +1)%4 , k ( m ′ +2)%4 , k ( m ′ +3)%4 are not the same). So we obtain the real coordinates by p ′ ∩ p .

Now we further refine the 4 P ′′ m, k by deleting the pair that contains the real coordinates we just found, which makes each P ′′ m, k have only M -1 pairs now. Then we loop the procedure above to find all the remaining real coordinates.

Note that all the information, including the final geometry, we obtained in the above proof is derived from only c 1 k . This means that when generating c 1 k , 4-WL injectively embeds the whole geometry using a HASH function. For two non-congruent geometric graph, their 4-tuple's (where 4 nodes are affinely independent) colors will always be different, giving the 4-WL the power to distinguish them.

For the degenerated situation (there does not exist 4 affinely independent nodes), the points form a 2D plane. Then with similar proof idea, we can prove that 3-WL can be universal on 2D plane. Since 3-WL is not stronger than 4-WL, 4-WL can also distinguish all the degenerated cases.

B.3.2 Proof of Completeness for Two-Round 3-FWL

Proof. As previous, our discussion is under the fact that the point cloud does not degenerates to a 2D or 1D point cloud. Otherwise, the problem is quite trivial.

We first find the tuple containing three nodes that are affinely independent from the 2-round 3-FWL's output { { c 2 i | i ∈ [ N ] 3 } } . Denote the founded tuple as k . We first calculate the 3 nodes' coordinates according to the c 0 k derived from c 2 k .

As the update procedure, c 2 k = HASH( c 1 k , { { ( c 1 Φ 0 ( k ,j ) , c 1 Φ 1 ( k ,j ) , c 1 Φ 2 ( k ,j ) ) | j ∈ N } } ) , each tuple ( c 1 Φ 0 ( k ,j ) , c 1 Φ 1 ( k ,j ) , c 1 Φ 2 ( k ,j ) ) in the multiset contains the distance d k 0 ,j , d k 1 ,j , d k 2 ,j . Thus, we can calculate at most two possible coordinates of node j , and the two coordinates are mirror-symmetric relative to plane k . Like in the proof of 4-WL, we denote the final calculaoed possible coordinates-pair multiset as P k = { {{ x (1) j , x (2) j } | j ∈ [ N ] } } .

We then further find a node j 0 that is not on plane k from the multiset term in c 2 k , and its colors' tuple is ( c 1 Φ 0 ( k ,j 0 ) , c 1 Φ 1 ( k ,j 0 ) , c 1 Φ 2 ( k ,j 0 ) ) . Note that the three colors are at time-step 1, which means that they already aggregated neighbors' information for one-round. So by repeating the above procedure, we can again get a possible coordinates-pair multiset P Φ m ( k ,j 0 ) from each of c 1 Φ m ( k ,j 0 ) where m ∈ [3] .

Note that the four plane, k , Φ 0 ( k , j 0 ) , Φ 1 ( k , j 0 ) , Φ 2 ( k , j 0 ) , do not intersect at a common point. So we can do as what we do in the proof of 4-WL, to repeatedly refine each multiset P and get all the real coordinates, thus reconstruct the whole geometry.

B.3.3 Proof of Completeness for Two-Round 3-E-DisGNN

Proof. As previous, our discussion is under the fact that the point cloud does not degenerates to a 2D or 1D point cloud. Otherwise, the problem is quite trivial.

So we first find the tuple containing three nodes that are affinely independent from the 2-round 3-E-DisGNN's output { { c 2 i | i ∈ [ N ] 3 } } , and refer to it as tuple k . Its output color c 2 k can derive the three points' coordinates. The update procedure is c 2 k = HASH( c 1 k , { { ( c 1 Φ 0 ( k ,j ) , d j, k 0 ) | j ∈ [ N ] } } , { { ( c 1 Φ 1 ( k ,j ) , d j, k 1 ) | j ∈ [ N ] } } , { { ( c 1 Φ 2 ( k ,j ) , d j, k 2 ) | j ∈ [ N ] } } ) . Since in the first multiset { { ( c 1 Φ 0 ( k ,j ) , d j, k 0 ) | j ∈ [ N ] } } , the element ( c 1 Φ 0 ( k ,j ) , d j, k 0 ) contains distances d j, k 0 , d j, k 1 , d j, k 2 , we can calculate a possible coordinates-pair multiset as P k = { {{ x (1) j , x (2) j } | j ∈ [ N ] } } . As we discussed before, the elements in P k may only have multiplicity value 1 or 2. So we discuss for two possible situations:

All the elements (coordinates pairs) in P k either have multiplicity value 2 or the two coordinates within the pair are the same. This means that we can determine all the real coordinates from P k , as we discussed in the proof of 4-WL.

There exists some element in P k , that have two distinct coordinates within the pair, and have multiplicity value 1. We denote the node corresponding to such element j 0 . We can uniquely find its relevant color (i.e., c 1 Φ 0 ( k ,j 0 ) , c 1 Φ 1 ( k ,j 0 ) , c 1 Φ 2 ( k ,j 0 ) ) among the three multisets of c 2 k 's input, because only one ( c 1 Φ m ( k ,j ) , d j, k m ) will derive the pair of possible coordinates of j 0 , which is unique. Now, as we discussed in the proof of 3-FWL, we can reconstruct the whole geometry.

B.3.4 Proof of Universality for Scalars

We have shown in the preceding three subsections that for each M ( θ ) ∈ MinMods , there exists a parameter θ 0 such that M ( θ 0 ) is complete for distinguishing all geometric graphs that are not congruent. Specifically, M ( θ 0 ) generates a unique K -dimensional intermediate representation for geometric graphs that are not congruent at the output block. According to Theorem 4.1 in Hordan et al. [2023], by passing the intermediate representation through a MLP, M ( θ 0 ) can achieve universality for scalars (note that MLPs are already included in the output block of M ).

B.3.5 Proof of Universality for Vectors

The key to proving universality for vectors of function f equiv out is Proposition 10 proposed by Villar et al. [2021]. The goal is to prove that function f = MLPs ( h m ) can approximate all functions that are O(3)-invariant and permutation-invariant with respect to all nodes except the m -th node (We denote the corresponding group as G m ).

First of all, note that h m is G m -invariant. That's because h v ( v ∈ V k ) is invariant under O(3) actions, and h m is obtained by pooling all the k -tuples whose first index is m . When permuting nodes except the m -th one, the multiset { { v | v 0 = m, v ∈ V k } } does not change. According to Theorem 4.1 in Hordan et al. [2023], to prove that f equiv out is universal, we need to show that h m can distinguish inputs X ∈ R n × 3 on different G m orbits.

To see this, recall that h m = f node ( { { h T v | v 0 = m } } ) , and as we demonstrated in the first three subsections, there always exists some h T v in the multiset { { h T v | v 0 = m } } that can reconstruct the entire geometry. For example, in the case of 3-F-DisGNN, if the point cloud does not degenerate into a 1-D point cloud, there must exist v where v 0 = m and which contains three nodes that are affinely independent (for arbitrary m ), which can reconstruct the entire geometry as demonstrated in B.3.2. It is worth noting that this reconstruction is started at the m -th node, thus if two point clouds share the same h m , they must be related by a G m action.

Finally, it is worth noting that an alternative perspective to understand the role of group G m is through the universality for node-wise representation , h m . The statement "Two geometric graphs are related by a G m action" is equivalent to the statement "the two geometric graphs are congruent, and the m -th nodes in the two geometric graphs are related by an autoisomorphism." The ability of models to generate node representations that can separate different G m orbits implies their capability to generate universal node-wise representations, which is indeed a more powerful property than global universality.

B.4 Proof of Theorem 6.5

Delle Rose et al. [2023] is our concurrent work, which demonstrated that geometric 3-round 2-FWL and 1-round 3-FWL can distinguish all non-congruent geometric graphs. We refer the readers to read the proof provided by the work, and in comparison to our previous proof, they further make good use of the global properties, i.e., the multiset of all the tuple representations, to get several desired conclusions. In this subsection, we make use of their findings: Since 3-round 2-F-DisGNN and 1-round 3-F-DisGNN are continuous versions of these methods and can achieve injectiveness with a parameter θ 0 , they can also achieve completeness.

It is important to note that the update function of 2-E-DisGNN can be used to implement that of 2-F-DisGNN: When the order is 2, the edge representation e t ij is simplified as e t ij = ( e ij , h ij ) . As a result, the neighbor component of 2-F-DisGNN { { ( h ik , h kj ) | k ∈ V } } is encompassed in the 1-neighbor component of 2-E-DisGNN { { ( h kj , e t ik ) | k ∈ V } } . Hence, the 3-round 2-E-DisGNN is also complete. When the round number is limited to 1, the update function of 3-E-DisGNN can also implement that of 3-F-DisGNN. This is due to the fact that during the initialization step, each 3-tuple is embedded based on the distance matrix. Consequently, the neighbor component of 3-F-DisGNN, { { ( h mjk , h imk , h ijm ) | m ∈ V } } , solely represents the distance matrix of ijkm , which can be obtained from any s -neighbor component ( s ∈ [3] ) of 3-E-DisGNN (It should be noted that when the iteration number is not 1, determining whether the update function of 3-E-DisGNN can implement that of 3-F-DisGNN is not trivial). Thus, the E-version DisGNNs are also complete. Moreover, as per our proof in Section B.3.4, they are also universal for scalars.

Discussion of universality for vectors. The primary obstacle to proving the universality of M∈ Minmods ext for vectors lies in the limited expressive power of individual nodes, which may not reconstruct the whole geometry. The problem cannot be easily solved by concating h m with the global universal representation h = { { h T v | v ∈ V k } }. Specifically, given ( h m 1 , h 1 ) and ( h m 2 , h 2 ) for two N -node point clouds, if the two tuples are the same, we know that the two point clouds are related by an E(3) transformation and a permutation of the N nodes. However, we cannot guarantee that the permutation (or there exists another permutation) maps the m -th node in the first point cloud to the m -th node in the second. For instance, even though two points in a point cloud may not be related by an autoisomorphism, they may still have the same representation after applying M∈ Minmods ext (Possible counterexamples may be derived from this situation). As a result, f may not be universal for all G m -invariant functions, and f equiv out may not be universal for all vector functions. We leave this as an open question.

C Detailed Model Design and Analysis

Radial basis function. In k -DisGNNs, we use radial basis functions (RBF) f rbf e : R → R H e to expand the distance between two nodes into an H e -dimension vector. This can reduce the number of learnable parameters and additionally provides a helpful inductive bias [Klicpera et al., 2020a]. The appropriate choice of RBF is beneficial, and we use nexpnorm RBF defined as

where β k , µ k are coefficients of the k th basis. Experiments show that this RBF performs better than others, such as the Bessel function used in Klicpera et al. [2020a], Gasteiger et al. [2021].

Initialization function. We realize f init of k -DisGNNs' initialization block in Section 5 as follows

where ⊙ represents Hadamard product and ⊙ represents Hadamard product over all the elements in the subscript. f emb z : Z + → R H z is a learnable embedding function, f rbf e : R + → R H e is a radial basis function, and f i z , f ij e are neural networks such as MLPs that maps the embeddings of z and e to the common continuous vector space R K .

f init can learn an injective representation for v as long as the embedding dimensions are high enough, just like passing the concatenation of z v i and e v i v j through an MLP. We chose this function form for the better experimental performance.

Note that by this means, tuples with different equality patterns [Maron et al., 2018] can be distinguished, i.e., get different representations, without explicitly incorporating the representation of equality pattern. This is because in the context of distance graphs, tuples with different equality patterns will have quite different distance matrices (elements are zero at the positions where two nodes are the same) and thus can be captured by f init .

Injective functions. We realize all the message passing functions as injective functions to ensure expressiveness. To be specific, we embed all the multisets in Equation (7, 9, 10) and in output function with the injective multiset function proposed in Xu et al. [2018], and use the matrix multiplication methods proposed in Maron et al. [2019] to implement the message passing functions of k -F-DisGNN (Equation (8)).

Inductive bias. Although theoretically there is no need to distinguish tuples with different equality patterns explicitly, we still do so to incorporate inductive bias into the model for better learning. Specifically, we modify the initialization function and the output function as follows: 1. At the initialization step, we learn embeddings for different equality patterns and incorporate them into the results of Equation (36) through a Hadamard product. 2. At the output step, we separate tuples where all sub-nodes are the same and the others into different multisets. These modifications are both beneficial for training and generalization.

D Detailed Experiments

D.1 Experimental Setup

Training Setting. For QM9, we use the mean squared error (MSE) loss for training. For MD17, we use the weighted loss function

where the force ratio ρ is fixed as 0.999 (For revised MD17, we set ρ to 0.99 for several targets since the data quality is better).

We follow the same dataset split as GemNet [Gasteiger et al., 2021]. We optimize all models using Adam [Kingma and Ba, 2014] with exponential decay and plateau decay learning rate schedulers, and also a linear learning rate warm-up. To prevent overfitting, we use early stopping on validation loss and an exponential moving average (EMA) with decay rate 0.99 for model parameters during validation and test. We follow DimeNet's [Klicpera et al., 2020a] setting to calculate ∆ ϵ by taking ϵ LUMO -ϵ HOMO and use the atomization energy for U 0 , U, H and G on QM9. Experiments are conducted on Nvidia RTX 3090 and Nvidia RTX 4090. Results are average of three runs with different random seeds. Detailed training setting can be referred to Table 4.

Model hyperparameters. The key model hyperparameters we coarsely tune are the rbf dimension, hidden dimension, and number of message passing blocks. For rbf dimension, we use 16 for MD17 and 32 for QM9. We choose the number of message passing blocks from { 4 , 5 , 6 } . For hidden dimension, we use 512 for 2-DisGNNs and 320 for 3-DisGNNs. The detailed hyperparameters can be found in our codes.

D.2 Supplementary Experimental Information

Discussion of results in MD17. In our experiments on MD17, most of the state-of-the-art performance we achieve is on force targets, and the loss on energy targets is relatively higher, see Table 1. On force prediction tasks, k -DisGNNs rank top 2 on force prediction tasks on average, and outperform the best results by a significant margin on several molecules, such as aspirin and malonaldehyde. However, the results on the energy prediction tasks are relatively lower, with 2-F-DisGNN ranking 2 nd and 3-E-DisGNN ranking 4 th .

This is due to the fact that we have assigned a quite high weight (0.999) to the force loss during training, similar to what GemNet[Gasteiger et al., 2021] does. In comparison, TorchMD [Thölke and De Fabritiis, 2021] assigned a weight of 0.8 to the force loss, resulting in better results on energy targets, but not as good results on force targets.

Actually, in molecular simulations, force prediction is a more challenging task. It determines the accuracy of molecular simulations and reflects the performance of a model better [Gasteiger et al., 2021] (One can see that the energy performance of different models does not vary much). Therefore, it makes more sense to focus on force prediction. Furthermore, previous researches [Batzner et al., 2022, Christensen and Von Lilienfeld, 2020] have found that models can achieve significantly lower

energy loss on the revised MD17 dataset than on the original MD17 dataset, while holding similar force accuracy on the two datasets. As analyzed in Batzner et al. [2022], this suggests that the noise floor on the original MD17 dataset is higher on the energies, indicating that better force prediction results are more meaningful than energy prediction results on the original MD17 datasets.

Revised MD17. For a comprehensive comparison, we also conducted experiments on the revised MD17 dataset, which has higher data quality than original MD17, and compared with two state-ofthe-art models, MACE [Batatia et al., 2022] and Allegro [Musaelian et al., 2023]. The results are shown in Table 5. 2-F-DisGNN achieved 10 best and 4 the second best results out of 20 targets. Note that MACE and Allegro are both models that leverage complex equivariant representations. This further demonstrates the high potential of distance-based models in geometry learning and molecular dynamic simulation.

Supplementary Results on QM9. We present the full results on QM9 in Table 6. We compare our model with 7 other models, including those that use invariant geometric features: SchNet [Schütt et al., 2018], DimeNet [Klicpera et al., 2020a], DimeNet++[Klicpera et al., 2020b], a model that uses group irreducible representations: Cormorant[Anderson et al., 2019], those that use first-order equivariant representations: TorchMD [Thölke and De Fabritiis, 2021], PaiNN [Schütt et al., 2021] and a model that uses local frame methods: GNN-LF [Wang and Zhang, 2022]. We calculate the average improvements of all the models relative to DimeNet and list them in the table.

Although our models exhibit significant improvements on the MD17 datasets, their performance on the QM9 datasets is relatively less remarkable. This discrepancy may be attributed to the fact that the QM9 datasets pose a greater challenge in terms of a model's generalization performance: Unlike MD17, QM9 comprises various kinds of molecules and regression targets. However, since k -DisGNNs rely purely on the fundamental geometric feature of distance, they may inherently require a larger amount of data to learn the significant geometric patterns within geometric structures and enhance their generalization performance, compared to models that explicitly incorporate equivariant representations or pre-calculated high-order geometric features. Nevertheless, we believe that such learning paradigm, which employs highly expressive models to extract information from the most basic components in a data-driven manner , has immense performance potential and will exhibit even greater performance gains when provided with more data points, increased model scales, better data quality or advanced pre-training methods, as already validated in fields such as natural language processing [Brown et al., 2020, Devlin et al., 2018], computer vision [Tolstikhin et al., 2021, Liu et al., 2021, Dosovitskiy et al., 2020], and molecule-related pre-training methods [Xia et al., 2023, Lu et al., 2023].

D.3 Time and Memory Consumption

We performed experiments on the MD17 dataset to compare the training and inference time, as well as the GPU memory consumption of our models with their counterparts DimeNet and GemNet. The experiments are conducted on Nvidia RTX 3090, and the results are shown in Table 7.

2-F-DisGNN demonstrates significantly better time and memory efficiency compared to DimeNet and GemNet. Even when compared to the quite efficient model, TorchMD, 2-F-DisGNN shows competitive efficiency. This is partly because it utilizes high-order tensors and employs dense calculation, resulting in significant acceleration on GPUs. However, due to the high theoretical complexity and high hidden dimension, 3-E-DisGNN shows relatively worse performance. Regarding this limitation, we have listed several possible solutions in Section 8 and recent work on simplifying and accelerating k -WL [Zhao et al., 2022, Morris et al., 2022] when k is large may also be applied, which is left for future work. Nonetheless, given that 2-F-DisGNN already satisfies the requirements for both theoretical (Section 6.3) and experimental (Section 7) performance, our primary focus lies on this model, and the analysis further emphasizes the practical potential of 2-F-DisGNN.

E Supplenmentary Related Work

There is a wealth of research in the field of interatomic potentials, which introduced various methods for modeling the atomic environment as representations [Bartók et al., 2010, 2013, Shapeev, 2016, Behler and Parrinello, 2007]. In particular, ACE (Atomic Cluster Expansion) [Drautz, 2019] has presented a framework that serves as a unifying approach for these methods and can calculate highorder complete polynomial basis features with small cost regardless of the body order. Building on Drautz [2019], Batatia et al. [2022] proposed a GNN variant called MACE that can effectively exploit the rich geometric information in atoms' local environment. Additionally, Joshi et al. [2023] proposed

geometric variants of the WL test to characterize the expressiveness of invariant and equivariant geometric GNNs for the general graph setting. All of Drautz [2019], Batatia et al. [2022], Joshi et al. [2023] leverage the many-body expansion and can include k -tuples during embedding atom's neighbors, allowing atoms to obtain k -order geometric information from their local environments like k -DisGNNs.

However, we note that there are key differences between these works (based on many-body expansion) and k -DisGNNs in terms of their message passing units and use of representations. The former assumes that the total energy can be approximated by the sum of energies of atomic environments of individual atoms and therefore use atoms as their message passing units, while k -DisGNNs pass messages among k -tuples and pool tuple representations to obtain global representation. In addition, while interatomic potential methods leverage equivariant representations to enhance their expressiveness, k -DisGNNs completely decouple E(3)-symmetry and only use the most basic geometric invariant feature, distance, to achieve universality and good experimental results.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes k-DisGNNs, that can use distances alone to represent higher-order interactions, instead of the direct use of angles, torsional angles, and dihedral angles. The authors try to build a bridge between geometric deep learning and the traditional expressiveness of graph representation learning without geometric information. They conduct experiments mainly on MD17 and show results on QM9 in the appendix. For MD17, 2F-Dis. and 3E-Dis. are comparable and SOTA, but for QM9 (still small geometric graphs), 2F-Dis. is worse than several previous works and 3E-Dis. can not be implemented due to high complexity, which means their method with higher k cannot extend to larger geometric graphs with more nodes.\n\nSTRENGTHS:\n1. The first work to use distances alone to represent higher order geometric information including angles, torsional angles, and dihedral angles.\n2. They establish a connection between geometric deep learning (GDL) and traditional graph representation learning (GRL).\n3. Good performances of k=2 model on MD17.\n4. Figures are good and informative.\n\nWEAKNESSES:\n1. Prohibitive Computational Costs for k > 2: The computational demands associated with the proposed method escalate dramatically when k > 2, as demonstrated in the appendix. Specifically, when k is set to 3, the model becomes computationally infeasible for the QM9 dataset, which comprises small molecules.\n\n2. Lack of Scalability and Performance for k = 2: When applied to the modestly-sized MD17 dataset, the method exhibits satisfactory performance with k = 2. However, this performance deteriorates markedly when transitioning to the more extensive QM9 dataset. The method's performance, as evaluated with 2F-Dis, fails to measure up to that of prior approaches such as TorchMD, GNN-LF, and PaiNN. \n\n3. The writing about the message passing mechanism lacks clarity and is hard to follow.\n\nMajor Concern: The major limitation of the proposed method is its inability to maintain the good performances when scaling to larger datasets, particularly QM9 concerning small molecules with not so many atoms. This significantly constrains the method's applicability and undermines its contribution.\n\nQUESTIONS:\nAs listed above in weaknesses.\n\nLIMITATIONS:\nThey mentioned the limitations in the main paper. More potential limitations are listed above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

comments: SUMMARY:
The paper proposes $k$-DistGNNs, a sequence of k-(F)WL-like models that operate on the distance matrix of a geometric graphs. These models are able to compute higher-order geometric invariants, generalise existent models, and for certain values of $k$ are shown to be able to approximate any geometrically invariant or equivariant functions. The model is evaluated against other SOTA models on standard molecular benchmarks. Additionally, the paper proposes a procedure to generate new counter-examples of geometric graphs that cannot be distinguished by MPNNs operating on the distance matrix.

STRENGTHS:
-	The paper is very well written and easy to follow. 
-	The expressive power of Geometric GNNs is a relatively new and unstudied topic, and the paper makes an important contribution in this regard. 
-	The proposed approach is a simple, but powerful, extension of k-WL to distance matrices and implicitly geometric graphs.
-	The evaluation shows that the proposed models perform very well on standard benchmarks against the latest geometric models (outperforming them in some tasks)
-	The universal approximation properties of the model are very attractive and back up the claim that the distance matrix is “enough”. Although, the question in the title is somewhat rhetorical since the distance matrix fully determines geometric graphs. So just from that, it follows the information contained inside is enough. 
-	The new counterexamples that the paper provides are very welcome as well as the general procedure used to generate them.

WEAKNESSES:
- Note that WL generalisations for Geometric GNNs were also studied in _On the Expressive Power of Geometric Graph Neural Networks_ (ICML 2023), which is not mentioned in the text as far as I could see.

QUESTIONS:
N/A

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors argue that the distance matrix associated with a geometric graph is sufficient to build powerful geometric neural networks.\nIn particular, the paper proposes a new family of graph architectures, named k-DisGNN, operating on a complete graph with node distances on the edges.\nThese architectures, inspired by the k-WL test (and its variations like the k-FWL or k-EWL tests), are able to model higher-order geometric information, which makes them more expressive than vanilla message-passing networks and ensures their universal approximator property.\n\nSTRENGTHS:\nThe paper is well written and well motivated.\nThe proposed ideas lead to a relatively simple yet expressive design to process geometric graphs.\nThe experiments support the theoretical results in the main paper.\n\n\nIn addition, the supplementary materials include a number of examples where message-passing networks fail to distinguish geometric graphs.\nThese interpretable examples can be useful to understand the limitations of existing and future architectures.\n\nWEAKNESSES:\nSome details about the proposed models were not completely clear to me, although this might be due to my little familiarity with Weisfeiler-Leman algorithms.\nStill, I think the presentation could be made a bit clearer, see Questions section.\n\nI don't see major flaws in this paper.\n\nQUESTIONS:\nI am not sure I understand the difference between Eq. 4 and 5.\nIt seems to me that the indices $j$ and $w$ are just swapped, but in the end the set contains the same tuples.\nCould you clarify this?\n\n\nlines 197-203: I think that writing the message passing block explicitly would improve the clarity of the paper (rather than explaining it by discussing how it differs from Eq. 4 and 5).\nSimilarly, in line 216, it would be useful to explicitly include the message passing equation of k-E-DisGNN.\n\n\nIn Sec. 5 k-DisGNN and k-F-DisGNN refer to two different models (using respectively the k-WL and k-FWL paradigms for message passing) but in Sec. 7 k-DisGNN seems to be used to refer to both 2-F-DisGNN and 3-E-DisGNN in an abstract way.\nThis seems confusing.\n\n\nAre you results valid only for 3D data or do these results generalize to any dimension?\n\nLIMITATIONS:\nThe authors explicitly addressed the limitations of the work in the main paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes neural networks for operating on distance matrices obtained from 3D/Euclidean point clouds. New architectures which generalise the k-WL/k-tuple higher order GNN framework (Morris et al.) to distance matrices are proposed, along with proofs of their universality/completeness for k=2 (scalar) and k=3 (vector predictions). The approach is evaluated to work as well or better than current geometric GNNs for small molecular tasks like MD17 and QM9.\n\nSTRENGTHS:\n- This work presents a **novel** k-WL framework for complete distance graphs from Euclidean point clouds. New architectures for processing k-tuple distance graphs are proposed, in a spirit similar to influential work by Morris et al. (https://arxiv.org/abs/1810.02244) on higher order GNNs for standard graphs.\n\n- I believe the theoretical results showing that architectures based on 2-WL and 3-WL variants of the proposed test are complete for scalar/vector predictions on distance graphs is **significant**.\n\n- The proposed architectures are shown to **work as well or better** than existing geometric GNNs for small molecular simulation tasks (QM9, MD17, rMD17).\n\n- The procedures to construct families of counterexamples for 1-EWL and vanilla disGNNs are described in detail and seem interesting.\n\nWEAKNESSES:\n- I don't have major concerns.\n\n- Ambiguous claim re. Geometric Deep Learning: \n    - Geometric Deep Learning is a pretty broad field beyond point clouds embedded in Euclidean space. It includes Groups, Gauges, Meshes, Grids, and Graphs. I personally felt the title and claims within the paper could be perceived ambiguously because the findings of this paper are not directly applicable across **the entire** spectrum covered by GDL; it only covers geometric *graph* learning.\n    - Is this work relevant to other areas of GDL? If so, how?\n\n- Scalability: One major downside of the proposed approach of adapting higher order tuple-based GNNs to operate on distance graphs will be the computational cost of working with anything other than small molecules.\n    - On MD17/rMD17, the models use a batch size of 2. I suppose this is due to high memory usage.\n    - Regarding the experiments on timings in the appendix: this seems like an unfair comparison b/c one would naturally expect using dense tensor-based implementations like yours to be much faster than PyG implementations of DimeNet/GemNet (especially for small molecule graphs with tens of nodes). \n    - **Is it possible to compare apples-to-apples the timing/memory usage of dense implementations of DimeNet/GemNet, too?**\n    - I see this paper's main contribution to be theoretical, so I don't think this is a major concern.\n\n- I had difficulty following Proof B.3.1. Perhaps proving a proof sketch or some warmup text outlining the proof techniques may be helpful, especially as it spans several pages. \n    - (It may be that I need more time to study the proof in further detail.)\n\nQUESTIONS:\nQuestions and clarifications:\n\n- In Related Work, for Hordan-etal, 2023: I realise this is probably concurrent work, but isn't the Gram matrix and distance matrix interchangable?\n\n- In Section 4, for the family of counterexamples for 1-EWL and vanilla disGNNs: how do you see these counterexamples being useful or informing future research? (Beyond telling us what we already know from the counterexample of Pozdnyakov-Ceriotti.)\n\nSuggestions:\n\n- Lines 85-88: I would suggest stating that these results are for complete graphs.\n\n- Lines 130-135: It was unclear what orientation information meant in this context without a coordinate system. Orientation w.r.t. what?\n\n- Line 1293 - typo.\n\n- I would personally go with Geometric Graph Learning instead of GDL. I realise this is opinionated.\n\nLIMITATIONS:\nThe authors have adequately addressed technical limitations but have not stated potential negative social impact. \n\nBeyond what the authors mention regarding scalability of their models, one major theoretical limitation of the proposed framework is that it is restricted to full geometric graphs, and the construction of complete/universal models for the general sparse graph setting remains an open question. This my be worth reiterating.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}