SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} = \mathbb {E} _ {q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L})} \left[ \mathbb {E} _ {q (\boldsymbol {y} ^ {(1: T)} | \boldsymbol {y})} \left[ \log p _ {\boldsymbol {\theta}} (\boldsymbol {y}, \boldsymbol {y} ^ {(1: T)} | G) - \log q (\boldsymbol {y} ^ {(1: T)} | \boldsymbol {y}) \right] - \log q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L}) \right] \\ = \mathbb {E} _ {q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L})} \left[ \mathbb {E} _ {q (\boldsymbol {y} ^ {(1: T)} | \boldsymbol {y})} \left[ \sum_ {t = 1} ^ {T} \log \frac {p _ {\boldsymbol {\theta}} (\boldsymbol {y} ^ {(t - 1)} | G , \boldsymbol {y} ^ {(t)})}{q (\boldsymbol {y} ^ {(t)} | \boldsymbol {y} ^ {(t - 1)})} + \log p (\boldsymbol {y} ^ {(T)}) \right] - \log q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L}) \right] \\ = \mathbb {E} _ {q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L})} \left[ \mathbb {E} _ {q (\boldsymbol {y} ^ {(1: T)} | \boldsymbol {y})} \left[ \sum_ {t = 1} ^ {T} \log \frac {p _ {\boldsymbol {\theta}} (\boldsymbol {y} ^ {(t - 1)} | G , \boldsymbol {y} ^ {(t)})}{q (\boldsymbol {y} ^ {(t - 1)} | \boldsymbol {y} ^ {(t)} , \boldsymbol {y})} \frac {q (\boldsymbol {y} ^ {(t - 1)} | \boldsymbol {y})}{q (\boldsymbol {y} ^ {(t)} | \boldsymbol {y})} + \log p (\boldsymbol {y} ^ {(T)}) \right] - \log q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L}) \right] \\ = \mathbb {E} _ {q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L})} \left[ \mathbb {E} _ {q (\boldsymbol {y} ^ {(1: T)} | \boldsymbol {y})} \left[ \sum_ {t = 1} ^ {T} \log \frac {p _ {\boldsymbol {\theta}} (\boldsymbol {y} ^ {(t - 1)} | G , \boldsymbol {y} ^ {(t)})}{q (\boldsymbol {y} ^ {(t - 1)} | \boldsymbol {y} ^ {(t)} , \boldsymbol {y})} + \log \frac {p (\boldsymbol {y} ^ {(T)})}{q (\boldsymbol {y} ^ {(T)} | \boldsymbol {y})} \right] - \log q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L}) \right] \\ = \mathbb {E} _ {q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L})} \left[ \sum_ {t = 1} ^ {T} \mathbb {E} _ {q (\boldsymbol {y} ^ {(1: T)} | \boldsymbol {y})} \left[ \log \frac {p _ {\boldsymbol {\theta}} (\boldsymbol {y} ^ {(t - 1)} | G , \boldsymbol {y} ^ {(t)})}{q (\boldsymbol {y} ^ {(t - 1)} | \boldsymbol {y} ^ {(t)} , \boldsymbol {y})} \right] + \mathbb {E} _ {q (\boldsymbol {y} ^ {(T)} | \boldsymbol {y})} \left[ \log \frac {p (\boldsymbol {y} ^ {(T)})}{q (\boldsymbol {y} ^ {(T)} | \boldsymbol {y})} \right] - \log q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L}) \right] \\ = \mathbb {E} _ {q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L})} \left[ \sum_ {t = 1} ^ {T} \mathcal {L} _ {\mathrm {D P M}} ^ {(t)} + C - \log q (\boldsymbol {y} _ {U} | \boldsymbol {y} _ {L}) \right], \\ \end{array}
$$
where  $\mathcal{L}_{\mathrm{DPM}}^{(t)}$  is a training objective for a  $t$  step, and  $C$  is a constant with respect to the parameters  $\pmb{\theta}$ . Following Ho et al. [16], we simplify  $\mathcal{L}_{\mathrm{DPM}}^{(t)}$  with a residual function  $\epsilon_{\pmb{\theta}}(\pmb{y}^{(t)}, G, t)$  in Equation (2).
$$
\mathcal {L} _ {\mathrm {D P M}} ^ {(t)} = C - \mathbb {E} _ {\boldsymbol {\epsilon} \sim \mathcal {N} (\mathbf {0}, \mathbf {I})} \left[ \frac {\beta_ {t} ^ {2}}{2 \sigma_ {t} ^ {2} \alpha_ {t} (1 - \bar {\alpha} _ {t})} \left\| \boldsymbol {\epsilon} - \boldsymbol {\epsilon} _ {\boldsymbol {\theta}} (\sqrt {\bar {\alpha} _ {t}} \boldsymbol {y} + \sqrt {1 - \bar {\alpha} _ {t}} \boldsymbol {\epsilon}, G, t) \right\| _ {2} ^ {2} \right],
$$
where  $C$  is a constant with respect to the parameters  $\theta$ . An additional suggestion from Ho et al. [16] is to set all weights of the mean squared error to one instead of  $\beta_t^2 / 2\sigma_t^2\alpha_t(1 - \bar{\alpha}_t)$ , and we follow this suggestion in this paper.
A.4 Manifold-constrained sampling
Algorithm 2 Manifold-constrained sampling
1: Input: Graph  $G$ , node attributes  $\mathbf{x}$ , labels  $\mathbf{y}_L$ , and temperature of randomness  $\tau^2$ .  
2: Get  $\mathbf{y}^{(T)} \sim \mathcal{N}(\mathbf{y}^{(T)}; \mathbf{0}, \tau^2 \mathbf{I})$   
3: for  $t = T - 1, \ldots, 0$  do  
4: Get  $\mathbf{z} \sim \mathcal{N}(\mathbf{z}; \mathbf{0}, \tau^2 \mathbf{I})$   
5: Set  $\tilde{\mathbf{y}}^{(t)} \gets \frac{1}{\sqrt{\alpha_t}} \left( \mathbf{y}^{(t+1)} - \frac{\beta_{t+1}}{\sqrt{1 - \bar{\alpha}_{t+1}}} \boldsymbol{\epsilon}_\theta (\mathbf{y}^{(t+1)}, G, t + 1) \right) + \sigma_{t+1} \mathbf{z}$   
6: Set  $\hat{\mathbf{y}}^{(t+1)} \gets \frac{1}{\sqrt{\alpha_{t+1}}} \left( \mathbf{y}^{(t+1)} - \frac{1 - \bar{\alpha}_{t+1}}{\sqrt{1 - \bar{\alpha}_{t+1}}} \boldsymbol{\epsilon}_\theta (\mathbf{y}^{(t+1)}, G, t + 1) \right)$   
7: Set  $\bar{\mathbf{y}}^{(t)} \gets \left( \tilde{\mathbf{y}}^{(t)} - \gamma \frac{\partial}{\partial \mathbf{y}^{(t+1)}} \| \mathbf{y}_L - \hat{\mathbf{y}}_L^{(t+1)} \|_2^2 \right)$   
8: Set  $\mathbf{y}_U^{(t)} \gets \bar{\mathbf{y}}_U^{(t)}$   
9: Get  $\mathbf{z} \sim \mathcal{N}(\mathbf{z}; \mathbf{0}, \tau^2 \mathbf{I})$   
10: Set  $\mathbf{y}_L^{(t)} \gets \sqrt{\bar{\alpha}_t} \mathbf{y}_L + \sqrt{1 - \bar{\alpha}_t} \mathbf{z}$   
11: Set  $\mathbf{y}^{(t)} \gets \mathbf{y}_L^{(t)} \cup \mathbf{y}_U^{(t)}$   
12: end for  
13: return  $\mathbf{y}_U^{(0)}$
To sample  $\pmb{y}_U$  from  $p_{\theta}(\pmb{y}_U|G,\pmb{y}_L)$ , we use a manifold-constrained sampling proposed by Chung et al. [18]. Here, the update rule of the reverse process for  $t = 0,\dots ,T - 1$  is defined as follows:
$$
\tilde {\boldsymbol {y}} ^ {(t)} = \frac {1}{\sqrt {\alpha_ {t}}} \left(\boldsymbol {y} ^ {(t + 1)} - \frac {\beta_ {t + 1}}{\sqrt {1 - \bar {\alpha} _ {t + 1}}} \boldsymbol {\epsilon} _ {\boldsymbol {\theta}} \left(\boldsymbol {y} ^ {(t + 1)}, G, t + 1\right)\right) + \sigma_ {t + 1} \boldsymbol {z}, \tag {3}
$$
$$
\boldsymbol {y} _ {U} ^ {(t)} = \bar {\boldsymbol {y}} _ {U} ^ {(t)}, \quad \bar {\boldsymbol {y}} ^ {(t)} = \tilde {\boldsymbol {y}} ^ {(t)} - \gamma \frac {\partial}{\partial \boldsymbol {y} ^ {(t + 1)}} \left\| \boldsymbol {y} _ {L} - \hat {\boldsymbol {y}} _ {L} ^ {(t + 1)} \right\| _ {2} ^ {2}, \tag {4}
$$
$$
\boldsymbol {y} _ {L} ^ {(t)} = \sqrt {\bar {\alpha} _ {t}} \boldsymbol {y} _ {L} + \sqrt {1 - \bar {\alpha} _ {t}} \boldsymbol {z}, \tag {5}
$$
where  $\mathbf{z}$  is sampled from  $\mathcal{N}(z;0,\tau^2\mathbf{I})$ . The Equation (3) is a temporal reverse diffusion step before applying the manifold-constrained samplings. The Equation (4) applies the manifold-constrained gradient  $\gamma \frac{\partial}{\partial\pmb{y}^{(t + 1)}}\| \pmb{y}_L - \hat{\pmb{y}}_L^{(t + 1)}\| _2^2$ . Here,  $\hat{\pmb{y}}^{(t + 1)}$  is the label estimate in  $t + 1$  steps defined as follows:
$$
\hat {\boldsymbol {y}} ^ {(t + 1)} = \frac {1}{\sqrt {\bar {\alpha} _ {t + 1}}} \left(\boldsymbol {y} ^ {(t + 1)} - \frac {1 - \bar {\alpha} _ {t + 1}}{\sqrt {1 - \bar {\alpha} _ {t + 1}}} \boldsymbol {\epsilon} _ {\boldsymbol {\theta}} (\boldsymbol {y} ^ {(t + 1)}, G, t + 1)\right),
$$
where  $\gamma$  is a hyper-parameter. We set  $\gamma$  to  $1 / \| \pmb{y}_L - \hat{\pmb{y}}_L^{(t + 1)}\| _2^2$ . The Equation (5) is a projection step. Additionally, we introduce a parameter  $\tau$  to control the randomness; when we set  $\tau$  to zero, the modified reverse step becomes deterministic. This allows us to control the randomness in obtaining samples. We describe the detailed sampling algorithm in Algorithm 2.
B WL test and GNN's expressiveness
In this section, we provide the proof of Theorem 1 in detail.
B.1 Preliminaries
Algorithm 3 1-dimensional Weisfeiler-Lehman algorithm
1: Input: Graph  $G = (\mathcal{V}^G, \mathcal{E}^G, \boldsymbol{x}^G)$  and the number of iterations  $T$ .  
2: Output: Color mapping  $\chi_G : \mathcal{V}^G \to \mathcal{C}$ .  
3: Initialize: Let  $\chi_G^0(v) \gets \text{hash}(x_v)$  for  $v \in \mathcal{V}^G$ .  
4: for  $t = 1, \ldots, T$  do  
5: for each  $v \in \mathcal{V}^G$  do  
6: Set  $\chi_G^t(v) \gets \text{hash}(\chi_G^{t-1}(v), \{\{\chi_G^{t-1}(u) : u \in \mathcal{N}_G(v)\}\})$   
7: end for  
8: end for  
9: Return:  $\chi_G^T$
We denote set as  $\{\}$ , and multiset as  $\{\{\}\}$ , which is a set allowing duplicate elements. We represent the cardinality of set or multiset  $S$  as  $|\mathcal{S}|$ . A graph is denoted as  $G = (\mathcal{V}^G, \mathcal{E}^G, \boldsymbol{x}^G)$ , where  $\mathcal{V}^G$  stands for the set of nodes,  $\mathcal{E}^G$  for the set of edges, and  $\boldsymbol{x}^G = \{x_v^G : v \in \mathcal{V}^G\}$  for the node-wise attributes. We also use  $\mathcal{N}_G(v) := \{u \in \mathcal{V}^G : \{v, u\} \in \mathcal{E}^G\}$  to denote the set of neighbor nodes of node  $v$  in graph  $G$ . We abbreviate a set of integers using the notation  $[m] := \{0, \dots, m\}$ . We also assume the hash functions are all injective and denote them by  $\mathrm{hash}(\cdot)$ .
Definition 1 (Vertex coloring). Vertex coloring  $\chi_G(v)$  is an injective hash function that maps a vertex  $v$  in graph  $G$  to a color  $c$  from an abstract color set  $\mathcal{C}$ . Next, with a slight abuse of notation, we let the graph color  $\chi_G$  indicate the multiset of node colors in graph  $G$ , i.e.,  $\chi_G := \{\{\chi_G(v): v \in \mathcal{V}^G\}\}$ .
Now we explain the 1-dimensional Weisfeiler-Lehman (1-WL) algorithm [19], a classical algorithm to distinguish non-isomorphic graphs. At a high level, the 1-WL test iteratively updates node colors based on their neighbors until a stable coloring is reached. To be specific, given a vertex  $v$ , the initial node color  $\chi_G^0(v)$  is set using an injective hash function on the node attribute  $x_v$ . At each iteration, each node color is refined based on the aggregation of neighbor node colors, e.g.,  $\{\chi_G^{t-1}(u) : u \in \mathcal{N}_G(v)\}$ . At the final  $T$ -th iteration, the algorithm returns the graph color  $\chi_G^T$ . We provide a detailed description in Algorithm 3.
The 1-WL test allows to compare a pair of graphs  $G, H$  using the refined graph colors  $\chi_G^T, \chi_H^T$ . If the two graph colors  $\chi_G^T$  and  $\chi_H^T$  are not equivalent, two graphs  $G, H$  are guaranteed to be nonisomorphic. Otherwise, the test is inconclusive, i.e., the two graphs  $G, H$  are possibly isomorphic. We provide an example run of the 1-WL test in Figure 8(a), where the test is inconclusive.
Related to the 1-WL algorithm, we first prove a characteristic of it that will be later used in our proof.
Lemma 1. Consider running 1-WL on two graphs  $G = (\mathcal{V}^G, \mathcal{E}^G, \boldsymbol{x}^G)$  and  $H = (\mathcal{V}^H, \mathcal{E}^H, \boldsymbol{x}^H)$ . If the initial graph colors of two graphs are distinct, i.e.,  $\chi_G^0 \neq \chi_H^0$ , the respective outputs of the 1-WL are also distinct, i.e.,  $\chi_G^T \neq \chi_H^T$ .
Proof. We first prove that  $\chi_G^t \neq \chi_H^t$  is satisfied when  $\chi_G^{t-1} \neq \chi_H^{t-1}$ . The rest of the proof is straightforward by induction on  $t$ . To be specific, the 1-WL updates  $\chi_G^t$  and  $\chi_H^t$  as follows:
$$
\chi_ {G} ^ {t} = \{\left\{\operatorname {h a s h} \left(\chi_ {G} ^ {t - 1} (v), \left\{\chi_ {G} ^ {t - 1} (u): u \in \mathcal {N} _ {G} (v) \right\}\right): v \in \mathcal {V} ^ {G} \right\},
$$
$$
\chi_ {H} ^ {t} = \left\{ \right.\left. \right.\left\{ \right.\operatorname {h a s h} \left( \right.\chi_ {H} ^ {t - 1} (v), \left\{ \right.\left\{\chi_ {H} ^ {t - 1} (u): u \in \mathcal {N} _ {H} (v) \right\}\left. \right) : v \in \mathcal {V} ^ {H} \left. \right.\left. \right.\left. \right\}.
$$
Since  $\mathrm{hash}(\cdot)$  is an injective function,  $\chi_G^t$  and  $\chi_H^t$  are distinct for  $\chi_G^{t - 1}\neq \chi_H^{t - 1}$ .
B.2 AGG-WL test
Next, we describe the newly proposed AGG-WL test, which is an analog of our DPM-SNC. Similar to the 1-WL test, our AGG-WL test assigns node colors and iteratively refines them based on the
neighbor node colors. However, AGG-WL generates augmented graphs at initialization, creating multiple "views" on the graph with diverse initial vertex coloring. Then it applies 1-WL on each of the augmented graphs to obtain the augmented graph colors. Finally, AGG-WL aggregates the augmented graphs colors.
The complete algorithm is described in Algorithm 4. Given a graph  $G = (\mathcal{V}^G, \mathcal{E}^G, \boldsymbol{x}^G)$  and the set of possible node-wise augmentations  $\mathcal{Z}^G$ , the algorithm generates augmented graph  $G^{m}$ . In particular, the augmented graph  $G^{m}$  is defined as follows:
$$
G ^ {m} = \left(\mathcal {V} ^ {G}, \mathcal {E} ^ {G}, \tilde {\boldsymbol {x}} ^ {G ^ {m}}\right), \quad \tilde {\boldsymbol {x}} ^ {G ^ {m}} = \left\{\left(x _ {v} ^ {G} \| z _ {v} ^ {m}\right): v \in \mathcal {V} \right\}, \quad \boldsymbol {z} ^ {m} = \left\{z _ {v} ^ {m}: v \in \mathcal {V} \right\}, \quad \boldsymbol {z} ^ {m} \in \mathcal {Z} ^ {G}
$$
where  $z^{m}$  is an augmentation method in  $\mathcal{Z}^G$ . The algorithm creates an augmented graph  $G^{m}$  by node-wise concatenation of  $z^{m}$  to its node attributes  $\pmb{x}^{G}$ , where  $\tilde{\pmb{x}}^{G^m}$  denotes the node-wise attribute augmented by  $z^{m}$ . The symbol  $\cdot \|\cdot$  denotes the concatenation of two elements. Furthermore,  $z^0$  is a unique token, where  $G^0$  has the same information as  $G$ . Also, with a slight abuse of notation, we let the graph color returned by the AGG-WL  $\chi_G^{\mathrm{AGG}}$  indicate the multiset of augmented graph colors, i.e.,  $\chi_{G}^{\mathrm{AGG}} := \{\{\chi_{G^{m}}^{T}: m \in [|Z^{G}]|\}\}$ . Here,  $|\mathcal{Z}^G|$  denotes the cardinality of set  $\mathcal{Z}^G$ .
Algorithm 4 Aggregation Weisfeiler-Lehman algorithm
1: Input: Graph  $G = (\mathcal{V}^G, \mathcal{E}^G, \boldsymbol{x}^G)$ , the number of iterations  $T$ , and the augmentation set  $\mathcal{Z}^G$ .  
2: Output: Color mapping  $\chi_G : \mathcal{V}^G \to \mathcal{C}$ .  
3: Initialize: Generate  $|\mathcal{Z}^G|$  augmented graphs  $G^m = (\mathcal{V}^G, \mathcal{E}^G, \boldsymbol{x}^{G^m})$  where  $\tilde{\boldsymbol{x}}^{G^m} = \{(x_v^G \| z_v^m) : v \in \mathcal{V}^G\}$  for  $z^m \in \mathcal{Z}^G$ . Let  $\chi_{G^m}^0(v) \gets \text{hash}(\tilde{\boldsymbol{x}}_v^{G^m})$  for  $v \in \mathcal{V}^G$ ,  $m \in [| \mathcal{Z}^G|]$ .  
4: for  $t = 1, \ldots, T$  do  
5: for each  $v \in V$  do  
6: for  $m \in [| \mathcal{Z}^G|]$  do  
7:  $\chi_{G^m}^t(v) \gets \text{hash}(\chi_{G^m}^{t-1}(v), \{\chi_{G^m}^{t-1}(u) : u \in \mathcal{N}_{G^m}(v)\})$   
8: end for  
9: end for  
10: end for  
11:  $\chi_{G}^{\mathrm{AGG}} \gets \{\chi_{G^m}^T : m \in [| \mathcal{Z}^G|]\}$   
12: Return:  $\chi_{G}^{\mathrm{AGG}}$
Our contribution is establishing the connection between the AGG-WL and the DPM-SNC: aggregation of refined graph colors for augmented graphs and marginalization over latent variables. More detail between DPM-SNC and the AGG-WL test is described in Appendix B.3.2.
We note that our algorithm bears some similarities with several previous studies. The DS-WL and DSS-WL test [41] defined a modified WL-test based on the set of subgraphs (instead of augmented graphs) with a modified edge set  $\mathcal{E}'$ . Hang et al. [13] uses a collective algorithm to consider pseudolabels as additional inputs to boost GNN expressiveness. However, they rely on the assumption that one can find an "optimal" pseudo-label to discriminate a pair of graphs, which may be hard to realize in practice. Zhang et al. [60] proposed labeling tricks that learn to capture dependence between nodes also by adding additional features, but mainly focus on link prediction in inductive settings.
Let us start by restating Theorem 1.
Proof. We divide the proof into two parts, Appendices B.3.1 and B.3.2. In Appendix B.3.1, we show that the AGG-WL test is strictly more powerful than the 1-WL test. In Appendix B.3.2, we show that DPM-SNC using a 1-WL-GNN is as powerful as the AGG-WL test. From these two proofs, one can conclude that DPM-SNC using a 1-WL-GNN is strictly more powerful than 1-WL-GNNs in distinguishing non-isomorphic graphs.  $\square$
B.3.1 Expressiveness of AGG-WL test
In this subsection, we prove that the AGG-WL test is strictly more powerful than the 1-WL test. The high-level idea is showing that (i) the AGG-WL test is at least as powerful as the 1-WL test, i.e., any two graphs distinguishable by the 1-WL test is also distinguishable by the AGG-WL test, and (ii) there exist two non-isomorphic graphs that are indistinguishable by the 1-WL test, but distinguishable by our AGG-WL test.
Lemma 2. AGG-WL is at least as powerful as WL in distinguishing non-isomorphic graphs, i.e., any two non-isomorphic graphs  $G, H$  distinguishable by WL are also distinguishable by AGG-WL.
Proof. We first note that both 1-WL and AGG-WL can discriminate two non-isomorphic graphs with distinct sizes in a straightforward way. Therefore, we focus on non-trivial cases for two graphs  $G = (\mathcal{V}^G, \mathcal{E}^G, \boldsymbol{x}^G)$ ,  $H = (\mathcal{V}^H, \mathcal{E}^H, \boldsymbol{x}^H)$  where the number of nodes and edges are same, i.e.,  $|\mathcal{V}^G| = |\mathcal{V}^H|, |\mathcal{E}^G| = |\mathcal{E}^H|$ . The cardinality of the augmentation sets are also the same,  $|\mathcal{Z}^G| = |\mathcal{Z}^H|$ . We prove if these two graphs  $G, H$  are distinguishable by the 1-WL, two graphs  $G, H$  are also distinguishable by the AGG-WL, i.e., if  $\chi_G^T \neq \chi_H^T$ , then  $\chi_G^{\mathrm{AGG}} \neq \chi_H^{\mathrm{AGG}}$ .
First, since  $\chi_G^T \neq \chi_H^T$  is satisfied,  $\chi_{G^0}^T \neq \chi_{H^0}^T$  is trivial. Additionally, the initial graphs colors of graph augmented by the unique token and any other augmentations are distinct, i.e.,  $\chi_{G^0}^0 \neq \chi_{H^1}^0, \ldots, \chi_{H^{|z^H|}}^0$ . This implies  $\chi_{G^0}^T \notin \chi_H^{\mathrm{AGG}} = \{\chi_{H^0}^T, \ldots, \chi_{H^{|z^H|}}^T\}$  according to Lemma 1 which shows distinct initial graph colors produce the distinct outputs of 1-WL. Since  $\chi_{G^0}^T \in \chi_G^{\mathrm{AGG}}$ , one can see that  $\chi_{G}^{\mathrm{AGG}} \neq \chi_{H}^{\mathrm{AGG}}$ .
Next, we show that there exist two graphs  $G$ ,  $H$  indistinguishable by the 1-WL test, but distinguishable by the AGG-WL test. We first show a simple example from Figure 8(a), then for a family of circular skip link (CSL) graphs [61].
Lemma 3. Two graphs  $A, B$  are indistinguishable by the 1-WL, but distinguishable by the AGG-WL.
Proof. In Figure 8(a), two graphs  $A, B$  with 6 nodes are given. Here, the graph color computed by the 1-WL are equivalent, i.e.,  $\chi_G^T = \chi_H^T$  (two light blue nodes and four gray nodes). Therefore, two graphs  $A, B$  are indistinguishable by the 1-WL.
Next, we prove that two graphs  $A, B$  are distinguishable by the AGG-WL. Since Lemma 1 shows 1-WL always computes a distinct graph color from distinct initial graph colors, we only need to show that there exists augmented graphs  $A^m, B^m$  with the same initial color has distinct graph colors computed by the 1-WL. Then, for two augmented graph sets  $\{A^m : m \in [|Z^{\mathcal{V}^A}]|\}$  and  $\{B^m : m \in [|Z^{\mathcal{V}^B}]|\}$ , we can show that the graph colors computed by AGG-WL are always distinct.
We show a simple case, where augmented graphs with non-zero augmentation on one node results distinct graph colors computed by the AGG-WL. We denote "single binary node feature augmentation",  $z^i$  as the case where node  $i$  is non-zero augmented and other nodes are augmented with zero, e.g.,  $z^3 = \{0, 0, 1, 0, 0, 0\}$ . In Figure 8(b) the non-zero augmented node  $i$  is colored in blue, and others are colored transparent. After applying the 1-WL algorithm on each graph, we obtain the graph color  $\chi_{A^i}^T$  and  $\chi_{B^i}^T$  for  $i \in \{1, \ldots, 6\}$ . It is clear that  $\{\{\chi_{A^i}^T : i \in \{1, \ldots, 6\}\} \} \neq \{\{\chi_{B^i}^T : i \in \{1, \ldots, 6\}\} \}$ , thereby two graphs  $A, B$  are distinguishable by AGG-WL.
Circular skip link graphs. CSL graph is denoted as  $\mathrm{CSL}(n,r)$ , where  $n$  is the number of nodes, i.e.,  $V = \{0,\dots ,n - 1\}$  and  $r$  is the skip connection length. For  $n$  and  $r$ ,  $r < n - 1$  must hold. There exists  $2n$  edges, between node  $i$  and  $(i + 1)$  mod  $n$  forming a cycle, and between nodes  $i$  and  $(i + r)$  mod  $n$  forming a skip link for  $i\in \{0,\ldots ,n - 1\}$ .
Lemma 4. For  $n \geq 8, r \in [3, n/2 - 1]$ , two graphs  $CSL(n, 2)$  and  $CSL(n, r)$  are indistinguishable by the 1-WL, but distinguishable by the AGG-WL.
Proof. We first give a brief proof of why  $\mathrm{CSL}(n,2)$  and  $\mathrm{CSL}(n,r)$  are indistinguishable by the 1-WL, then prove they are distinguishable by AGG-WL. We consider a non-trivial case where the node attributes are all same, i.e.,  $\chi_G^0 (v) = \chi_G^0 (u)$  for  $v,u\in \mathcal{V}^{G}$ .
Let the initial color be  $c_{0}$  for all the nodes for graphs  $\mathrm{CSL}(n,2)$  and  $\mathrm{CSL}(n,r)$ ,  $\chi_G^0 (v) = c_0\forall v\in \mathcal{V}^G,G\in \{\mathrm{CSL}(n,2),\mathrm{CSL}(n,r)\}$ . Then, for all nodes  $v\in \mathcal{V}^{G}$ , the color refinement process can
(a) 1-WL
(b) AGG-WL
Figure 8: An example of two non-isomorphic graphs  $A, B$  indistinguishable by 1-WL, but distinguishable by the AGG-WL. (a) Two graphs  $A, B$  are associated with nodes sharing the same attribute. The 1-WL computes the same graph color for  $A, B$ , i.e.,  $\chi_A^T = \chi_B^T$ . (b) Applying AGG-WL with "single binary node feature augmentation" to the graphs  $A$  and  $B$ . This results in six cases for each graph. One can see that the multiset of augmented graph colors are different, i.e.,  $\{\chi_{A^i}^T : i \in \{1, \ldots, 6\}\} \neq \{\chi_{B^i}^T : i \in \{1, \ldots, 6\}\}$ , thereby two graphs  $A, B$  are distinguishable by the AGG-WL.
than be written as  $\chi_G^1 (v) = \mathrm{hash}(c_0,\{\{c_0,c_0,c_0,c_0\}\})$ . All nodes have identical colors  $c_{1}$  for both graphs  $\operatorname {CSL}(n,2)$ ,  $\operatorname {CSL}(n,r)$ , refining the initial graph color. Therefore  $\operatorname {CSL}(n,2)$  and  $\operatorname {CSL}(n,r)$  are indistinguishable by the 1-WL.
Now we consider the AGG-WL. Again, we only need to show that there exists augmented graphs  $A^i, B^j$  with the same initial color but having distinct graph colors refined by the 1-WL. Given symmetry, there is only one case of an augmented graph that adds a non-zero augmentation to one node. We use  $v_i$  to denote the  $i$ -th node in graph and denote the augmented node as  $v_0$ . We denote each augmented graph as  $\mathrm{CSL}(n,2)^1, \mathrm{CSL}(n,r)^1$ , and let the initial color  $\chi_G^0(v_0) = c_1$ , and  $\forall i \in \{1, \ldots, n-1\}, \forall G \in \{\mathrm{CSL}(n,2)^1, \mathrm{CSL}(n,r)^1\} \chi_G^0(v_i) = c_0$ .
Iteration 1. We focus on the four nodes connected to  $v_{0}$ . Since  $r \in [3, n/2 - 1]$ , two node  $v_{n-r}, v_{r}$  are distinct. The color refinement can be written as following:
- For  $v \in \{v_1, v_2, v_{n-1}, v_{n-2}\}$ ,  $\chi_{\mathrm{CSL}(n,2)^1}^1(v) = \mathrm{hash}(c_0, \{c_1, c_0, c_0, c_0\}) = c_2$ .
- For  $v \notin \{v_{1}, v_{2}, v_{n - 1}, v_{n - 2}\}$ ,  $\chi_{\mathrm{CSL}(n,2)^{1}}^{1}(v) = \mathrm{hash}(c_{0}, \{c_{0}, c_{0}, c_{0}, c_{0}\}) = c_{3}$ .
- For  $v \in \{v_1, v_r, v_{n-1}, v_{n-r}\}$ ,  $\chi_{\mathrm{CSL}(n,r)^1}^1(v) = \mathrm{hash}(c_0, \{c_1, c_0, c_0, c_0\})$  =  $c_2$ .
- For  $v \notin \{v_1, v_r, v_{n-1}, v_{n-r}\}$ ,  $\chi_{\mathrm{CSL}(n,r)^1}^1(v) = \mathrm{hash}(c_0, \{c_0, c_0, c_0, c_0\})$  =  $c_3$ .
(a) Original graph and 1-WL coloring
(b) AGG-WL coloring on  $\mathrm{CSL}(8,2)^1$ ,  $\mathrm{CSL}(8,3)^1$
Figure 9: Example on two non-attribute CSL(8,2), CSL(8,3) graph. (a) 1-WL returns same graph color for two graphs, i.e.,  $\chi_{\mathrm{CSL}(8,2)}^T = \chi_{\mathrm{CSL}(8,3)}^T$ . (b) Considering symmetry, there is only one type of augmented graph that with non-zero augmentation on one node, we denoted the augmented graph as  $\mathrm{CSL}(8,2)^1$  and  $\mathrm{CSL}(8,3)^1$ . In the process of 1-WL assigning colors to augmented graph, they result different graph color from the second iteration, and one can conclude  $\chi_{\mathrm{CSL}(8,2)^1}^T \neq \chi_{\mathrm{CSL}(8,3)^1}^T$ . Therefore  $\mathrm{CSL}(8,2)$ ,  $\mathrm{CSL}(8,3)$  are distinguishable by the AGG-WL.