SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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_i^1\mid i\in [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_{i}^{0}) = 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_{i}^{t} = \mathrm{HASH}(c_{i}^{t - 1},C_{0,i}^{t - 1},C_{1,i}^{t - 1},C_{2,i}^{t - 1},C_{3,i}^{t - 1})$  where  $C_{m,i}^{t - 1}$  denotes the color multiset of tuple  $i$ 's  $m$ -th neighbors, there also exists a series functions  $f^{t}$  such that  $f^{t}(c_{i}^{t}) = c_{i}^{t - 1}$ . 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}\circ 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_{\pmb{k}}^{1}$ .
First, we can reconstruct the 4 points' 3D coordinates (given any arbitrary center and orientation) from  $\pmb{k}$ 's distance matrix (Theorem 3.1). Formally, there exists a function  $f^{D}(D_{k}) = X_{k}$ , where  $X_{k} \in \mathbb{R}^{4*3}$  represents the coordinates. And since  $f^{1} \circ f^{0}(c_{k}^{1}) = D_{k}$ , we have that  $f^{1} \circ f^{0} \circ f^{D}(c_{k}^{1}) = X_{k}$ .
The update function of  $c_{\pmb{k}}$  is  $c_{\pmb{k}}^{1} = \mathrm{HASH}(c_{\pmb{k}}^{0}, C_{0,\pmb{k}}^{0}, C_{1,\pmb{k}}^{0}, C_{2,\pmb{k}}^{0}, C_{3,\pmb{k}}^{0})$ , where  $C_{m,\pmb{k}}^{0} = \{\{c_{\Phi_m(\pmb{k}, j)}^{0} \mid j \in [N]\}\}$ ,  $m \in [4]$  and  $\Phi_m(\pmb{k}, j)$  replaces the  $m$ -th element in tuple  $\pmb{k}$  with  $j$ . Since the function is HASH function, we can also reconstruct each of  $C_{m,\pmb{k}}^{0}$  from  $c_{\pmb{k}}^{1}$ . In  $C_{m,\pmb{k}}^{0}$ , there exists  $N$  4-tuples' initial color, from which we can reconstruct the 4-tuple's distance matrix. Note that given a color  $c_{\Phi_m(\pmb{k}, j)}^{0}$  in the multiset, we can reconstruct 3 distances, namely  $d(x_j, x_{\pmb{k}_{(m+1)\%4}})$ ,  $d(x_j, x_{\pmb{k}_{(m+2)\%4}})$ ,  $d(x_j, x_{\pmb{k}_{(m+3)\%4}})$ , where  $x_j$  represents the 3D coordinates of point  $j$  and  $d(x, y)$  calculates the  $l_2$  norm of  $x, y$ .
There is a strong geometric constraint in 3D space: Given the distances between a point and 3 affinely independent points (whose positions are known), one can calculate at most 2 possible positions of the point, and the 2 possible positions 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_{m,k}^{0}$ , 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_{j}^{(1)}, x_{j}^{(2)}\} \mid j \in [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}$ :
1. Given arbitrary two elements (pairs) of  $P_{m,\pmb{k}}$ , denoted as  $p_{j_1}$  and  $p_{j_2}$ , we have  $p_{j_1} = p_{j_2}$  or  $p_{j_1} \cap p_{j_2} = \emptyset$ . These correspond to two possible situations:  $j_1$  and  $j_2$  are either mirror-symmetric relative to the plane formed by  $\pmb{k}_{(m + 1)\%4}, \pmb{k}_{(m + 2)\%4}, \pmb{k}_{(m + 3)\%4}$ , or not.
2. Arbitrary element  $p_j$  of  $P_{m,\mathbf{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  $\mathbf{k}_{(m + 1)\%4}$ ,  $\mathbf{k}_{(m + 2)\%4}$ ,  $\mathbf{k}_{(m + 3)\%4}$ , or does not.
The follows proves that the four 0-dimension manifolds determined by  $P_{m,\pmb{k}}, m \in [4]$  can intersect at a unique solution, which is the real geometry.
We first "refine" the multiset  $P_{m,\pmb{k}}$ . There are 3 kinds of pairs in  $P_{m,\pmb{k}}$ :
1. The pair  $\{x_j^{(1)}, x_j^{(2)}\}$  that have multiplicity value 2. This means that there are two mirror-symmetric 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_j^{(1)}, x_j^{(2)}$  both exist.
2. The pair where  $x_{j}^{(1)} = x_{j}^{(2)} = 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).
3. 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}^{\prime}$ . Note that now in each  $P_{m,k}^{\prime}$ :
1. Arbitrary two elements intersect at  $\varnothing$
2. Arbitrary element has two distinct possible coordinates, one is real and the other is fake.
3. The two properties above also mean that elements in multiset  $\bigcup_{p\in P_m^{\prime},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}^{\prime}$  by removing the pairs that contain coordinates in  $(A_0 \cup A_1 \cup A_2 \cup A_3) - A_m$ . Note that according to the property of  $P_{m,k}^{\prime}$  mentioned above, each coordinates in  $(A_0 \cup A_1 \cup A_2 \cup A_3) - A_m$  will and will only correspond to one pair in  $P_{m,k}^{\prime}$ . Denote the further refined sets as  $P_{m,k}^{\prime \prime}$ . After the refinement, each  $P_{m,k}^{\prime \prime}$  has equally  $M = N - |A_0 \cup A_1 \cup A_2 \cup A_3|$  elements (pairs), corresponding to  $M$  undetermined real coordinates.
We use  $p \cap P$  to denote the multiset  $\{(p' \mid p' \in P, p' \cap p \neq \emptyset)\}$ . Now we prove that  $\exists m \in [4], \exists p \in P_{m,k}''$ ,  $\exists m' \in [4] - m$ , s.t.  $|p \cap P_{m',k}''| = 1$  by contradiction.
Note that  $\forall m\in [4],\forall p\in P_{m,k}^{\prime \prime},\forall m^{\prime}\in [4] - m,|p\cap P_{m^{\prime},k}^{\prime \prime}|\in \{1,2\}$ . That's because each element in  $P_{m,k}^{\prime \prime}$  contains one real coordinates and one fake coordinates, each real/fake coordinates can at most intersect with one pair in  $P_{m^{\prime},k}^{\prime \prime}$ , and there must exist one element in  $P_{m^{\prime},k}^{\prime \prime}$  that contains the same real coordinates.
So opposite to the conclusion is:  $\forall m\in [4],\forall p\in P_{m,\pmb{k}}^{\prime \prime},\forall m^{\prime}\in [4] - m,|p\cap P_{m^{\prime},\pmb{k}}^{\prime \prime}| = 2$  , we assume this holds. Now there are also two properties:
1. For all  $m \in [4]$ , elements in multiset  $\bigcup_{p \in P_{m,k}^{\prime\prime}} p$  has multiplicity 1.
2. For all  $m \in [4]$ ,  $\bigcup_{p \in P_{m,k}^{\prime\prime}} p$  with different  $m$  shares the same  $M$  real coordinates.
So we have that:  $\bigcup_{p\in P_{m,k}^{\prime \prime}}p$  is the same across all  $m\in [4]$ . Note that for each  $p\in P_{m,k}^{\prime \prime}$ , since the two coordinates in  $p$  are mirror-symmetric relative to plane  $\pmb{k}_{(m + 1)\% 4},\pmb{k}_{(m + 2)\% 4},\pmb{k}_{(m + 3)\% 4}$  the midpoint of the two points is on the plane. This means that by averaging all the coordinates in  $\bigcup_{p\in P_{m,k}^{\prime \prime}}p$ , we have a center point on the plane  $\pmb{k}_{(m + 1)\% 4},\pmb{k}_{(m + 2)\% 4},\pmb{k}_{(m + 3)\% 4}$ . However, since  $\bigcup_{p\in P_{m,k}^{\prime \prime}}p$  for all  $m\in [4]$  are the same, we have one same center point on each four planes  $\pmb{k}_{(m + 1)\% 4},\pmb{k}_{(m + 2)\% 4},\pmb{k}_{(m + 3)\% 4},m\in [4]$ . This contradicts to the assumption that  $\pmb{k}_1,\pmb{k}_2,\pmb{k}_3,\pmb{k}_4$  are affinely independent points. So the conclusion that  $\exists m\in [4],\exists p\in P_{m,k}^{\prime \prime},\exists m^{\prime}\in [4] - m$ , s.t.  $|p\cap P_{m',k}''| = 1$  holds.
Since the conclusion holds, we can obtain the  $p$  and the unique element  $p' \in p \cap P_{m',k}''$ . Moreover, we have that  $|p' \cap p| = 1$ . This is because, the two pairs share one real coordinates, and the fake coordinates must be different (The two planes formed by  $\mathbf{k}_{(m+1)\%4}$ ,  $\mathbf{k}_{(m+2)\%4}$ ,  $\mathbf{k}_{(m+3)\%4}$  and  $\mathbf{k}_{(m'+1)\%4}$ ,  $\mathbf{k}_{(m'+2)\%4}$ ,  $\mathbf{k}_{(m'+3)\%4}$  are not the same). So we obtain the real coordinates by  $p' \cap p$ .
Now we further refine the  $4P_{m,k}^{\prime \prime}$  by deleting the pair that contains the real coordinates we just found, which makes each  $P_{m,k}^{\prime \prime}$  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_{k}^{1}$ . This means that when generating  $c_{k}^{1}$ , 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 degenerate 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_i^2\mid i\in [N]^3\}\}$ . Denote the founded tuple as  $k$ . We first calculate the 3 nodes' coordinates according to the  $c_{k}^{0}$  derived from  $c_{k}^{2}$ .
As the update procedure,  $c_{\pmb{k}}^{2} = \mathrm{HASH}(c_{\pmb{k}}^{1}, \{(c_{\Phi_{0}(\pmb{k}, j)}^{1}, c_{\Phi_{1}(\pmb{k}, j)}^{1}, c_{\Phi_{2}(\pmb{k}, j)}^{1}) \mid j \in N\})$ , each tuple  $(c_{\Phi_{0}(\pmb{k}, j)}^{1}, c_{\Phi_{1}(\pmb{k}, j)}^{1}, c_{\Phi_{2}(\pmb{k}, j)}^{1})$  in the multiset contains the distance  $d_{\pmb{k}_{0}, j}, d_{\pmb{k}_{1}, j}, d_{\pmb{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  $\pmb{k}$ . Like in the proof of 4-WL, we denote the final calculaoed possible coordinates-pair multiset as  $P_{\pmb{k}} = \{\{x_{j}^{(1)}, x_{j}^{(2)}\} \mid j \in [N]\}$ .
We then further find a node  $j_0$  that is not on plane  $\pmb{k}$  from the multiset term in  $c_{\pmb{k}}^{2}$ , and its colors' tuple is  $(c_{\Phi_0(\pmb{k},j_0)}^1, c_{\Phi_1(\pmb{k},j_0)}^1, c_{\Phi_2(\pmb{k},j_0)}^1)$ . 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_{\Phi_m(\pmb{k},j_0)}$  from each of  $c_{\Phi_m(\pmb{k},j_0)}^1$  where  $m \in [3]$ .
Note that the four plane,  $\pmb{k}, \Phi_0(\pmb{k}, j_0), \Phi_1(\pmb{k}, j_0), \Phi_2(\pmb{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
So we first find the tuple containing three nodes that are affinely independent from the 2-round 3-E-DisGNN's output  $\{\{c_i^2\mid i\in [N]^3\}\}$ , and refer to it as tuple  $\pmb{k}$ . Its output color  $c_{k}^{2}$  can derive the three points' coordinates. The update procedure is  $c_{k}^{2} = \mathrm{HASH}(c_{k}^{1},\{(c_{\Phi_{0}(\pmb {k},j)}^{1},d_{j,\pmb{k}_{0}})\mid j\in$ $[N]\} ,\{(c_{\Phi_1(\pmb {k},j)},d_{j,\pmb {k}_1})\mid j\in [N]\} ,\{(c_{\Phi_2(\pmb {k},j)}^{1},d_{j,\pmb {k}_2})\mid j\in [N]\} \}$ . Since in the first multiset  $\{(c_{\Phi_0(\pmb {k},j)},d_{j,\pmb {k}_0})\mid j\in [N]\}$ , the element  $(c_{\Phi_0(\pmb {k},j)}^1,d_{j,\pmb {k}_0})$  contains distances  $d_{j,\pmb {k}_0},d_{j,\pmb {k}_1},d_{j,\pmb {k}_2}$  we can calculate a possible coordinates-pair multiset as  $P_{k} = \{\{x_{j}^{(1)},x_{j}^{(2)}\} \mid j\in [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:
1. 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.
2. 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_{\Phi_0(\pmb{k},j_0)}^1, c_{\Phi_1(\pmb{k},j_0)}^1, c_{\Phi_2(\pmb{k},j_0)}^1$ ) among the three multisets of  $c_{\pmb{k}}^{2}$ 's input, because only one  $(c_{\Phi_m(\pmb{k},j)}^1, d_j, \pmb{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  $\mathcal{M}(\theta) \in \mathrm{MinMod}_s$ , there exists a parameter  $\theta_0$  such that  $\mathcal{M}(\theta_0)$  is complete for distinguishing all geometric graphs that are not congruent. Specifically,  $\mathcal{M}(\theta_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,  $\mathcal{M}(\theta_0)$  can achieve universality for scalars (note that MLPs are already included in the output block of  $\mathcal{M}$ ).
B.3.5 Proof of Universality for Vectors
The key to proving universality for vectors of function  $f_{\mathrm{out}}^{\mathrm{equiv}}$  is Proposition 10 proposed by Villar et al. [2021]. The goal is to prove that function  $f = \mathrm{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  $\mathcal{G}_m$ ).
First of all, note that  $h_m$  is  $\mathcal{G}_m$ -invariant. That's because  $h_v$  ( $v \in 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  $\{\langle v \mid v_0 = m, v \in V^k \rangle\}$  does not change. According to Theorem 4.1 in Hordan et al. [2023], to prove that  $f_{\mathrm{out}}^{\mathrm{equiv}}$  is universal, we need to show that  $h_m$  can distinguish inputs  $\mathbf{X} \in \mathbb{R}^{n \times 3}$  on different  $\mathcal{G}_m$  orbits.
To see this, recall that  $h_m = f_{\mathrm{node}}\big(\{\{h_v^T \mid v_0 = m\}\} \big)$ , and as we demonstrated in the first three subsections, there always exists some  $h_v^T$  in the multiset  $\{\{h_v^T \mid 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  $\mathbf{v}$  where  $\mathbf{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  $\mathcal{G}_m$  action.
Finally, it is worth noting that an alternative perspective to understand the role of group  $\mathcal{G}_m$  is through the universality for node-wise representation,  $h_m$ . The statement "Two geometric graphs are related by a  $\mathcal{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  $\mathcal{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  $\theta_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_{ij}^{t}$  is simplified as  $e_{ij}^{t} = (e_{ij}, h_{ij})$ . As a result, the neighbor component of 2-F-DisGNN  $\{(h_{ik}, h_{kj}) \mid k \in V\}$  is encompassed in the 1-neighbor component of 2-E-DisGNN  $\{(h_{kj}, e_{ik}^{t}) \mid k \in 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}) \mid m \in V\}$ , solely represents the distance matrix of  $ijkm$ , which can be obtained from any  $s$ -neighbor component ( $s \in [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  $\mathcal{M} \in \mathrm{Minmod}_{\mathrm{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 concatenating  $h_m$  with the global universal representation  $h = \{\{h_{\pmb{v}}^T \mid \pmb{v} \in V^k\}$ . Specifically, given  $(h_{m1}, h_1)$  and  $(h_{m2}, 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  $\mathcal{M} \in \mathrm{Minmod}_{\mathrm{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_{\mathrm{out}}^{\mathrm{equiv}}$  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_{\mathrm{e}}^{\mathrm{rbf}}: \mathbb{R} \to \mathbb{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
$$
f _ {e} ^ {\operatorname {r b f}} \left(e _ {i j}\right) [ k ] = e ^ {- \beta_ {k} \left(\exp \left(- e _ {i j}\right) - \mu_ {k}\right) ^ {2}}, \tag {35}
$$
where  $\beta_{k},\mu_{k}$  are coefficients of the  $k^{\mathrm{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_{\mathrm{init}}$  of  $k$ -DisGNNs' initialization block in Section 5 as follows
$$
f _ {\text {i n i t}} (\boldsymbol {v}) = \bigodot_ {i \in [ k ]} f _ {z} ^ {i} \left(f _ {z} ^ {\text {e m b}} \left(z _ {v _ {i}}\right)\right) \odot \bigodot_ {i, j \in [ k ], i <   j} f _ {e} ^ {i j} \left(f _ {\mathrm {e}} ^ {\text {r b f}} \left(e _ {v _ {i} v _ {j}}\right)\right), \tag {36}
$$
where  $\odot$  represents Hadamard product and  $\odot$  represents Hadamard product over all the elements in the subscript.  $f_{z}^{\mathrm{emb}}:\mathbb{Z}^{+}\to \mathbb{R}^{H_{z}}$  is a learnable embedding function,  $f_{\mathrm{e}}^{\mathrm{rbf}}:\mathbb{R}^{+}\to \mathbb{R}^{H_{e}}$  is a radial basis function, and  $f_{z}^{i},f_{e}^{ij}$  are neural networks such as MLPs that maps the embeddings of  $z$  and  $e$  to the common continuous vector space  $\mathbb{R}^K$ .
$f_{\mathrm{init}}$  can learn an injective representation for  $\pmb{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_{\mathrm{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)).