SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lecheng Kong $^{1}$  Jiarui Feng $^{1}$  Hao Liu $^{1}$  Dacheng Tao $^{2}$
Yixin Chen<sup>1</sup> Muhan Zhang<sup>3</sup>
{jerry.kong, feng.jiarui, liuhao, ychen25}@wustl.edu,
dacheng.tao@gmail.com, muhan@pku.edu.cn
While Graph Neural Networks (GNNs) recently became powerful tools in graph learning tasks, considerable efforts have been spent on improving GNNs' structural encoding ability. A particular line of work proposed subgraph GNNs that use subgraph information to improve GNNs' expressivity and achieved great success. However, such effectivity sacrifices the efficiency of GNNs by enumerating all possible subgraphs. In this paper, we analyze the necessity of complete subgraph enumeration and show that a model can achieve a comparable level of expressivity by considering a small subset of the subgraphs. We then formulate the identification of the optimal subset as a combinatorial optimization problem and propose Magnetic Graph Neural Network (MAG-GNN), a reinforcement learning (RL) boosted GNN, to solve the problem. Starting with a candidate subgraph set, MAG-GNN employs an RL agent to iteratively update the subgraphs to locate the most expressive set for prediction. This reduces the exponential complexity of subgraph enumeration to the constant complexity of a subgraph search algorithm while keeping good expressivity. We conduct extensive experiments on many datasets, showing that MAG-GNN achieves competitive performance to state-of-the-art methods and even outperforms many subgraph GNNs. We also demonstrate that MAG-GNN effectively reduces the running time of subgraph GNNs.
Recent advances in Graph Neural Networks (GNNs) greatly assist the rapid development of many areas, including drug discovery [2], recommender systems [31], and autonomous driving [6]. The power of GNNs has primarily been attributed to their Message-Passing Paradigm [12]. Message-Passing Paradigm simulates a 1-dimensional Weisfeiler-Lehman (1-WL) algorithm for graph isomorphism testing. Such a simulation allows GNN to encode rich structural information. In many fields, structural information is crucial to determine the properties of a graph.
However, as Xu et al. [32] pointed out, GNN's structure encoding capability, or its expressivity, is also upper-bounded by the 1-WL test. Specifically, a message-passing neural network (MPNN) cannot recognize many substructures like cycles and paths and fails to properly learn and distinguish regular graphs. Meanwhile, these substructures are significant in areas including chemistry and biology. To overcome this limitation, considerable effort was spent on investigating more-expressive GNNs. A famous line of work is subgraph GNNs [33, 34, 37]. Subgraph GNNs extract rooted subgraphs around every node in the graph and apply MPNN onto the subgraphs to obtain subgraph representations. The subgraph representations are summarized to form the final representation of the graph. Such an approach is theoretically proved to be more expressive than MPNN and achieved superior empirical results. Later work found that subgraph GNNs are still bounded by the 3-dimensional WL test (3-WL) [10].
Figure 1: Comparison of two simple graphs.
This spurs the research on extensions to edge-rooted subgraph GNNs with better expressivity [15]. In fact, by linking subgraph GNNs to the  $k$ -dimensional WL ( $k$ -WL) test hierarchy, we observe that for subgraph GNNs to exceed the limits of k-WL expressivity, it needs to enumerate and apply MPNN to subgraphs exponential to  $k$  [26, 10]. Hence, the higher expressivity, unfortunately, comes with a higher computational cost, which makes it challenging to apply subgraph GNNs to even moderately sized graphs (e.g.,  $\sim 100$  nodes).
This makes us wonder, is it necessary for subgraph GNNs to enumerate all rooted subgraphs to achieve
higher expressivity? For example, an MPNN fails to distinguish graphs A and B in Figure 1, as they are 2-regular graphs with identical 1-hop subtrees. Meanwhile, a subgraph GNN will see different subgraphs around nodes in the two graphs. These subgraphs are distinguishable by MPNN, allowing a subgraph GNN to differentiate between graphs A and B. However, we can observe that many subgraphs from the same graph are identical. Specifically, graph A has two types of subgraphs, while graph B only has triangle subgraphs. As a result, locating a non-triangle subgraph in the top graph enables us to run MPNN once on it to discern the difference between the two graphs. On the contrary, a subgraph GNN takes eight extra MPNN runs for the remaining nodes. This graph pair shows that we can obtain discriminating power equal to that of a subgraph GNN without enumerating all subgraphs. We also include advanced examples with more complex structures in Section 3.
Therefore, we propose Magnetic Graph Neural Network (MAG-GNN), a reinforcement learning (RL) based method, to leverage this property and locate the discriminative subgraphs effectively. Specifically, we start with a candidate set of subgraphs randomly selected from all rooted subgraphs. The root node features of each subgraph are initialized uniquely. In every step, each target subgraph in the candidate set is substituted by a new subgraph with more distinguishing power. MAG-GNN achieves this by mapping each target subgraph to a Q-Table, representing the expected reward of replacing the target subgraph with another potential subgraph. It then selects the subgraph that maximizes the reward. MAG-GNN repeats the process until it identifies the set of subgraphs with the highest distinguishing power. The resulting subgraph set is then passed to a prediction GNN for downstream tasks. MAG-GNN reduces subgraph GNN's exponentially complex enumeration procedure to an RL searching process with constant steps. This potently constrains the computational cost while keeping the expressivity.
We conduct extensive experiments on synthetic and real-world graph datasets and show that MAG-GNN achieves competitive performance to state-of-the-art (SOTA) methods and even outperforms subgraph GNNs on many datasets with a shorter runtime. Our work shows that partial subgraph information is sufficient for good expressivity, and MAG-GNN smartly locates expressive subgraphs and achieves the same goal with better efficiency.
A graph can be represented as  $G = \{V, E, X\}$ , where  $V$  is the set of nodes and  $E \subseteq V \times V$  is the set of edges. Let  $V(G)$  and  $E(G)$  represent the node and edge sets of  $G$ , respectively. Nodes are associated with features  $X = \{\pmb{x}_v | \forall v \in V\}$ . An MPNN  $g$  can be decomposed into  $T$  layers of COMBINE and AGGREGATE functions. Each layer uses the COMBINE function to update the current node embedding from its previous embedding and the AGGREGATE function to process the node's neighbor embeddings. Formally,
$$
\boldsymbol {m} _ {v} ^ {(t)} = \text {A G G R E G A T E} ^ {(t)} \left(\left\{\left\{\boldsymbol {h} _ {u} ^ {(t - 1)}, u \in \mathcal {N} (v) \right\} \right\}\right), \quad \boldsymbol {h} _ {v} ^ {(t)} = \text {C O M B I N E} ^ {(t)} \left(\boldsymbol {m} _ {v} ^ {(t)}, \boldsymbol {h} _ {v} ^ {(t - 1)}\right) \tag {1}
$$
where  $\pmb{h}_v^{(t)}$  is the node representation after  $t$  iterations,  $\pmb{h}_v^{(0)} = \pmb{x}_v, \mathcal{N}(v)$  is the set of direct neighbors of  $v$ ,  $\pmb{m}_u^{(t)}$  is the message embedding.  $\pmb{h}_v^{(T)}$  is used to form node, edge, and graph-level representations. We use  $H = g(G)$  to denote the generated node embeddings of MPNN. MPNN's variants differ mainly by their AGGREGATE and COMBINE functions but are all bounded by 1-WL in expressivity.
This paper adopts the following formulation for subgraph GNNs. For a graph  $G$ , a subgraph GNN first enumerates all  $k$ -order node tuples  $\{\pmb{v} | \pmb{v} \in V^k(G)\}$  and creates  $|V^k(G)|$  copies of the graph. A
graph associated with node tuple  $\pmb{v}$  is represented by  $G(\pmb{v})$ . Node-rooted subgraph GNNs adopt a 1-order policy and have  $O(V(G))$  graphs; edge-rooted subgraph GNNs adopt a 2-order policy and have  $O(V^{2}(G))$  graphs. Note that we are not taking exact subgraphs here, so we need to mark the node tuples on the copied graphs to maintain the subgraph effect. Specifically,
$$
[ X (\boldsymbol {v}) ] _ {l, p} = \left\{ \begin{array}{l l} c ^ {+} & \text {i f} v _ {l} = [ \boldsymbol {v} ] _ {j} \text {a n d} p = j \\ c ^ {-} & \text {o t h e r w i s e} \end{array} \quad \boldsymbol {v} \in V ^ {k} (G), \right. \tag {2}
$$
$X(\pmb{v}) \in \mathbb{R}^{|V| \times k}$  and  $G(\pmb{v}) = \{V, E, X \oplus X(\pmb{v})\}$ , where  $\oplus$  means row-wise concatenation. We use square brackets to index into a sequence. All entries in  $X(\pmb{v})$  are labeled as  $c^{-}$  except those appearing at corresponding positions of the node tuple. An MPNN  $g$  is applied to every graph, and we use a pooling function to obtain the collective embedding of the graph:
$$
f _ {s} (G) = R ^ {(P)} \left(\left\{g \left(G (\boldsymbol {v})\right) \mid \forall \boldsymbol {v} \in V ^ {k} (G) \right\}\right), \quad P \in \{G, N \}. \tag {3}
$$
$f_{s}(G)$  can be a vector of graph representation if  $R^{(P)}$  is a graph-level pooling function ( $P$  equals  $G$ ) or a matrix of node representations if  $R^{(P)}$  is node-level ( $P$  equals  $N$ ). This essentially implements  $k$ -dimensional ordered subgraph GNN ( $k$ -OSAN) defined in [26] and captures most of the popular subgraph GNNs, including NGNN[35] and ID-GNN[33]. Furthermore, the expressivity of subgraph GNNs increases with larger values of  $k$ . Since node tuples, node-marked graphs, and subgraphs refer to the same thing, we use these terms interchangeably.
The Weisfeiler-Lehman hierarchy (WL hierarchy). The  $k$ -dimensional Weisfeiler-Lehman algorithm is vital in graph isomorphism testing. Earlier work established the hierarchy where the  $(k + 1)$ -WL test is more expressive than the  $k$ -WL test [4]. Xu et al. [32] and Morris et al. [23] connected GNN expressivity to the WL test and proved that MPNN is bounded by the 1-WL test. Later work discovered that all node-rooted subgraphs are bounded by 3-WL expressivity [10], which cannot identify strongly regular graphs and structures like 4-cycles. Qian et al. [26] introduced  $k$ -dimensional ordered-subgraph WL ( $k$ -OSWL) hierarchy which is comparable to the  $k$ -WL test.
Deep Q-learning (DQN). DQN [19] is a robust RL framework that uses a deep neural network to approximate the Q-values, representing the expected rewards for a specific action in a given state. Accumulating sufficient experience with the environment, DQN can make decisions in intricate state spaces. For a detailed introduction to DQN, please refer to Appendix B.
Figure 1 shows that while an MPNN encodes rooted subtrees, a subgraph GNN encodes rooted subgraphs around each node. This allows subgraph GNN to differentiate more graphs at the cost of  $O(|V|)$  MPNN runs. Meanwhile, subgraph GNNs are still bounded by 3-WL tests. Hence, we may need at least 2-node-tuple-rooted (e.g., edge-rooted) subgraph GNNs, requiring  $O(|V^2|)$  MPNN runs to obtain better expressivity. In fact, the required complexity of subgraph GNNs to break beyond  $k$ -WL expressivity grows exponentially with  $k$ . The high computational cost of high expressivity models prevents them from being widely applied to real-world datasets. A natural question is, can we consider only a small subset of all the subgraphs to
Figure 2: Sorted scores.
obtain similar expressivity, just like in Figure 1, where one subgraph is as powerful as the collective information of all subgraphs? This leads us to the following experiment.
We focus on the SR25 dataset. It contains 15 different strongly-regular graphs with the same configuration, each of 25 nodes. The goal is to do multi-class classification to distinguish all pairs of graphs. Since node-rooted subgraph GNNs are upper-bounded by 3-WL and 3-WL cannot distinguish any strongly regular graphs with the same configuration, node-rooted subgraph GNNs will generate identical representations for the 15 graphs while performing 25 MPNN runs each. 2-node-tuple subgraph GNNs have expressivity beyond 3-WL and can distinguish any pair of graphs from the dataset, but it takes 625 MPNN runs.
To test if every subgraph is required, we train an MPNN on randomly sampled 2-node-marked graphs to minimize the expected loss to label  $y$  of the unmarked graph  $G$ ,
$$
\min  _ {g _ {p}} \mathbb {E} _ {\boldsymbol {v}} [ \mathcal {L} (\operatorname {M L P} (f _ {r} (G, \boldsymbol {v})), y) ], \quad f _ {r} (G, \boldsymbol {v}) = R ^ {(G)} \left(g _ {p} (G (\boldsymbol {v}))\right), \quad \boldsymbol {v} \in V ^ {2} (G), \tag {4}
$$
where  $\mathcal{L}$  is the loss function, MLP is a multi-layer perceptron,  $g_{p}$  is an MPNN, and  $R^{(G)}$  pools the node representations to graph representations. Unlike 2-node-tuple-rooted subgraph GNNs that run the MPNN  $|V^2|$  times for each graph, this model runs the MPNN exactly once. During testing, for each of the 15 graphs, we randomly sample one 2-node-marked graph for classification. We perform ten independent tests, and the average test accuracy is  $66.8\%$ . Using 2-node-marked graphs, with only one GNN run, it already outperforms node-rooted subgraph GNNs that fail on this dataset. More interestingly, for each graph  $G$ , we can sort the classification score of its  $|V^2|$  possible node-marked graphs and plot them in Figure 2 (C14 and C15 are the plots for 14-th and 15-th graphs in the dataset). Note that the horizontal axis is not the number of subgraphs; it is the index of subgraphs after sorting by their individual classification scores. We see that each original graph has many marked graphs with a classification score close to one. That means even in one of the most difficult datasets, we still can find particular node-marked graphs that uniquely distinguish the original graph from others. Moreover, unlike the example in Figure 1 with only two types of subgraphs, these marked graphs fall into many different isomorphism groups, meaning that the same observation holds in more complex graphs and can be applied to a wide range of graph classes. We prove that such a phenomenon exists in most regular graphs, which cannot be distinguished by MPNNs (proof in Appendix A).
Theorem 1. Let  $G_{1}$  and  $G_{2}$  be two graphs uniformly sampled from all  $n$ -node,  $r$ -regular graphs where  $3 \leq r < \sqrt{2\log n}$ . Given an injective pooling function  $R^{(G)}$  and an MPNN  $g_{p}$  of I-WL expressivity, with a probability of at least  $1 - o(1)$ , there exists a node (tuple)  $\mathbf{v} \in V(G)$  whose corresponding node marked graph's embedding,  $f_{r}(G_{1}, \mathbf{v})$ , is different from any node marked graph's embedding in  $G_{2}$ .
These observations show that by finding discriminative subgraphs effectively, we only need to apply MPNN to a much smaller subset of the large complete subgraph set to get a close level of expressivity.
4 Magnetic graph neural network
We formulate the problem of finding the most discriminative subgraphs as a combinatorial optimization problem. Given a budget of  $m$  as the number of subgraphs,  $k$  as the order of node tuples, and  $g_{p}$  an MPNN that embeds individual subgraphs, we minimize the following individual loss to graph  $G$ ,
$$
\min  _ {U = \left(\boldsymbol {v} _ {1}, \dots , \boldsymbol {v} _ {m}\right) \in \left(V ^ {k} (G)\right) ^ {m}} \mathcal {L} (\operatorname {M L P} \left(f _ {p} (G, U)\right), \boldsymbol {y}) \tag {5}
$$
$$
f _ {p} (G, U) = R ^ {(P)} \left(\left\{g _ {p} (G (\boldsymbol {v})) | \forall \boldsymbol {v} \in U \right\}\right),
$$
Note that this formulation resembles that of subgraph GNNs in Equation (3); we are only substituting  $V^{k}$  with  $U$  to reduce the number of MPNN runs. Witnessing the great success of deep RL in combinatorial optimization, we adopt Deep Q-learning (DQN) to our setting. We introduce each component of our DQN framework as follows.
State Space: For graph  $G$ , a state is  $m = |U|$  node tuples, their corresponding node-marked graphs, and a  $m$ -by-  $w$  matrix  $W$  to record the state of  $m$  node tuples. Our framework should generalize to arbitrary graphs. Hence, a state  $s$  is defined as,
$$
s = (G, U, W) = (G, \left(\boldsymbol {v} _ {1}, \dots , \boldsymbol {v} _ {m}\right), W), s \in S = \mathcal {G} \times \left(\mathcal {V} ^ {k}\right) ^ {m} \times \left(\mathbb {R} ^ {m \times w}\right) \tag {6}
$$
$S$  is the state space,  $\mathcal{G}$  is the set of all graphs, and  $\mathcal{V}^k$  is the set of all possible  $k$  node tuples of  $\mathcal{G}$ . To generate the initial state during training, we sample one graph  $G$  from the training set and randomly sample  $m$  node tuples from  $V^k(G)$ . The state matrix  $W$  is initialized to 0. The expressivity grows as  $k$  grows. Generally, MAG-GNN with larger  $k$  produces more unique graph embeddings, which is harder to train but might require smaller  $m$  and fewer RL steps to represent the graph, leading to better inference time. However, for some datasets, such expressivity is excessive and poses great challenges to training. A smaller  $k$  can reduce the sample space and stabilize training in this case.
Action Space: We define one RL agent action as selecting one index from one node tuple and replacing the node on that index with another node in the graph. This replaces the node-marked
Figure 3: MAG-GNN's pipeline. An RL agent iteratively updates node tuples for better expressivity.
graph corresponding to the original node tuple with the one corresponding to the modified node tuple. Specifically, an action  $a_{i,j,l}$  on state  $s = (G,U,W)$  does the following on  $U$ :
$$
U ^ {\prime} = a _ {i, j, l} (U) = \left(\boldsymbol {v} _ {1} \dots , \boldsymbol {v} _ {i - 1}, \boldsymbol {v} _ {i} ^ {\prime}, \boldsymbol {v} _ {i + 1}, \dots \boldsymbol {v} _ {m}\right), \boldsymbol {v} _ {i} ^ {\prime} = \left([ \boldsymbol {v} _ {i} ] _ {1}, \dots , [ \boldsymbol {v} _ {i} ] _ {j - 1}, v _ {l}, [ \boldsymbol {v} _ {i} ] _ {j + 1}, \dots [ \boldsymbol {v} _ {i} ] _ {k}\right) \tag {7}
$$
The agent selects a target node tuple  $v_{i}$ , whose  $j$ -th node is replaced with node  $v_{l} \in V$ .  $W$  is then updated by an arbitrary state update function  $W' = f_W(s, U')$  depending on the old state and new node tuples. The update function  $f_{W}$  is not necessarily trainable (e.g., It can simply be a pooling function on embeddings of the marked nodes over states). The next state is  $s' = (G, U', W')$ . The action space is then  $A = [m] \times [k] \times \mathcal{V}$ . Actions only change the node tuple while keeping the graph structure, and the state matrix serves as a tracker of past states and actions. Unlike stochastic RL systems, our RL actions have deterministic outcomes.
The intuition behind the design of the action space is that it limits the number of actions for each node tuple to  $O(m|V|k)$ , which is linear in the number of nodes, and  $k$  is usually small ( $k = 2$  is already more expressive than most subgraph GNNs). We can further reduce the action space to  $O(|V|k)$  if the agent does not need to select the update target but uses a Q-network to do Equation (7) on a given  $v_{i}$ . In such a case, we either sequentially or simultaneously update all node tuples. Since the agent learns to be stationary when a node tuple should not be updated, we do not lose the power of our agent by the reduction. The overall action space design allows efficient computation of Q-values. We include a detailed discussion on the action space and alternative designs in Appendix C.1.
**Reward:** In Section 3, we show that a proper node-marked graph significantly boosts the expressivity. Hence, an optimal reward choice is the increased expressivity from the action. However, expressivity is itself vaguely defined, and we can hardly quantify it. Instead, since the relevant improvement in expressivity should affect the objective value, we choose the objective value improvement as the instant reward. Specifically, let  $s = (G, U, W)$  be the current state and let  $s' = (G, U', W') = a(s)$  be the outcome state of action  $a$ , the reward  $r$  is
$$
r (s, a, s ^ {\prime}) = \mathcal {L} (\operatorname {M L P} \left(f _ {p} (G, U)\right), \boldsymbol {y}) - \mathcal {L} (\operatorname {M L P} \left(f _ {p} (G, U ^ {\prime})\right), \boldsymbol {y}) \tag {8}
$$
This reward correlates the action directly with the objective function, allowing our RL agent to be task-dependent and more flexible for different levels of tasks.
Q-Network: Because our state includes graphs, we require an equivariant Q-network to output consistent Q-tables for actions. Hence, we choose MPNN to parameterize the Q-network. Specifically, for current state  $s = (G,U,W)$  and the target node tuple  $\pmb{v} \in U$ , we have the Q-table as,
$$
[ Q (s, \boldsymbol {v}) ] _ {l} = \operatorname {M L P} ([ g _ {r l} (G (\boldsymbol {v})) ] _ {l} \oplus \sum_ {\boldsymbol {v} \in U} R ^ {(G)} (g _ {r l} (G (\boldsymbol {v}))) \oplus R ^ {(W)} (W)) \tag {9}
$$
Row  $l$  in the Q-table is computed by the embedding of node  $v_{l}$  in the node-marked graph by an MPNN  $g_{rl}$ , the current overall graph representation across all node tuples, and the state matrix  $W$  summarized by a pooling function  $R^{(W)}$ .  $[Q]_{l,j}$  represents the expected reward of replacing the node on index  $j$  of node tuple  $\mathbf{v}$  with node  $v_{l}$ . The best action  $a_{j,l}$  is then chosen by,
$$
\underset {j, l} {\arg \max } [ Q (s, \boldsymbol {v}) ] _ {l, j} \tag {10}
$$
Note that because we assign different initial embeddings based on the node tuple, the MPNN distinguishes otherwise indistinguishable graphs.
As demonstrated in Figure 3, our agent starts with a set of random node tuples and their corresponding subgraphs. In each step, the agent uses an MPNN-parameterized Q-network to update one slot in one node tuple such that the new node tuple set results in a better objective value. The agent repeats for a fixed number of steps  $t$  to find discriminative subgraphs. We do not assign a terminal state during training. Instead, the Q-Network will learn to be stationary when all other action decreases the objective value. This process is like throwing iron balls (marked nodes) into a magnetic field (geometry of the graph) and computing the force that the balls receive along with the interactions among balls (Q-network). We learn to move the balls and reason about the magnetic field's properties. Hence, we dub our method Magnetic GNN (MAG-GNN). To show the effectiveness of our method in locating discriminative node tuples, we prove the following theorem (proof in Appendix A).
Theorem 2. There is a MAG-GNN whose action is more powerful in identifying discriminative node tuples than random actions.
MAG-GNN is at least as powerful as random actions since we can adopt a uniform-output MPNN for the MAG-GNN, yielding random actions. The superiority proof identifies cases where MAG-GNN requires fewer steps to locate the discriminative node tuples. The overall inference complexity is the number of MPNN runs,  $O(mtT|V^2|)$ . A more detailed complexity analysis is in Appendix D.
Some previous works also realize the complexity limitation of subgraph GNNs and propose sampling-based methods, and we discuss their relationship to MAG-GNN. PF-GNN [7] uses particle-filtering to sample from canonical labeling tree. MAG-GNN and PF-GNN do not overlap exactly. However, we show that using the same resampling process, MAG-GNN captures PF-GNN (Appendix A).
Theorem 3. MAG-GNN captures PF-GNN using the same resampling method in PF-GNN.