SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Basic methods. Assume that we want to learn  $\mathcal{O}_{\mathrm{tar}} = f_{\mathrm{tar}}(\mathcal{I}_{\mathrm{tar}})$  with our function  $\mathcal{O}_{\mathrm{fit}} = f_{\mathrm{fit}}(\mathcal{I}_{\mathrm{fit}})$ . In our proof, the form of  $\mathcal{O}_{\mathrm{tar}}$  and  $\mathcal{O}_{\mathrm{fit}}$ , as well as the form of  $\mathcal{I}_{\mathrm{tar}}$  and  $\mathcal{I}_{\mathrm{fit}}$ , are quite different. For example, consider the case where  $\mathcal{O}_{\mathrm{fit}}$  is an embedding  $h_{abc}$  for a 3-tuple  $abc$  and  $\mathcal{I}_{\mathrm{fit}} = \mathcal{I}_{\mathrm{fit}}^{(abc)}$  contains the information of all the neighbors of  $abc$ , while  $\mathcal{O}_{\mathrm{tar}}$  is an embedding  $m_{ab}$  for a 2-tuple  $ab$  and  $\mathcal{I}_{\mathrm{tar}} = \mathcal{I}_{\mathrm{tar}}^{(ab)}$  contains the information of all the neighbors of  $ab$ . Therefore, directly learning functions by  $f_{\mathrm{fit}}$  that produces exactly the same output as  $f_{\mathrm{tar}}$  is inappropriate. Instead, we will learn functions that can calculate several different outputs in a way  $f_{\mathrm{tar}}$  does and appropriately embeds them into the output of  $f_{\mathrm{fit}}$ . For example, we still consider the case mentioned before, and we want to learn a function  $f_{\mathrm{fit}}$  that can extract  $\mathcal{I}_{\mathrm{tar}}^{(ab)}, \mathcal{I}_{\mathrm{tar}}^{(ac)}, \mathcal{I}_{\mathrm{tar}}^{(bc)}, \mathcal{I}_{\mathrm{tar}}^{(ba)}, \mathcal{I}_{\mathrm{tar}}^{(ca)}, \mathcal{I}_{\mathrm{tar}}^{(cb)}$  from  $\mathcal{I}_{\mathrm{fit}}^{(abc)}$  respectively, and calculates  $m_{ab}, m_{ac}, m_{bc}, m_{ba}, m_{ca}, m_{cb}$  with these information like  $f_{\mathrm{tar}}$ , and embed them into the output  $h_{abc}$  in an injective way.
Since we realize  $f_{\mathrm{fit}}$  as a universal function approximator (such as MLPs and deep multisets),  $f_{\mathrm{fit}}$  can always learn the function we want. So the only concern is that whether we can extract exact  $\mathcal{I}_{\mathrm{tar}}$  from  $\mathcal{I}_{\mathrm{fit}}$ , i.e., whether there exists an injective function  $(\mathcal{I}_{\mathrm{tar}}^{(1)}, \mathcal{I}_{\mathrm{tar}}^{(2)}, \dots, \mathcal{I}_{\mathrm{tar}}^{(n)}) = f_{\mathrm{ext}}(\mathcal{I}_{\mathrm{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  $\mathcal{I}$  with the same superscript and subscript to represent the input of a function. We use the superscript  $(a_{1}a_{2}\dots a_{k})$  to represent some input  $\mathcal{I}$  if it's for tuple  $(a_{1}a_{2}\dots a_{k})$ . If there exists an injective function  $\mathcal{I}_{\mathrm{tar}} = f_{\mathrm{ext}}(\mathcal{I}_{\mathrm{fit}})$ , then we denote it by  $\mathcal{I}_{\mathrm{fit}} \to \mathcal{I}_{\mathrm{tar}}$ , meaning that  $\mathcal{I}_{\mathrm{tar}}$  can be derived from  $\mathcal{I}_{\mathrm{fit}}$ . If some geometric information  $\mathcal{I}_{\mathrm{geo}}$  (such as distance and angles) is contained in the distance matrix of a tuple  $a_{1}a_{2}\dots a_{k}$ , then we denote it by  $\mathcal{I}_{\mathrm{geo}} \in a_{1}a_{2}\dots 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
$$
m _ {a b} = f _ {\text {i n i t}} ^ {\mathrm {G}} \left(z _ {a}, z _ {b}, d _ {a b}\right). \tag {17}
$$
At the initial step, 3-E-DisGNN follows the paradigm outlined below:
$$
h _ {a b c} = f _ {\text {i n i t}} ^ {\mathrm {3 E}} \left(z _ {a}, z _ {b}, z _ {c}, d _ {a b}, d _ {a c}, d _ {b c}\right). \tag {18}
$$
$$
\mathcal {I} _ {\mathrm {i n i t}} ^ {\mathrm {3 E}, (a b c)} \to (\mathcal {I} _ {\mathrm {i n i t}} ^ {\mathrm {G}, (a b)}, \mathcal {I} _ {\mathrm {i n i t}} ^ {\mathrm {G}, (a c)}, \mathcal {I} _ {\mathrm {i n i t}} ^ {\mathrm {G}, (b c)}, \mathcal {I} _ {\mathrm {i n i t}} ^ {\mathrm {G}, (b a)}, \mathcal {I} _ {\mathrm {i n i t}} ^ {\mathrm {G}, (c a)}, \mathcal {I} _ {\mathrm {i n i t}} ^ {\mathrm {G}, (c b)}),
$$
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_{\mathrm{init}}^{\mathrm{3E}}$ . And thanks to the universal property of  $f_{\mathrm{init}}^{\mathrm{3E}}$ , we can approximate a function that accurately calculates these variables and injectively embeds them into  $h_{abc}$ , such that  $h_{abc} \rightarrow (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_{\mathrm{init}}^{3\mathrm{E}}$ , it is also possible to learn all the  $u_i, i \in \{a, b, c\}$  and embed them into  $h_{abc}$  at the same time with some function, since  $\mathcal{I}_{\mathrm{init}}^{3\mathrm{E},(abc)}$  contains  $z_i, i \in \{a, b, c\}$ . Therefore,  $h_{abc} \to (u_a, u_b, u_c)$  holds.
Initialization of geometric information. It is an important observation that since  $f_{\mathrm{init}}^{\mathrm{3E}}$  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_{\mathrm{init}}^{\mathrm{3E}}$  which can correctly calculate  $m_{ij}$  ( $i \neq j, i, j \in \{a, b, c\}$ ),  $u_i$  ( $i \in \{a, b, c\}$ ) and all the geometric information  $\mathcal{I}_{\mathrm{geo}}$  within the triangle  $abc$ , and injectively embed them into  $h_{abc}$ :
$$
h _ {a b c} \rightarrow \Big (\big (m _ {i j} \mid i \neq j, i, j \in \{a, b, c \} \big), \big (u _ {i} \mid (i \in \{a, b, c \}) \big), \big (\mathcal {I} _ {\mathrm {g e o}} \mid \mathcal {I} _ {\mathrm {g e o}} \in a b c \big) \Big).
$$
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
$$
u _ {a} = f _ {\text {a t o m}} ^ {\mathrm {G}} \left(\left\{\left(m _ {k a}, e _ {\mathrm {R B F}} ^ {(k a)}\right) \mid k \in [ N ] \right\}\right). \tag {19}
$$
Since this function involves the process of pooling, it cannot be learned by  $f_{\mathrm{init}}^{3\mathrm{E}}$ . However, it can be easily learned by the basic message passing layers of 3-E-DisGNN  $f_{\mathrm{MP}}^{3\mathrm{E}}$ , which is formulated as
$$
h _ {a b c} = f _ {\mathrm {M P}} ^ {3 \mathrm {E}} \left(h _ {a b c}, \{\{(h _ {k b c}, e _ {a k}) \mid k \in [ N ] \} \}, \{\{(h _ {a k c}, e _ {b k}) \mid k \in [ N ] \} \}, \{\{(h _ {a b k}, e _ {c k}) \mid k \in [ N ] \} \}\right), \tag {20}
$$
$$
\text {w h e r e} e _ {a b} = f _ {\mathrm {e}} ^ {3 \mathrm {E}} \left(d _ {a b}, \{\{h _ {k a b} \mid k \in [ N ] \} \}, \{\{h _ {a k b} \mid k \in [ N ] \} \}, \{\{h _ {a b k} \mid k \in [ N ] \} \}\right). \tag {21}
$$
We now want to learn a function  $f_{\mathrm{MP}}^{3\mathrm{E}}$  that updates the  $u_{a}, u_{b}, u_{c}$  embedded in  $h_{abc}$  like  $f_{\mathrm{atom}}^{\mathrm{G}}$  and keep the other variables and information unchanged. Note that  $h_{abc}$  is the first input of  $f_{\mathrm{MP}}^{3\mathrm{E}}$ , 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_{\mathrm{MP}}^{3\mathrm{E}}$ 's input. In fact, the following derivation holds
$$
\begin{array}{l} \mathcal {I} _ {\mathrm {M P}} ^ {\mathrm {3 E}, (a b c)} \rightarrow \{\{(h _ {a k c}, e _ {b k}) \mid k \in [ N ] \} \} \rightarrow \{\{h _ {a k c} \mid k \in [ N ] \} \} \\ \rightarrow \{\{(m _ {k a}, d _ {k a}) \mid k \in [ N ] \} \} \rightarrow \{\{(m _ {k a}, e _ {\mathrm {R B F}} ^ {(k a)}) \mid k \in [ N ] \} \} = \mathcal {I} _ {\mathrm {a t o m}} ^ {\mathrm {G}, (a)}. \\ \end{array}
$$
Note that  $d_{ka} \in abc$ , and  $e_{\mathrm{RBF}}^{(ka)}$  can be calculated from  $d_{ka}$ . Similarly, we can derive that  $\mathcal{I}_{\mathrm{MP}}^{3\mathrm{E},(abc)} \to \mathcal{I}_{\mathrm{atom}}^{\mathrm{G},(b)}$  and  $\mathcal{I}_{\mathrm{MP}}^{3\mathrm{E},(abc)} \to \mathcal{I}_{\mathrm{atom}}^{\mathrm{G},(c)}$ . 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
$$
\mathrm {T} - \mathrm {M P}: m _ {a b} = f _ {\mathrm {T M P}} ^ {\mathrm {G}} \left(\left\{\left(m _ {k b}, e _ {\mathrm {R B F}} ^ {(k b)}, e _ {\mathrm {C B F}} ^ {(a b k)}\right) \mid k \in [ N ], k \neq a \right\}\right) \tag {22}
$$
$$
\mathrm {Q} - \mathrm {M P}: m _ {a b} = f _ {\mathrm {Q M P}} ^ {\mathrm {G}} \left(\left\{\left(m _ {k _ {2} k _ {1}}, e _ {\mathrm {R B F}} ^ {(k _ {2} k _ {1})}, e _ {\mathrm {C B F}} ^ {(b k _ {1} k _ {2})}, e _ {\mathrm {S B F}} ^ {(a b k _ {1} k _ {2})}\right) \mid k _ {1}, k _ {2} \in [ N ], k _ {1} \neq a, k _ {2} \neq b, a \right\}\right) \tag {23}
$$
Note that what we need to construct is a function  $f_{\mathrm{MP}}^{\mathrm{3E}}$  that can update the information about  $m_{ij}, i, j \in \{a, b, c\}$ ,  $i \neq j$  embedded in  $h_{abc}$  just like what  $f_{\mathrm{TMP}}^{\mathrm{G}}$  and  $f_{\mathrm{QMP}}^{\mathrm{G}}$  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
$$
\begin{array}{l} \mathcal {I} _ {\mathrm {M P}} ^ {\mathrm {3 E}, (a b c)} \rightarrow \left(h _ {a b c}, \{\{(h _ {k b c}, e _ {a k}) \mid k \in [ N ] \}\}\right)\rightarrow \{\{(h _ {a b c}, h _ {k b c}, e _ {a k}) \mid k \in [ N ] \} \} \\ \rightarrow \{\{(m _ {k b}, d _ {k b}, d _ {a b}, \phi_ {a b k}) \mid k \in [ N ] \} \} \rightarrow \{\{(m _ {k b}, e _ {\mathrm {R B F}} ^ {(k b)}, e _ {\mathrm {C B F}} ^ {(a b k)}) \mid k \in [ N ], k \neq a \} \} = \mathcal {I} _ {\mathrm {T M P}} ^ {\mathrm {G}, (a b)}. \\ \end{array}
$$
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}$ ,  $\phi_{abk}$  can be obtained from it, and since  $e_{\mathrm{RBF}}^{(kb)}$  and  $e_{\mathrm{CBF}}^{(abk)}$  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:
$$
m _ {a b} = f _ {\mathrm {Q M P ^ {\prime}}} ^ {\mathrm {G}} \left(\left\{\left\{\left(m _ {k _ {2} k _ {1}}, e _ {\mathrm {R B F}} ^ {(k _ {2} k _ {1})}, e _ {\mathrm {C B F}} ^ {(b k _ {1}, k _ {2})}, e _ {\mathrm {S B F}} ^ {(a b k _ {1} k _ {2})}\right) \mid k _ {1} \in [ N ], k _ {1} \neq a \right\} \mid k _ {2} \in [ N ], k _ {2} \neq b, a \right\}\right). \tag {24}
$$
Note that the expressiveness of two-step pooling is not less than two-order pooling, i.e.
$$
\begin{array}{l} \{\{\{(m _ {k _ {2} k _ {1}}, e _ {\mathrm {R B F}} ^ {(k _ {2} k _ {1})}, e _ {\mathrm {C B F}} ^ {(b k _ {1} k _ {2})}, e _ {\mathrm {S B F}} ^ {(a b k _ {1} k _ {2})}) \mid k _ {1} \in [ N ], k _ {1} \neq a \} \} \mid k _ {2} \in [ N ], k _ {2} \neq b, a \} \} \\ \rightarrow \left\{\left(m _ {k _ {2} k _ {1}}, e _ {\mathrm {R B F}} ^ {(k _ {2} k _ {1})}, e _ {\mathrm {C B F}} ^ {(b k _ {1} k _ {2})}, e _ {\mathrm {S B F}} ^ {(a b k _ {1} k _ {2})}\right) \mid k _ {1}, k _ {2} \in [ N ], k _ {1} \neq a, k _ {2} \neq b, a \right\}. \\ \end{array}
$$
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_{\mathrm{QMP}}^{\mathrm{G}}$ , we first learn a function by  $f_{\mathrm{MP}}^{3\mathrm{E}}$  that calculates an intermediate variable  $w_{c} = f_{\mathrm{inter1}}^{3\mathrm{E}}\big(\{(m_{ck},e_{\mathrm{RBF}}^{(ck)},e_{\mathrm{CBF}}^{(bkc)},e_{\mathrm{SBF}}^{(abkc)})|k\in [N],k\neq a\} \} \big)$  by pooling all the  $m_{ck}$  at index  $k$ , which is feasible because the following derivation holds
$$
\begin{array}{l} \mathcal {I} _ {\mathrm {M P}} ^ {\mathrm {3 E}, (a b c)} \rightarrow \left(h _ {a b c}, \{\{(h _ {a b k}, e _ {c k}) \mid k \in [ N ] \}\}\right)\rightarrow \{\{(h _ {a b c}, h _ {a b k}, e _ {c k}) \mid k \in [ N ] \} \} \\ \rightarrow \{\{(m _ {c k}, d _ {c k}, d _ {b k}, \phi_ {a b k}, \theta_ {a b k c}) \mid k \in [ N ] \} \} \\ \rightarrow \left\{\left(m _ {c k}, e _ {\mathrm {R B F}} ^ {(c k)}, e _ {\mathrm {C B F}} ^ {(b k c)}, e _ {\mathrm {S B F}} ^ {(a b k c)}\right) \mid k \in [ N ], k \neq a \right\} = \mathcal {I} _ {\mathrm {i n t e r}} ^ {3 \mathrm {E}, (c)}. \\ \end{array}
$$
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_{\mathrm{MP}}^{3\mathrm{E}}$  but this time we learn a function that just pools all the  $w_{c}$  in  $h_{abc}$  and finally updates  $m_{ab}$ :
$$
\mathcal {I} _ {\mathrm {M P}} ^ {\mathrm {3 E}, (a b c)} \to \{\{(h _ {a b k}, e _ {c k}) \mid k \in [ N ] \} \} \to \{\{h _ {a b k} \mid k \in [ N ] \} \} \to \{\{w _ {k} \mid k \in [ N ], k \neq b, a \} \}.
$$
This means we can realize  $f_{\mathrm{QMP}}^{\mathrm{G}}$  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_{\mathrm{MP}}^{3\mathrm{E}}$  in Appendix B.1.1. The second procedure can be formulated as
$$
m _ {a b} = f _ {\text {s e l f - i n t e r}} ^ {\mathrm {G}} \left(u _ {a}, u _ {b}, m _ {a b}\right). \tag {25}
$$
It is obvious that we can update  $m_{ab}$  in  $h_{abc}$  by  $f_{\mathrm{MP}}^{3\mathrm{E}}$  because the following derivation holds
$$
\mathcal {I} _ {\mathrm {M P}} ^ {\mathrm {3 E}, (a b c)} \to h _ {a b c} \to \left(u _ {a}, u _ {b}, m _ {a b}\right) = \mathcal {I} _ {\mathrm {s e l f - i n t e r}} ^ {\mathrm {G}, (a b)}.
$$
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
$$
t = \sum_ {a \in [ N ]} W _ {\text {o u t}} \left(f _ {\text {a t o m}} ^ {\mathrm {G}} \left(\left\{\left(m _ {k a}, e _ {\mathrm {R B F}} ^ {(k a)}\right) \mid k \in [ N ] \right\}\right)\right), \tag {26}
$$
where  $W_{\mathrm{out}}$  is a learnable matrix.
And our output function is
$$
t = f _ {\text {o u t p u t}} ^ {3 \mathrm {E}} \left(\{\{h _ {a b c} \mid a, b, c \in [ N ] \} \}\right). \tag {27}
$$
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.
$$
\mathcal {I} _ {\mathrm {o u t p u t}} ^ {\mathrm {3 E}, (a b c)} \to \left\{\left\{h _ {a a a} \mid a \in [ N ] \right\} \right\} \to \left\{\left\{u _ {a} \mid a \in [ N ] \right\} \right\}.
$$
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 hold 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
$$
m _ {a b} = f _ {\text {i n i t}} ^ {\mathrm {D}} \left(z _ {a}, z _ {b}, d _ {a b}\right). \tag {28}
$$
What 2-F-DisGNN does at the initialization step is
$$
h _ {a b} = f _ {\text {i n i t}} ^ {2 \mathrm {F}} \left(z _ {a}, z _ {b}, d _ {a b}\right). \tag {29}
$$
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
$$
\mathcal {I} _ {\mathrm {i n i t}} ^ {\mathrm {2 F}, (a b)} \to (\mathcal {I} _ {\mathrm {i n i t}} ^ {\mathrm {D}, (a b)}, \mathcal {I} _ {\mathrm {i n i t}} ^ {\mathrm {D}, (b a)}).
$$
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
$$
m _ {a b} = f _ {\mathrm {M P}} ^ {\mathrm {D}} \left(\left\{\left(m _ {k a}, e _ {\mathrm {R B F}} ^ {(a b)}, e _ {\mathrm {C B F}} ^ {(k a b)}\right) \mid k \in [ N ], k \neq b \right\}\right). \tag {30}
$$
And the message passing framework of 2-F-DisGNN can be formulated as
$$
h _ {a b} = f _ {\mathrm {M P}} ^ {2 \mathrm {F}} \left(h _ {a b}, \{(h _ {k b}, h _ {a k}) \mid k \in [ N ], k \neq b \}\right). \tag {31}
$$
Now we want to learn a function  $f_{\mathrm{MP}}^{2\mathrm{F}}$  that can updates the  $m_{ab}$  and  $m_{ba}$  embedded in  $h_{ab}$  like what  $f_{\mathrm{MP}}^{\mathrm{D}}$  does. We need to check if  $f_{\mathrm{MP}}^{2\mathrm{F}}$  have sufficient information to update  $m_{ab}$  and  $m_{ba}$  embedded in its output  $h_{ab}$ . The derivation holds:
$$
\begin{array}{l} \mathcal {I} _ {\mathrm {M P}} ^ {2 \mathrm {F}, (a b)} \rightarrow \left\{\left\{\left(h _ {a b}, h _ {k b}, h _ {a k}\right) \mid k \in [ N ] \right\}\right\}\rightarrow \left\{\left\{\left(m _ {k a}, d _ {a b}, d _ {k a}, \phi_ {k a b}\right) \mid k \in [ N ], k \neq b \right\}\right. \\ \rightarrow \{\{(m _ {k a}, e _ {\mathrm {R B F}} ^ {(a b)}, e _ {\mathrm {C B F}} ^ {(k a b)}) \mid k \in [ N ], k \neq b \} \} = \mathcal {I} _ {\mathrm {M P}} ^ {\mathrm {D}, (a b)}. \\ \end{array}
$$
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  $\mathcal{I}_{\mathrm{MP}}^{2\mathrm{F},(ab)} \to \mathcal{I}_{\mathrm{MP}}^{\mathrm{D},(ba)}$ .
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
$$
u _ {a} = f _ {\text {a t o m}} ^ {\mathrm {D}} \left(\left\{\left(m _ {k a}, e _ {\mathrm {R B F}} ^ {(k a)}\right) \mid k \in [ N ] \right\}\right), \tag {32}
$$
$$
t = \sum_ {a \in [ N ]} u _ {a}. \tag {33}
$$
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
$$
\begin{array}{l} \mathcal {I} _ {\mathrm {M P}} ^ {2 \mathrm {F}, (a b)} \rightarrow \left\{\left(h _ {k b}, h _ {a k}\right) \mid k \in [ N ] \right\}\rightarrow \left\{\left. h _ {a k} \mid k \in [ N ] \right\}\right\} \\ \rightarrow \{\{(m _ {k a}, d _ {k a}) \mid k \in [ N ] \} \} \rightarrow \{\{(m _ {k a}, e _ {\mathrm {R B F}} ^ {(k a)}) \mid k \in [ N ] \} \} = \mathcal {I} _ {\mathrm {a t o m}} ^ {\mathrm {D}, (a)}. \\ \end{array}
$$
Similarly we can get  $\mathcal{I}_{\mathrm{MP}}^{2\mathrm{F},(ab)} \to \mathcal{I}_{\mathrm{atom}}^{\mathrm{D},(b)}$ .
And the output block of 2-F-DisGNN is formulated as follows
$$
t = f _ {\text {o u t p u t}} ^ {2 \mathrm {F}} \left(\left\{\left\{h _ {a b} \mid a, b \in [ N ] \right\} \right\}\right). \tag {34}
$$
Note that the following derivation holds
$$
\mathcal {I} _ {\text {o u t p u t}} ^ {2 \mathrm {F}, (a b)} \rightarrow \left\{\left\{h _ {a a} \mid a \in [ N ] \right\}\right\}\rightarrow \left\{\left\{u _ {a} \mid a \in [ N ] \right\}\right\}.
$$
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  $\mathcal{M}(\theta) \in \mathrm{MinMod}_{\mathrm{s}} = \{1\text{-round 4 - DisGNN}, 2\text{-round 3 - E - DisGNN}, 2\text{-round 3 - F - DisGNN}\}$  with parameters  $\theta$ . Denote  $h_m = f_{\mathrm{node}}\left(\left\{\left[h_{\mathbf{v}}^T \mid \mathbf{v} \in V^k, \mathbf{v}_0 = m\right]\right)\right) \in \mathbb{R}^{K'}$  as node  $m$ 's representation produced by  $\mathcal{M}$  where  $f_{\mathrm{node}}$  is an injective multiset function, and  $\mathbf{x}_m^c$  as node  $m$ 's coordinates w.r.t. the center. Denote  $\mathrm{MLPs}(\theta') : \mathbb{R}^{K'} \to \mathbb{R}$  as a multi-layer perceptron with parameters  $\theta'$ .
1. (Completeness) Given arbitrary two geometric graphs  $\mathbf{X}_1, \mathbf{X}_2 \in \mathbb{R}^{3 \times n}$ , then there exists  $\theta_0$  such that  $\mathcal{M}(\mathbf{X}_1; \theta_0) = \mathcal{M}(\mathbf{X}_2; \theta_0) \iff \mathbf{X}_1$  and  $\mathbf{X}_1$  are congruent.
2. (Universality for scalars) Let  $f: \mathbb{R}^{3 \times n} \to \mathbb{R}$  be an arbitrary continuous,  $E(3)$  invariant and permutation invariant function over geometric graphs, then for any compact set  $M \subset \mathbb{R}^{3 \times n}$  and  $\epsilon > 0$ , there exists  $\theta_0$  such that  $\forall \mathbf{X} \in M, |f(\mathbf{X}) - \mathcal{M}(\mathbf{X};\theta_0)| \leq \epsilon$ .
3. (Universality for vectors) Let  $f: \mathbb{R}^{3 \times n} \to \mathbb{R}^3$  be an arbitrary continuous,  $O(3)$  equivariant, permutation and translation invariant function over geometric graphs, let  $f_{\mathrm{out}}^{\mathrm{equiv}} = \sum_{m=1}^{|V|} \mathrm{MLPs}\big(h_m\big)\mathbf{x}_m^c$ , then for any compact set  $M \subset \mathbb{R}^{3 \times n}$  and  $\epsilon > 0$ , there exists  $\theta_0$  for  $\mathcal{M}$  and  $\theta_0'$  for MLPs such that  $\forall \mathbf{X} \in M, |f(\mathbf{X}) - f_{\mathrm{out}}^{\mathrm{equiv}}(\mathbf{X};\theta_0,\theta_0')| \leq \epsilon$ .
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  $\theta$  of  $\mathcal{M}(\theta)$  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.