SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Proof. We consider the case where  $k = 2$  and  $m = 1$ . That is, we use one node tuple of size two. Let MAG-GNN implement the following policy: move one node in the node tuple to the node in the graph that is farthest away from the other node. This is easy to be implemented by MAG-GNN since
both nodes are marked, and we can compute the distance from all nodes in the graph to the two nodes with a sufficient number of MPNN layers. According to the properties of  $G(S)$ , MAG-GNN will move the node tuples so that they sit on the opposite graphs in the super graph. They identify the type of the two graphs, determining if  $G(S)$  is positive. Note that, in this case, we execute MAG-GNN once to obtain the discriminating node tuple independent of the super graph diameter.
On the other hand, if we use random actions, because node markings only help identify the graphs that the markings land on, marking nodes on any location other than the two opposite graphs will not help identify the super graph. The probability of random actions choosing the opposite graphs is  $\frac{1}{4n - 1}$ . Consequently, the expected number of random actions required to identify the super graph is  $4n - 1$ .
We first describe the PF-GNN algorithm; we refer the readers to Dupty et al. [7] for more details.
PF-GNN starts with a set of particles for a graph  $G$
$$
\left\langle \left(G, H _ {1} ^ {m}\right), w _ {1} ^ {k} \right\rangle_ {m = 1: M} \tag {14}
$$
where  $H$  is the node embedding of  $G$  using a MPNN  $g$ . For particle  $m$  in step  $t$ , PF-GNN uses a policy function  $P(V|H_t^m)$  to compute a distribution of nodes and sample a node to individualize. After using an MLP to update the embedding of the individualized node, PF-GNN uses GNN to refine all node embeddings. Specifically,
$$
v \sim P \left(V \mid H _ {t} ^ {m}\right), \left[ H _ {t} ^ {m} \right] _ {v} = \operatorname {M L P} \left(\left[ H _ {t} ^ {m} \right] _ {v}\right), H _ {t + 1} ^ {m} = g \left(H _ {t} ^ {m}\right) \tag {15}
$$
Then PF-GNN updates the particle weights as follows,
$$
w _ {t + 1} ^ {m} = \frac {f _ {o b s} \left(H _ {t + 1} ^ {m}\right) \cdot w _ {t} ^ {m}}{\sum_ {k} f _ {o b s} \left(H _ {t + 1} ^ {m}\right) \cdot w _ {t} ^ {m}} \tag {16}
$$
where  $f_{obs}$  is a function that measures the quality of the refined embeddings to a scalar value. PF-GNN then performs particle resampling from the distribution of  $\langle w_{t+1}^m \rangle_{m=1:M}$  and reset weights to  $1/M$ . Eventually, all  $H_{t+1}^m$  node representations undergo a readout function to generate graph embedding.
Proof. To prove MAG-GNN captures PF-GNN, we need to show that in each step, the node representation generated by PF-GNN and MAG-GNN are identical, resulting in the same node sampling distribution, resampling process, and weight update. Specifically, for a  $M$ -particle,  $T$ -step PF-GNN, we implemented a MAG-GNN with  $M$  node tuples of order  $T$  using  $T$  RL steps. The node tuples are still randomly initialized. The state matrix  $W$  is  $k$ -by- $(1 + T)$ , where  $[W]_{m,i} = 0, i \neq T + 1$  indicates the  $i$ -th node in node tuple  $m$  has not been individualized,  $[W]_{m,i} = j, i \neq T + 1$  indicates the  $i$ -th node in node tuple  $m$  is individualized at time step  $j$ , and  $[W]_{m,T + 1}$  is the weight in PF-GNN.
Let  $\pmb{v} = (v_{1},\dots,v_{t})$  be the sequence of nodes that PF-GNN used to individualize. We first note that there exists an MPNN applied onto the node marked graph  $G(\pmb{v})$  with a sufficient number of layers that output the same value as the final output,  $H_{t}$ , of PF-GNN. Specifically, let  $T$  be the number of layers of  $g$  in the MPNN, and then we can build a  $tT + t$ -layers MPNN  $g_{p}$  such that it alternates between  $T$  layers MPNN that is a copy of  $g$  and one layer MPNN that only update node features on marked nodes using the MLP of PF-GNN. Then, as long as we know the sequence, we can mimic a PF-GNN with an MPNN on node-marked graphs.
We then implement an update function  $f_{W}$  to log the step when a node  $v_{i}$  in the tuple  $m$  is marked that is  $[W]_{m,t} = t$ . We then set the Q-network to never output positive scores to the tuple index whose corresponding entries in  $W$  are greater than zero so that we fix the individualized nodes.  $f_{W}$  also update  $[W]_{m,T + 1}$  to be equal to the weight after the resampling process. Then, by using the  $tT + t$ -layers MPNN consists of  $g$  as  $g_{p}$  and another  $tT + t$ -layers MPNN consists of the MPNN of  $P(V|H_t)$  as  $g_{rl}$ , MAG-GNN can replicate the output of PF-GNN.
B Introduction to deep Q-learning
Deep Q-learning (DQN) is a neural network variant of Q-learning which is a popular algorithm for solving the Markov decision process (MDP). MDP is defined by a tuple  $(S, A, P, R, \gamma)$ , where:
-  $S$  is a set of states.
-  $A$  is a set of actions.
-  $P$  is a transition function that defines the probability of outcomes states after a particular action. For MAG-GNN, the transition is deterministic.
-  $R$  is a reward function defining the reward received by an action in a particular state.
-  $\gamma$  is a discount factor determining future rewards' importance.
Agents in an MDP aim to learn a policy that maps states to actions with maximal expected rewards, and Q-learning generates such agents. Specifically, Q-learning maintains a Q-table of Q-values  $Q(s,a)$ , representing the expected cumulative rewards of taking action  $a$  in state  $s$ . During training the Q-values are updated by,
$$
Q (s, a) = Q (s, a) + \alpha \left(r + \gamma * \max  _ {a ^ {\prime}} \left(Q \left(s ^ {\prime}, a ^ {\prime}\right)\right) - Q (s, a)\right) \tag {17}
$$
where  $\alpha$  is the learning rate,  $r$  is the instant reward of action  $a$  in  $s$  and  $s'$  is the resulting state of  $a$ . In more complex systems, keeping enormous Q-tables for all actions and states is infeasible, and DQN uses deep neural networks to build a Q-network to approximate  $Q(s,a)$ . The Q-network takes  $s$  as input and output  $Q(s,a)$  for all possible actions  $a$ . A loss function to train the Q-network is,
$$
\mathcal {L} = \mathbb {E} \left[ \left(r + \gamma * \max  _ {a ^ {\prime}} Q \left(s ^ {\prime}, a ^ {\prime}\right) - Q (s, a)\right) ^ {2} \right] \tag {18}
$$
A widely adopted method to train the Q-network is the Experience Replay. In each step, the RL agent interacts with the environment to store experience tuples  $E = (s, a, r, s')$ , and sample a set of previous experiences to train the Q-network. This technique effectively improves the sample efficiency. We adopt  $\epsilon$ -greedy training strategy, where the agent takes random actions with a probability of  $\epsilon$ . We start with  $\epsilon = 1$  and gradually decrease it to zero. This ensures that the agent explores the graph and learns the good policy.
C More discussion on the design of MAG-GNN
C.1 Action space
In Section C.1, we propose an action to substitute one node in the node tuple with another node. Such action requires a  $O(m|V|k)$  Q-table for selecting the node tuple  $(m)$ , selecting the node tuple index  $(k)$ , and selecting the node in the graph. We also propose not to select node tuple to reduce the  $m$  factor, then the action in 7 becomes
$$
\boldsymbol {v} _ {i} ^ {\prime} = a _ {j, l} (\boldsymbol {v} _ {i}) = \left(\left[ \boldsymbol {v} _ {i} \right] _ {1}, \dots , \left[ \boldsymbol {v} _ {i} \right] _ {j - 1}, v _ {l}, \left[ \boldsymbol {v} _ {i} \right] _ {j + 1}, \dots \left[ \boldsymbol {v} _ {i} \right] _ {k}\right) \tag {19}
$$
The action space becomes  $A = [k] \times \mathcal{V}$ . In such a case, we need to enforce an order to update each node tuple sequentially or simultaneously update them. In the experiments, we choose to update all node tuples for better parallelizability simultaneously.
A seemingly more natural design choice is directly replacing the current node tuple with another one. However, it has  $O(|V^k|)$  possible actions for one node tuple. Its corresponding Q-table size is also exponential to  $k$ . The cost of computing such a table may equal that of running the complete subgraph GNN, which violates our motivation for complexity reduction.
The state matrix  $W$  serves as a tracker of previous states. The choice of update function  $f_{W}$  can vary. It can simply be a mean pooling function on the embeddings of marked nodes. It can also be an LSTM pooling function. While using matrix  $W$  has good theoretical properties as demonstrated in Appendix A, we found it is not as helpful practically and becomes a confounding factor under the SIMUL training scheme because the states matrix stored in the memory does not properly reflect the prediction GNN. Hence, we use zero-initialized  $W$  and an identity update function during training.
C.2 Training paradigms
We provide pseudo-codes for each training paradigm of MAG-GNN. All training paradigms are based on experience replay, where in each step the agent obtains experiences from the environment, stores the experiences into memory, and retrieves other experiences from memory to optimize the
Algorithm 1 RL-Experience
Input: Node tuples  $U$ , batch of data  $B$ , buffer  $D$ , RL agent  $a$  and GNN  $g_{rl}$
Output: New node tuples.
1: Use  $a$  and  $g_{rl}$  to generate a new node tuples  $U'$  according to equation (9) and (10)
2: Obtain rewards  $r$  from  $U^{\prime}, U$ , and  $B$  according to equation (8)
3: Store the experience tuple  $(U', U, B, r)$  to  $D$
4: return  $U'$
Algorithm 2 ORD-Train
Input: A target dataset of  $(\mathcal{G},\mathcal{Y})$ . A graph prediction model  $f_{p}$ , an RL agent  $a$  with RL GNN  $g_{rl}$ .
Output: A trained RL agent  $a$  that output the most discriminative node tuples and a corresponding prediction GNN  $f_{p}$ .
1: while  $f_{p}$  is not converged do
2:  $B\gets$  a batch from  $(\mathcal{G},\mathcal{V})$
3: Sample random node tuples  $U$  for graphs in  $B$
4: Optimize  $f_{p}$  according to equation (4) w.r.t  $B$  and  $V$ .
5: end while
6:  $D\gets$  memory buffer
7: Call RL-Experience until  $D$  has enough memory
8: while  $g_{rl}$  is not converged do
9:  $B\gets$  a batch from  $(\mathcal{G},\mathcal{V})$
10: Sample random node tuples  $U$  for graphs in  $B$
11: RL-Experience  $(U,B,D,a,g_{rl})$
12: Sample experience tuples  $E$  from  $D$
13: Optimize  $g_{rl}$  using the  $E$
14: end while
15: return  $f_{p}, a, g_{rl}$
Q-learning agent. Since we are training the agent (RL GNN) while the environment (prediction GNN) is evolving, we propose different training paradigms to better model both sides. Different training paradigms differ mainly in the schedule to train the two sides. ORD paradigm first trains the environment without the agent and then trains the agent in the fixed environment. SIMUL paradigm trains the agent and environment together, where an environment training step is followed by the agent collecting experience from the current environment and optimizing according to past experience. The SIMUL paradigm is proposed because the trained environment might overfit the data and hence give incorrect rewards to the agent. We provide pseudo-codes for the paradigms.
C.3 Other sampling-based methods
The advantage of MAG-GNN over other subgraph sampling methods, including PF-GNN [7] and DropGNN [24], is that it effectively chooses the most discriminative subgraphs. DropGNN drops nodes with equal probability; PF-GNN uses an MPNN to model the drop probability of nodes which reduces to the uniform distribution when the graph is regular and has no features. According to Figure 2, there is still a high chance that the randomly sampled subgraphs are not informative. On the contrary, MAG-GNN starts with random subgraphs but uses previous knowledge about the correlation between subgraphs to refine the likelihood of a subgraph being informative. RL training helps the model acquire such knowledge. Just as Theorem 2 points out, previous methods are essentially MAG-GNN with random action. In contrast, learned actions can significantly improve efficiency and effectiveness.
C.4 Limitations
As mentioned in Section 4.1, MAG-GNN does not train as well on datasets with labels, and the training time is longer than subgraph GNNs despite better inference complexity. Note that we are employing the basic Q-learning techniques in MAG-GNN to demonstrate the framework's power,
Algorithm 3 SIMUL-Train
1:  $D\gets$  memory buffer
2: Call RL-Experience until  $D$  has enough memory
3: while  $f_{p}$  and  $g_{rl}$  are not converged do
4:  $B\gets$  a batch from  $(\mathcal{G},\mathcal{V})$
5: Sample random node tuples  $U$  for graphs in  $B$
6:  $U^{\prime} = \mathbb{R}\mathbb{L}$  -Experience  $(U,B,D,a,g_{r\bar{l}})$
7: Optimize  $f_{p}$  according to equation (4) w.r.t  $B$  and  $U'$
8:  $B^{\prime}\gets$  a batch from  $(\mathcal{G},\mathcal{V})$
9: Sample random node tuples  $U$  for graphs in  $B^{\prime}$
0: RL-Experience(U,B,D,a,grl)
1: Sample experience tuples  $E$  from  $D$
2: Optimize  $g_{rl}$  using the  $E$
3: end while
4: return  $f_{p}, a, g_{rl}$
Algorithm 4 PRE-Train
Input: A target dataset of  $(\mathcal{G},\mathcal{Y})$ . A graph prediction model  $f_{p}$ , an RL agent  $a$  with RL GNN  $g_{rl}$  trained with either SIMUL or ORD on a different dataset.
1: while  $f_{p}$  is not converged do
2:  $B\gets$  a batch from  $(\mathcal{G},\mathcal{V})$
3: Sample random node tuples  $U$  for graphs in  $B$
4:  $U^{\prime} = \mathrm{RL}$  -Experience(U,B,D,a,grl)
5: Optimize  $f_{p}$  according to equation (4) w.r.t  $B$  and  $U'$
6: end while
7: return  $f_{p}, a, g_{rl}$
and we can adopt more robust training methods from the large RL toolbox to improve training time. We leave this to future work.
An advantage of subgraph GNN is its equivariance property. Given the same graph up to isomorphism, subgraph GNNs always generate the same embedding for the graph, whereas MAG-GNN sacrifices such property for better efficiency. However, MAG-GNN still provides a higher level of certainty than other sampling-based and randomization-based non-equivariant GNNs, because when the initial node marking is fixed, MAG-GNN can also generate the same embedding for the graph. Then, suppose we can deterministically generate initial node markings; we also grant MAG-GNN equivariance. Such initial node marking can be computed from canonical labeling or previous knowledge about the graph.
While we see MAG-GNN's potential in node-level tasks in Section 6, MAG-GNN is primarily defined at the graph level, and its reward and actions might not be the most appropriate for node-level tasks. We conduct the evaluation in Section 3 on the Cycle counting dataset; we observe that even for 6-CYCLE counting, there exist node marked graphs that result in near zero error while our Q-learning framework fails to identify such node tuple. We suspect this is because the MAG-GNN's reward is defined at graph-level, and taking a plain average of node error does not faithfully capture the environment. This leads to exciting research paths extending MAG-GNN to node-level or even edge-level tasks.
Table 7: BREC dataset results (↑)
Model	Basic (60)	Regular (140)
Number	Accuracy	Number	Accuracy
NGNN [35]	59	98.3	48	34.3
SSWL+ [36]	60	100	50	35.7
I2GNN [15]	60	100	100	71.4
OSAN [26]	56	93.3	8	5.7
MAG-GNN	49	81.6	61	43.5
The complexity of MAG-GNN is straightforward. It runs MPNN multiple times; hence, the complexity is a multiple of MPNN  $(O(T|V^2))$ . Specifically, for a MAG-GNN with  $m$ $k$ -tuples that runs  $t$  steps, there is a total number of  $O(mt)$  MPNN runs, incurring  $O(mtT|V^2|)$  total complexity. The complexity of computing the Q-table is equal to the size of the Q-table, which is  $O(m|V|k)$ . We compute Q-tables  $O(t)$  times, and the overall complexity for computing the Q-table is  $O(mkt|V|)$ . The overall complexity is  $O(mkt|V| + mtT|V^2|)$ . In practice,  $O(mkt) < O(T|V|)$  as can be observed in Table 10 and hence the complexity can be reduced to  $O(mtT|V^2|)$ . Note that subgraph GNNs have the complexity of  $O(T|V^k|)$ , which grows exponentially with  $k$ , while MAG-GNN's complexity grows linearly with  $k$  even when  $k$  is large.
In equation 9, we propose to use graph embedding function  $f_{p}$  to represent the current state; however, an alternative is to use  $f_{rl}$  to describe the current state. This allows us not to compute graph embedding for prediction but only compute that for Q-table. This does not change the complexity but can reduce the number of MPNN runs in half, improving efficiency.
Figure 8: Reward increases over time.
Table 7 shows the results on the BREC dataset's basic and regular graph subsets [28]. BREC dataset contains 400 pairs of graphs that require different levels of expressiveness to distinguish. The regular graph dataset contains strongly regular and distance-regular graphs that cannot be distinguished by 3-WL GNNs. We apply a fixed 2-node-tuple MAG-GNN to these pairs and test 1000 times with different initial node tuples. We say the model distinguishes the pair if the model correctly classifies the pair at least 950 times. We then count the ratio of correctly distinguished pairs. We observe that MAG-GNN falls short on the basic graph, possibly due to the excessive expressiveness and the difficulty it poses to training. However, we also observe that MAG-GNN performed well on regular graphs, outperforming NGNN and SSWL+ with  $< 3$ -WL expressiveness. While outcompeted by I²GNN, MAG-GNN is considerably more efficient. We also noticed that MAG-GNN significantly improves the performance on regular graphs over OSAN, further validating that RL-selected subgraphs carry more information.
Move markings apart
Agent stagnant
Move marking apart
Agent keeps moving
Common neighbors w/ edge
Common neighbors w/o edge
Figure 9: MAG-GNN example on Shrikhande graph (Top) and Rood 4-by-4 graph (Bottom). Red arrows show how MAG-GNN moves the node tuples. Green edges show the difference in the subgraphs defined by the resulting node tuples.
In Figure 8, we show the reward improvement over the training step on the SR25 dataset. We plot the mean reward of every 500 steps and plot the standard deviation of the 500 steps by the shaded range. We can see that the reward is close to zero at the initial stage of training because the  $\epsilon$  is set to 1, and the agent essentially takes random actions. After a period of searching where the reward stays at around 0, the output node tuples gradually become informative. Eventually, the step-average reward converges at 0.05. We see a few sudden drops in reward during training due to the "overestimation bias" commonly seen in RL. Nevertheless, we see that the agent can correct the mistake and eventually converge. More advanced deep Q-learning techniques like double Q-learning might resolve the issue, and we leave this as future work.
We conducted the experiments on the Shrikhande graph versus the Rook 4-by-4 graph toy example. The task is to differentiate the two strongly regular graphs with the same configuration and hence cannot be distinguished by MPNN and node-rooted subgraph GNNs. We train a MAG-GNN with 2-node markings to break the 3-WL expressivity bound. After training, the MAG-GNN generates the trajectory given the initial 1-distance random node markings as shown in Figure 9. We can observe that MAG-GNN first learns to push the two marked nodes apart in both graphs. However, in the Shrikhande graph, after it finds a marked node pair whose common neighbors are connected, the agent becomes stagnant, while in the Rook 4-by-4 graph, the agent is never stagnant and keeps moving one marked node, because none of the common neighbors of 2-distanced node pairs in the Rook 4-by-4 graph are connected. Such information is easy to be learned by a downstream MPNN.
F More experiment details
F.1 Datasets statistics
We took the CYCLE dataset from [37] repository. All training and testing are on standardized targets following the previous practice [37, 15]. The task is regression, and we use continuous value to approximate the discrete number of cycles. We use the QM9 dataset provided by Pytorch-Geometric [9], and we use a train-valid/test split ratio of  $0.8 / 0.1 / 0.1$ . All QM9 targets are standardized during training but mapped back to their actual value for prediction and testing. We use the ZINC dataset provided by Pytorch-Geometric [9] and use the official split. We take OGBG-MOLHIV dataset from the Open Graph Benchmark package [14] and use their official split. For CSL and EXP, we follow
Table 8: Hyperparameters Summary
Method	QM9	ZINC	ZINC-FULL	CYCLES
Learning rate	0.001
Weight decay	0	0	0	0
# GNN Layers	5	6	6	6
Jumping Knowledge	Sum	Sum	Sum	Sum
m	{2,3,4}	{2,3,4}	{2,3,4}	{2,3,4}
k	{3,4,5,6}	{3,4,5,6}	{3,4,5,6}	{3,4,5,6}
t	2	2	2	2
# Epochs	2000	4000	500	5000
Batch size	512	512	512	500
Memory size	3M	300K	300K	300K
Sync rate	3000	3000	3000	3000
# runs	4	10	3	10
Table 9: Hyperparameters Summary (cont.)
Method	MOLHIV	EXP	SR25	CSL
Learning rate	0.001
Weight decay	0.001	0	0	0
# GNN Layers	6	4	4	4
Jumping Knowledge	Sum	None	None	None
m	{2,3,4}	{1,2,3}	{1,2,3}	{1,2,3}
k	{3,4,5,6}	{1,2,3}	{1,2,3}	{1,2,3}
t	2	4	4	4
# Epochs	1000	500	10000	500
Batch size	128	64	32	64
Memory size	300K	30K	30K	30K
Sync rate	3000	100	150	150
# runs	10	10	10	10
Table 10: Datasets statistics.
Dataset	#Graphs	Avg. #Nodes	Avg. #Edges	Task type
CYCLE	5,000	18.8	31.3	Node regression
QM9	130,831	18.0	18.7	Graph regression
ZINC	12,000	23.2	24.9	Graph regression
ZINC-FULL	249,456	23.1	24.9	Graph regression
ogbg-molhiv	41,127	25.5	27.5	Graph classification
previous practices and use 10-fold cross-validation to train the model and report the mean results. For SR25, since there are only 15 different graphs, we follow previous practices and use the complete data for training, valid, and test sets to examine the expressivity of the model.
F.2 Training details
We summarize the hyperparameters used for different datasets in Table 8 and 9.
Curly bracket means the range of hyper-parameters we search. Memory size is the size of the memory buffer for experience and replay. Sync rate is the number of optimization steps before the target Q-network sync with the train Q-network. For ZINC datasets, we additionally use learning rate decay with a rate of 0.5 for every 200 epochs. All models are implemented in DGL [29] and PyTorch [25]. We only use Pytorch Geometric [9] to get the datasets.
Since there is a random factor in MAG-GNN (initial node tuples), all tests are repeated ten times, and we take the average score as the test score of one run. We ran experiments multiple times with different random seeds to take the average scores.