MAG-GNN: Reinforcement Learning Boosted Graph Neural Network
Authors: Lecheng Kong
Jiarui Feng
Hao Liu
Dacheng Tao
Yixin Chen
Muhan Zhang
Venue: NeurIPS 2023 poster (2023)
Keywords: Graph Neural Network; Reinforcement Learning
TL;DR: We propose a Reinforcement Learning method to find most expressivie subgraphs and boost the expressivity of GNNs.

Abstract: 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.

MAG-GNN: Reinforcement Learning Boosted Graph Neural Network
Lecheng Kong'
Jiarui Feng' Hao Liu'
Dacheng Tao?
Yixin Chen Muhan Zhang²
(jerry. kong, jiarui, liuhao, dacheng, tao@gmail com, muhan@pku edu,
Washington University St. Louis 2JD Explore Academy *Peking University
Abstract
While Graph Neural Networks (GNNs) recently became powerful tools in graph learning tasks, considerable efforts have been spent improving GNNs' structural encoding ability. A particular line 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 identifica- tion of the optimal subset as a combinatorial optimization problem and propose Magnetic Graph Neural Network (MAG-GNN), reinforcement learning (RL) boosted GNN, to solve the problem. Starting with a candidate subgraph set, MAG- GNN employs an RL agent 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 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.
Introduction
Recent advances in Graph Neural Networks (GNNs) greatly assist the rapid development of many ar-
eas, including drug discovery 22) recommender systems [31], and autonomous driving (6) The power
of GNNs has primarily been attributed to their Message-Pastie Passing Paradigm [12]. Message-Passing
Paradigm simulates I-dimensional Weisfeiler-Lehmee
(1-WI.) algorithm for graph isomorphism
testing. Such simulation allows GNN to encode rich structural information. In many fields, structural information crucial to determine the properties of graph.
However, as Xu et al. 332 pointed out, GNN's structure encoding capability, or its expressiv- ity, is also upper-bounded by the 1-WL test. Specifically, 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 includ- ing chemistry and biology. To overcome this limitation, considerable effort was spent on in- vestigating more-expressive GNNs. famous line of work is subgraph GNNs B3 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 sum- marized 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) (TO].
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

This spurs the research on extensions to edge-rooted subgraph GNNs with better expressivity [[5]. In fact, by linking subgraph GNNs to the k-dimensional WI. (k-WI) 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 ex- ponential to k [26 [O]. Hence, the higher expressivity, unfortunately, comes with a higher computational cost, which makes it challenging to apply subgraph GNNs to even moderately sized graphs (c.g., 100 nodes).
Beyyaar Grogle
Stip 1-Eap Subgraphs
B
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 as they are 2-regular graphs with identical 1 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 run MPNN once 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 enumcrating all subgraphs. We also include advanced examples with more complex structures in Section 3
Figure 1: Comparison of two simple graphs.
Therefore, we propose Magnetic Graph Neural Network (MAG-GNN), 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 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-GNI repeats the process until it identifies the set of subgraphs with the highest distinguishing power. The resulting subgraph set is then passed prediction GNN for downstream tasks. MAG-GNN reduces subgraph GNN's exponentially complex enumeration procedure to RL searching process with constant steps. This potently constrains the computational cost while keeping the expressivity.
We conduct extensive experiments synthetic and real-world graph datasets and show that MAG- GNN achieves competitive performance to state-of-the-arr (SOTA) methods and even outperforms subgraph GNNs many datasets with 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.
2 Preliminaries
A graph can be represented as G = EV,E, where V is the set of nodes and E C V x V is the set of edges. Let (G) and E(G) represent the node and edge sets of G, respectively. Nodes are associated with features X = (x.Ive E V). An MPNN g can be decomposed into T layers of COMBINE and AGGREGATE functions. Each layer uses the COMBINE function update the current node embedding from its previous embedding and the AGGREGATE function to process the node's neighbor embeddings. Formally,
m AGGGEGATEE ( (((h((-)) uu (()))), h COMBINE (m (m.".,((-))) ((!1))
(I)
where h is the node representation after t iterations, h (0) = Z,., (()) is the set of direct neighbors of v, m. is the message embedding. (? is used form node, edge, and graph-level representations.
We use H = gG) 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 graph G, a subgraph GNN first enumerates all o-rreer node tuples (v/v V*(()) and creates (G)) copies of the graph. A
2

---

graph associated with node tuple is represented by G(v). Node-rooted subgraph GNNs adopt l-rreer policy and have ((V(G)) graphs; edge-rooted subgraph GNNs adopt a 2-order policy and have (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,
[(())lpp = f(--
c
if = |ol, and p j VE (G),
(2)
otherwise
where @ means row-wise concatenation. We use square brackets to index into a sequence. All entries in X ((v) are labeled e except those appearing at corresponding positions of the node tuple. An MPNN g is applied to every graph, and we use pooling function to obtain the collective embedding of the graph:
(v) E R and G(v)= EV, @ ((v)).
f=(G) = R( ((o((())) E E (G))),
PE[G,N].
f. fG) can be a vector of graph representation if ((P) is graph-level pooling function (P equals G) or matrix of node representations if R(r) 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 (3). Furthermore, the expressivity of subgraph GNNs increases with larger values k. Since node tuples, node-marked graphs, and subgraphs refer to the same thing, we use these terms interchangeably.
The Weisfeiler-Lchman hierarchy (WL hierarchy). The k-dimensional Weisfeiler-Lehmar al- gorithm 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 el al. [32] and Morris 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- -ubgraph WL (k-OSWL) hierarchy which is comparable to the k-WL. test.
Deep Q-learning (DQN). DQN [9 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 given state. Accumulating sufficient experience with the environment, DQN can make decisions in intricate state spaces. For detailed introduction to DQN, please refer Appendix B
Motivation
Figure I shows that while an MPNN encodes rooted subtrees, subgraph GNN encodes rooted subgraphs around each node. This allows subgraph GNN to differentiate more graphs at the cost of O(VV)) MPNN runs. Mcanwhile, subgraph GNNs are still bounded by 3-WL tests. Hence, we may need at least (e.g., edge-rooted) subgraph GNNs, requir- ing 1 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 com- putational 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 obtain similar expressivity, just like in Figure [I where subgraph is powerful as the collective information of all subgraphs? This leads us to the following experiment.
1.0
0.8
Softmax Score
0.6
2-node-tuple-ron
0.4
0.2
C15
Acc
0.0
280
400
600
Node Tuple Indek
Figure 2: Sorted scores.
We focus on the SR25 dataset. It contains 15 different strongly configuration, each of 25 nodes. The goal graphs. Since node-rooted subgraph GNNs are upper- 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 subgraph GNNs have expressivity beyond 3-WL. and can distinguish any pair of graphs from the dataset, but it takes 625 MPNN runs.
graphs with the same do multi-class classification to distinguish all pairs of
3

---

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 L,LLL e)), v),, (. R (g) V BE V 9p
(4)
where L is the loss function, MLP is multi-layer perceptron. 9p is an MPNN, and RC)) pools the node representations to graph representations. Unlike 2-node-tuple-rooted subgraph GNNs that run the MPNN |V²| 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-marker graph for classification. We perform graphs, with
ten independent tests, and the average test accuracy is 66.8%. Using 2-node- only one GNN run, it already outperforms node-rooted subgraph GNNs that fail this dataset. More interestingly, for each graph G, we can sort the classification score of its |V² possible node-marked graphs and plot them in Figure2(C14 and C15 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 [ 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 wide range of graph classes. We prove that such phenomenon exists in most regular graphs, which cannot be distinguished by MPNNS (proof in Appendix[A] (A].
Theorem 1. Let G1 and G2 be two graphs uniformly sampled from all n-node, r-regular graphs where 3 S r < v2log n. Given an injective pooling function R(G) and an MPNN 9p of I-WL expressivity, with a probability of at least 1 o(1), there exists node (tuple) v (G) whose corresponding node marked graph's 's embedding, fr (G1, v), is different from any node marked graph s embedding in Gz-
These observations show that by finding discriminative subgraphs effectively, we only need to apply MPNN much smaller subset of the large complete subgraph set to get close level of expressivity.
4 Magnetic graph neural network
We formulate the problem of finding the most discriminative subgraphs as a combinatorial optimiza- tion problem. Given a budget of m as the number of subgraphs, is as the order of node tuples, and gp an MPNN that embeds individual subgraphs, we minimize the following individual loss to graph G,
min U= DE (G)) TM [(MLP(,(CC,U),
(5)
fz(G,U) = RP (1g, e U}),
Note that this formulation resembles that of subgraph GNNs in Equation 3 we are only substituting V 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, state is m [0] 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 is defined as,
(G. = (G, (wi..., v.)) W), S G (V&)m x (R
(6)
S is the state space, G is the set of all graphs, and D is the set of all possible k node tuples of 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 (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 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 l: can reduce the sample space and stabilize training in this case.
Action Space: We define one RI. 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

---

Step
Step 2
Step t
Original graph without node marking
RL agent node tuples for steps.
Graph with one node marked
Randomly mark node
Inference
Train
Target difference
Prediction
L
Reward
f( (
Q-Table
arginax
Marked node
RL GNN
RL. RLA.gem
Prediction GNN
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, action aiji. on
(G, U, W) does the following on U:
(7) The agent selects target node tuple v.. whose j j-th node is replaced with node on E V. W is then updated by arbitrary state update function W = fw(s: U²)) depending
U' (U) = (W)..., 11 v/,, P.. 0y ([v)),,.., ((0])),,.,,..))_,,,,,,, wt, [u:]]+1,....o.)p]]
the old state and new node tuples. The update function fw 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] x [к] x 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 ((ml V{k,, which is linear the number nodes, and l is usually small (k: is already more expressive than most subgraph GNNs). We further reduce the action space O ([V])) k)) if
the agent does not need to select the update target but uses Q-network to do Equation ( on a given
0;- In such case, either sequentially 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 detailed discussion on the action space and alternative designs in Appendix C.I
In Section 3] show that proper mode-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 expressivity should affect the objective value, we choose the objective value improvement as the instant reward. Specifically, let (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
Reward:
hardly quantify it. Instead, since the relevant improvement in
r(s,a, a, s') = ((ML
(GML
(8)
This reward correlates the action directly with the objective function, allowing our RL agent to be desk-depnndent and more flexible for different levels of tasks.
Because our state includes graphs, we require an equivariant Q-network to output consistent Q-tables for actions. Hence, we choose MPNN parameterize the Q-network. Specifically, for current state (G, U, W) and the target node tuple U, we have the Q-table
Q-Network:
Q(ss. v)] MLP(([,, C Σ R (g.((()) @ R(W((W)
(9)
NEU
Row / in the Q-table is computed by the embedding of node vy in the node-marked graph by an MPNN Get, the current overall graph representation across all node tuples, and the state matrix W summarized by pooling function R(W), QQi, represents the expected reward of replacing the node on index j of node tuple v with node uz. The best action aj,1 is then chosen by.
arg max max[Q((s v)lt.j
(10)
j.1
Note that because we assign different initial embeddings based on the node tuple, the MPNN distinguishes otherwise indistinguishable graphs.
5

---

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 MPPNN-parametert 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 & 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- (MAG-GNN). To show the effectiveness of our method in locating discriminative node tuples, we prove the following theorem (proof in Appendix A}.
Q-newwork to update one slot in
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 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 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- G-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-GNT captures PF-GNN (Appendix (A).
Theorem 3. MAG-GNN captures PF-GNN using the same resampling method in PF-GNN.
K-OSAN (26] proposes a data-driven subgraph sampling strategy to find informative subgraphs by another MPNN. This strategy reduces to random sampling when the graph requires higher expressivity (c.g., regular graphs) and has no features because the MPNN will generate the same embedding for all nodes and hence cannot identify subgraphs that most benefit the prediction like MAG-GNN can. MAG-GNN does not solely depend on the data and finds more expressive subgraphs even without node features. Moreover, sampled subgraphs in previous methods are essentially independent. In contrast, MAG-GNN also models their correlation using RL. This allows MAG- GNN obtain better expressivity with fewer samples and better consistency (More discussions in Appendix C.3)
4.1 Training MAG-GNN
With the state and action space, reward, and Q-network defined, we can use any DQN techniques to train the RL agent. However, to evaluate the framework's capability, we select the basic Experience Replay method 19 to train the Q-network. MAG-GNN essentially consists of two systems, an RL agent and a prediction MPNN. Making them coordinate is more critical to the method. The most natural way to train our system is first to train 9p. introduced in Section 3 with random node tuples. We then use Ip as part of fp. the marked-graphe encoder, and treat fp as the fixed environment to train our Q-network. The advantage of the paradigm is that the environment is stable. Consequently, all experiences stored in the memory have the correct reward value for the action. This encourages stability and fast convergence during RL training. We term this paradigm ORD for ordered training.
However, not all Ip trained from random node tuples are good. When we train Ip for the ZINC dataset and evaluate all node-marked graphs to as in Section 3 The average standard deviation among all scores of node-marked graphs is -0.003, and the average difference between the worst and best score is ~0.007. Hence, the maximal improvement is minimal if we use this MPNN as the environment. Intuitively, when the graph has informative initial features, X, like those in the ZINC dataset, the MPNN quickly recognizes patterns from these features while gradually learning ignore node marking features X (v,), as not all node markings carry helpful information. In such cases, we need to adjust 9p while the RL agents become more powerful in finding discriminative subgraphs.
One way is train the RL agent and 9p simultaneously. Concretely, we sample state s, run the RL agent t steps to update it to state s', and train Jp on the marked graphs of s'. Then, in the same step, the RL agent samples a different state and treats the current 9p as the environment to generate experience. Lastly, the RL agent is optimized on the sampled previous experiences. Because we adjust 9p to capture node tuple information better, the score standard deviation of the node-marked ZINC graphs is kept at -0.036. We term this paradigm SIMUL. Compared to ORD, SIMUL makes

---

the RL agent more difficult to train when by evolves rapidly. Nevertheless, we observe that as Ip gradually becomes stable, the RL agent can still learn effective policies.
One of the critical goals of MAG-GNN is to identify complex structural information that MPNN cannot. Hence, instead of training the agent on real-world datasets from scratch, we can transfer the knowledge from synthetic expressivity data to real-world data. As mentioned above, training MAG-GNN on graphs with features is difficult. Alternatively, we first use the ORD paradigm train the agent on expressivity data without node features such as SR25. On these datasets, 9p relies on the node markings to make correct predictions. We then fix the RL agent and only use the output state from the agent to train a new 9, for the actual tasks, such as ZINC graph regression. Using this paradigm, we only need to train one MAG-GNN with good expressivity and adapt different tasks without worrying about overfitting and the aforementioned stability issue, we name this paradigm PRE. We experimented with different schemes Section
Related work
More expressive GNNs. A substantial amount of work strive to improve the expressivity of GNNs.
They can be classified into the following categories: (1) Just like MPNN simulates the 1-WL test, Higher-order GNNs design GNNs to simulate higher-order WL tests. They include k-GNN [18).
RingGNN [5], and PPGN (20). These methods perform message passing node tuples and have complexity that grows exponentially with k and hence does not scale well to large graphs. (2) Realizing the soauettr--breaknng power of subgraphs, Subgraph GNNs, including NGNN (35. GNN-AK 37, KPGNN R and ID-GNN B3 use MPNNs to encode subgraphs instead of subtrees around graph nodes. Later works, like 12-GNN 15 further use 2-order(edge) subgraph information to improve the expressivity of subgraph GNNs beyond 3-WL. Recent works, such as OSAN 26) and SUN [0], unify subgraph GNNs into the WL hierarchy, showing improvement in expressivity for subgraph GNNs also requires exponentially more subgraphs. (3) Substrcuture counting methods, including GSN 3 and MotifNet (22), employ substructure counting GNNs. They count predefined substructures undetectable by the 1-WL test as features for MPNN to break its expressivity limit, However, the design of the relevant substructures usually requires human expertise.
(4) Many previous works also realize the complexity issue of more expressive GNNs and strive to reduce it. SetGNN (38 proposes to reduce node tuple to set and thus reduce the number of nodes in message passing. GDGNN [16] designs geodesic pooling functions to have strong expressivity without running MPNN multiple times. (5) -quivvavvar GNNs. Abboud el al. [1 proves the universality of MPNN with randomly initialized node features, but due to the large search space, such expressivity is hardly achieved. DropGNN (24 randomly drops out nodes from the graph to break symmetries in graphs. PF-GNN (7) implements neural version of canonical labeling and uses particle filtering to sample branches in the labeling process. OSAN (26) proposes to use input features to select important subgraphs. Agent-based GNN (21) initializes multiple agents on a graph without the message- sasssagepassingg paradigm to iteratively update node embeddings. MAG-GNN also falls into this category.
Reinforcement Learning and GNN. There has been a considerable effort to combine RL and GNN. Most work is on application. GraphNAS [I] uses GNN encode neural architecture and use reinforcement learning to search for the best network. Wang et al. (30] uses GNN to model circuits and RL to adjust the transistor parameters. On graph learning, most works use RL to optimize particular parameters in GNN. For example, SUGAR (7 uses Q-learning to learn the best top-k subgraphs for aggregations; Policy-GNN [17] learns to pick the best number of hops to aggregate node features. GPA [33 uses Deep Q-learning to locate the valuable nodes in active search. These works leverage the node feature to improve the empirical performance but fail to identify graphs with symmetries, while MAG-GNN has good expressivity without node features. To the best of our knowledge, our work is the first to apply reinforcement learning to the graph expressivity problem.
Experimental results
In the experiment | we answer the following questions: QI: Does MAG-GNN have good expressivity, and is the RL agent output more expressive than random ones? Q2: MAG-GNN is graph-level;
'The code can be found attn:://gtttooooo//ccccoggnnnggg

---

Table Synthetic results. (1)
Table 2: Cycle counting results. ())
EXP
CSL
SR25
Method
3-CYCLES
4-CYCLES
5-CYCLES
6-CYCLES
GIN
32
50.0
10.0
6.7
GIN 32
0.3515
0.2742
0.2088
0.1555
RNI 1
99.7
16.0
6.7
RNM
0.0041
0.0129
0.0351
0.0327
NGNN 35
100
100
6.7
ID-GNN 33
0.0006
0,0022
0.0490
0.0495
GNNAK+
37
100
100
6.7
NGNN RS
0.0003
0.0013
0.0402
0.0439
SSWL+
36
100
100
6.7
GNNAK+ 37
0.0004
0,0041
0.0133
0.0238
RNM
100
100
93.8
PGNNN [5
0,0003
0.0016
0.0028
0,0082
GNNN
15
100
100
100
MAG-GNN
100
MAG-GNN
0.0029
0.0037
0.0097
0.0286
can the expressivity generalize to node-level tasks? Q3: How is the RL method's performance on real-world datasets? Q4: Does MAG-GNN have the claimed computational advantages?
For the experiment, we update all node tuples simultaneously as it allows more parallelizable computation. More experiment and dataset details can be found in Appendix
6.1 Discriminative power
Dataset.To To answer Q1, we use synthetic datasets (Accuracy) to test the expressivity of models. (1) EXP contains 600 pairs of non-isomorphic graphs that cannot be distinguished by 1- WL/2-WL bounded GNN. The task is to differentiate all pairs. (2) SR25 contains 15 strongly regular graphs with the same configuration and cannot be distinguished by 3-WL bounded GNN. (3) CSL contains 150 circular skip link graphs in 10 isomorphism classes. (4) BREC contains 400 pairs of synthetic graphs to test the fin-grained expressivity of GNNs (Appendix
We use the ORD training scheme.
Results. We compare to MPNN 62 : Random Node Marked (RNM) GNN with the same hyperparameter search space MAG-GNN, subgraph G55 15 36 371, and Non- equivariant GNNs [1 as baselines. Table shows that MAG- GNN achieved perfect score on all datasets, which verifies our observation and theory. Note that subgraph GNNs like NGNN take at least |V] MPNN runs, while MAG-GNN takes constant times of MPNN runs. However, MAG-GNN success- fully distinguished all strongly regular graphs SR25, which NGNN failed. RNI, despite being universal, is challenging to train and can only make random guesses on SR25. Compared to the RNM approach, MAG-GNN finds more representative subgraphs for the SR25 dataset and performs better. Figure shows graph in the EXP dataset. We observe that MAG-GNN moves the initial node marking on same component to dif- ferent components, allowing better structure learning. Another example of strongly regular graphs is in Appendix E In Figure 5 also plot the performance of MAG-GNN against the number of MAG-GNN search steps when we only use one node tuple of size two. We see that node tuples from MAG-GNN are significantly more expressive than random node tuples (step 0). On EXP and CSL datasets, MAG-GNN can achieve perfect scores in one or two steps, whereas in SR25, it takes six steps but with a consistent performance increase over steps. We plot the reward curve during training in Appendix B
Figure 4: Initial Random Node Marking (Left). MAG-GNN gen- erated markings. (Right)
1.0
Accuracy
EXP CSL
0.7
Agent Steps
Figure 5: Performance versus the
number of RL steps.
6.2 Node-level tasks.
Datasets. answer Q2, we adopt the synthetic cycle counting dataset, CYCLE (MAE), Zhao el al. 337. The task is to count the number of (3,4.5.6)-cycles on each node. MPNN cannot count these structures. Node-rooted GNN can count (3,4)-cycles, while only models with expressivity beyond 3-WL can count all cycles. We use the ORD training scheme.

---

Table QM9 experimental results on all targets. (L)
Target
RNM
1-2-3-GNN 23
PPGN 20
NGNN 351
P--GNN 15
MAG-GNN
Comp.
((
O(T(V
O
O(TIV)))
O(T[V*))
O(mfT(V2))
0.426
0.476
0.231
0.428
0.428
0.353
a
0.306
0.27
0,382
0,230
0.230
0.226
CHOM
0.00258
0.00337
0,00276
0.00265
0.00261
0.00257
ÉLUMO
0.00269
0.00351
0.00287
0.00297
0.00267
0.00252
Ac
0.0047
0.0048
0,00406
0.0038
0.0038
0.0035
(R*)
20.9
22.9
16.07
20.5
18.64
15.44
ZPVE
0.0002
0.00019
0.0064
0.0002
0.00014
0.0002
Ug
0,281
0.0427
0.234
0.295
0.211
0.111
U
0.193
0.111
0.234
0.361
0.206
0.105
0.384
0.0419
0.229
0.305
0.269
0.089
0.250
0.0469
0.238
0.489
0.261
0.116
C.
0.177
0.0944
0.184
0.174
0.073
0.093
Results. Following Huang al. 15 say model has the required counting power if the error below 0.01 and report the result in Table We compare to MPNN baseline [32, RNM GNN, that MAG- GNN successfully counts cyccls, which indicates that node-marking marking also helps non-marked nodes to count cycles. We also note that MAG-GNN does not count 6-cycles well, although > 2-order node tuples. We suspect this is because MAG-GNN takes the average score improvement of nodes the reward, which might not be the best reward for node-level task. Even so, observe MAG-GNN's performance improvement over NGNN, which shows that MAG-GNN with larger node tuples indeed increases expressivity. We leave the design of more proper node-level reward to future works.
node-level subgraph GNNs 33 37 35], and edge-level subgraph GNNs [15] We
still
6.3 Real-world datasets
Datasets. To answer Q3, we adopt the following real-world molecular property prediction datasets: (I) QM9 (MAE) con- tains 130k molecules for twelve graph regression targets. (2) ZINC and ZINC-FULL (MAE) include 12k and 250k chem- ical compounds for graph regression. (3) OGBG-MOLHII
Table 3: Transfer learning.
ZINCJ ZINC-FULLJ MOLHIVI
No PRE
0.106
0.030
77.12
3-CYCLES
0.104
0,030
76.98
(AUROC) contains 41k molecules, and the goal is to classify whether molecule inhibits HIV. We use the SIMUL training scheme for a fair comparison to other methods.
6-CYCLES
SR25
0.103
0.028
77.26
Table 4: Inference time.
Results. On QM9 targets, MAG-GNN significantly outper- forms NGNN on all targets with an average of 33% MAE reduction. It also outperforms I2-GNN, with partially >3-WL expressivity, on most of the targets (16% average MAE reduc- tion), showing that with much fewer MPNN runs, MAG-GNN can still achieve better performance. This is because despite using fewer subgraphs, we use node tuples with size greater than two, which grants MAG-GNN the power distinguish graphs that require higher expressivity. MAG-GNN also performs comparably to 1-2-3-GNN, which simulates 3-W We observe that 1-2-3-GNN performs well on the last five targets (which are global properties hard for subgraph GNNs) while outcompeted by MAG-GNN on the rest of the targets. We suspect that 1 1-2-3-GNN has constant high expressivity, which might easily lead to overfitting during training. At the same time, MAG-GNN can automatically adjust the expressivity by subgraph selection, reducing the risk of overfitting.
Time ZINC-FULL CYCLE
QM
MPNN
100.1
58.4
222.9
NGNN
402,9
211.7
776.8
PGNN
1864.1
1170.4
3524.0
PPGN
2097.3
1196.8
4108.7
MAG-GNN
385.8
155.1
704.9
The results on other molecular datasets are shown in Table 6 We see that MAG-GNN outperforms base GNN by large margin showing its better expressivity. Another important comparison is between MAG-GNN and other non-equivariant GNNs, including PF-GNN, K-OSAN, and RNM. We see that MAG-GNN achieves significantly better results on ZINC, where only models with high expressivity get good performance. This verifies that using an RL agent to capture the inter-subgraph relation is essential in finding expressive subgraphs. MAG-GNN does not perform well on the

---

Table 6: Molecular datasets results.
# Params
ZINC (1)
ZINC-FULL (4)
OGBG-MOLHIV (1)
GIN
0.163+0.004
0.088:10.002
77.07=1.49
PNA
0.188=0.004
7..05=1.32
O-OAAN
0.155=0.004
PF-GNN
0.122==0.010
80.15=0.68
RNM
453k
0.128+000227
0.062++0.004
76.79+0094
GSN
0.115=0.012
7..8
CIN
-00k
0.009=000006
0.022+00.002
00.94+0.57
NGNN
-50k
0.111=0.003
0.029
78.34
GNNAK+
-500k
0.080=00001
1.16
SUN
526k
0.083=00003
00.5
KPGNN
489k
0.093+00007
PGNN
0.083++0.001
0.023+0.001
78.68
SSWL+
387k
0.000=00.005
0.022+00002
79.58
MAG-GNN
496k
0. 0.106+0..014
0.030+0.002
1..12+1.13
MAG-GNN-PRE
496k
0.096+000009
0.023+00.022
78.30+
OGBG-MOLHIV dataset. We observe that I-WL bounded methods PNA also achieves good results on datasets, meaning that the critical factor determining the performance on this dataset likely the implementation of the base GNN but not the expressivity. MAG-GNN is highly adaptable to any base GNN, potentially improving MAN'SNN's performance further. We leave this to future work.
We use the PRE training scheme to conduct transfer learning tasks on ZINC, ZINC-FULL, and OGBG-MOLHIV datasets. We pre-train the expressivity datasets shown on the left column of Table 3 and train the attached MPNN head using the datasets on the top row. We see that pretraining consistently brings performance improvement all datasets. Models pre-trained on CYCLE are generally better than the one pretrained on SR25, possibly due to the abundance of cycles in molecular graphs.
6.4 Runtime comparison
We conducted runtime analysis on previously mentioned datasets. Since is difficult match the number of parameters for all models strictly, we fixed the number of GNN layers to five and the embedding dimension to 100. We a GB memory budget for all models and measured their inference time on the test datasets. We use m = and T = 2 for MAG-GNN, The results in Table 4 show that MAG-GNI is more efficient than all subgraph GNNs and is significantly faster than the edge-rooted subgraph GNN, GNN.. NGNN achieves comparable efficiency to MAG-GNN because takes a fixed-hop subgraph around nodes, reducing the subgraph size. However, MAG-GNN outperforms NGNN most targets with its better expressivity.
Limitations. Despite the training schemes in Section 4.1 MAG-GNN is harder train. Also, MAG-GNN's design for node-level tasks might not be proper. This motivates research on extending MAG-GNN to node-level even edge-level tasks. We discuss this further in Appendix C.4
7 Conclusions
In this work, we closely examine one popular GNN paradigm, subgraph GNN, and discover that a small subset of subgraphs is sufficient for obtaining high expressivity. We then design MAG-GNN, which uses RL to locate such subset, and propose different schemes to train the RL agent effectively. Experimental results show that MAG-GNN achieved very competitive performance to subgraph GNNs with significantly less inference time. This opens up new pathways to design efficient GNNs.
Acknowledgement. Lecheng Kong, Jiarui Feng, Hao Liu, and Yixin Chen are supported by NSF grant CBF-2225809. Muhan Zhang is partially supported by the National Natural Science Foundation of China (62276003) and Alibaba Innovative Research Program.
10

---

References
[1] Ralph Abboud, Ismail Ilkan Ceylan, Martin Grohe, and Thomas Lukasiewicz. The surprising power of graph neural networks with random node initialization. arXiv preprint arXiv:2010.01179, 2020.
[2] Pietro Bongini, Monica Bianchini, and Franco Scarselli. Molecular generative graph neural networks for drug discovery. Neurocomputing, 450:242-252, 2021.
[3] Giorgos Bouritsas, Fabrizio Frasca, Stefanos Zafeiriou, and Michael Bronstein. Improving graph neural network expressivity via subgraph isomorphism counting. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022.
[4] Jin- Yi Cai, Martin Fürer, and Neil Immerman. An optimal lower bound on the number of variables for graph identifications. Combinatorica, 12(4))389-41) 410, 1992.
[5] Zhengdao Chen, Soledad Villar, Lei Chen, and Joan Bruna. On the equivalence between graph isomorphism testing and function approximation with gnns. Advances in neural information processing systems, 32. 2019.
[6] Sikai Chen, Jiqian Dong, Paul Ha, Yujie Li, and Samuel Labi. Graph neural network and reinforcement learning for multi-agent cooperative control of connected autonomous vehicles. Computer-Aided Civil and Infrastructure Engineering, 36(7):838-83) 857777,,888 2021.
[7] Mohammed Haroon Dupty, Yanfei Dong, and Wee Sun Lee. Pf-gnn: Differentiable particle
filtering based approximation of universal graph representations. In International Conference Learming Representations, 2021.
[8] Jiarui Feng. Yixin Chen, Fuhai Li, Anindya Sarkar, and Muhan Zhang, How powerful are k-hop message passing graph neural networks. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[9] Matthias Fey and Jan Eric Lenssen. Fast graph representation learning with pytorch geometric. arXiv preprint arXiv:1903.000 020428, 2019.
[10] Fabrizio Frasca, Beatrice Bevilaequa, Michael Bronstein, and Haggai Maron. Understanding and extending subgraph gnns by rethinking their symmetries. Advances in Neural Information Processing Systems, 35:31376-311 376-31390, 2022.
[11] Yang Gao, Hong Yang. Peng Zhang, Chuan Zhou, and Yue Hu. Graphnas: Graph neural architecture search with reinforcement learning. arXiv preprint arXiv: htt:::904.099.. 2019.
[12] Justin Gilmer, Samuel S Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference machine learning, pages 1263-1272. PMLR, 2017.
[13] Shengding Hu, Zheng Xiong, Meng Qu, Xingdi Yuan, Marc-Alexande Côté, Zhiyuan Liu, and Jian Tang. Graph policy network for transferable active learning on graphs. Advances in Neural Information Processing Systems, 33:10174-10185 2020.
[14] Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in neural information processing systems, 33:22118 11-12133, 2020.
[15] Yinan Huang, Xingang Peng, Jianzhu Ma, and Muhan Zhang. Boosting the cycle counting power of graph neural networks with i2 P-gnns. arXiv preprint arXiv:22 /3978, 2022.
[16] Lecheng Kong, Yixin Chen, and Muhan Zhang. Geodesic graph neural network for efficient graph representation learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[17] Kwei-Herng Lai, Daochen Zha, Kaixiong Zhou, and Xia Hu. Policy-gnn: Aggregation op- timization for graph neural networks. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 461-471, 2020.
11

---

[18] Yujia Li, Chenjie Gu, Thomas Dullien, Oriol Vinyals, and Pushmeet Kohli. Graph matching networks for learning the similarity of graph structured objects. In International conference on machine learning, pages 3835-3845. PMLR, 2019.
[19] Long-Ji Lin. Reinforcement learning for robots using neural networks. Carnegie Mellon University, 1992.
[20] Haggai Maron, Heli Ben-Hami Hadar Serviansky, and Yaron Lipman. Provably powerful graph networks. Advances in neural information processing systems, 32, 2019.
[21] Karolis Martinkus, Pál András Papp, Benedikt Schesch, and Roger Wattenhofer. Agent-based graph neural networks. In The Eleventh International Conference on Learning Representations, 2023.
[22] Federico Monti, Karl Otness, and Michael M Bronstein. Motifnet: a motif-based graph convolutional network for directed graphs. In 2018 IEEE Data Science Workshop (DSW), pages 225-228. IEEE, 2018.
[23] Christopher Morris, Martin Ritzert, Matthias Fey, William L Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe Weisfeiler and leman go neural: higher-order graph neural networks. Proceedings of the Thirty-Third AAAI Conference on Artificial Intelligence and Thirty- Innovative Applications of Artificial Intelligence Conference and Ninth AAAI Symposium on Educational Advances in Artificial Intelligence, pages 4602 2602-460,,55 2019.
[24] Pál András Papp. Karolis Martinkus, Lukas Faber, and Roger Wattenhofer. Dropgnn: Random dropouts increase the expressiveness of graph neural networks. Advances in Neural Information Processing Systems, 34:21997-22009, 2021.
[25] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[26] Chendi Qian, Gaurav Rattan, Floris Geerts, Mathias Niepert, and Christopher Morris. Ordered subgraph aggregation networks. Advances in Neural Information Processing Systems, 35:21030- 21045, 2022.
[27] Qingyun Sun, Jianxin Li, Hao Peng, Jia Wu, Yuanxing Ning, Philip S Yu, and Lifang He. Sugar: Subgraph neural network with reinforcement pooling and self-supervised mutual information mechanism. In Proceedingg of the Web Conference 2021, pages 2081-2091, 2021.
[28] Yanbo Wang and Muhan Zhang. Towards better evaluation of gnn expressiveness with brec dataset, 2023.
[29] Minjie Wang, Da Zheng, Zihao Ye, Quan Gan, Mufei Li, Xiang Song, Jinjing Zhou, Chao Ma, Lingfan Yu, Yu Gai, Tianjun Xiao, Tong He, George Karypis, Jinyang Li, and Zheng Zhang. Deep graph library: A graph-centric, ggghly-performatt package for graph neural networks. arXiv preprint 1rXiv:190.001315. 2019.
[30] Hanrui Wang, Kuan Wang, Jiacheng Yang, Linxiao Shen, Nan Sun, Hae-Seung Lee, and Song Han. Gen-rl circuit designer: Transferable transistor sizing with graph neural networks and reinforcement learning. In 2020 57th ACM/IEEE Design Automation Conference (DAC), pages 1-6. IEEE, 2020.
[31] Shiwen Wu, Fei Sun, Wentao Zhang, Xu Xie, and Bin Cui. Graph neural networks recom- mender systems: survey, ACM Computing Surveys (CSUR), 2020.
[32] Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations, 2018.
[33] Jiaxuan You, Jonathan Gomes-Selman, Rex Ying, and Jure Leskovee. Identity- graph neural networks. CoRR, abs/2101.10320, 2021.
12

---

[34] Muhan Zhang and Yixin Chen. Link prediction based on graph neural networks. Advances in Neural Information Processing Systems, 31:5165-5175, 2018.
[35] Muhan Zhang and Pan Li. Nested graph neural networks. In A. Beygelzimer, Y. Dauphin,
P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[36] Bohang Zhang, Guhao Feng, Yiheng Du, Di He, and Liwei Wang. A complete expressiveness hierarchy for subgraph gnns via subgraph weisfeiler-lehman tests, 2023.
[37] Lingxiao Zhao, Wei Jin, Leman Akoglu, and Neil Shah. From stars to subgraphs: Uplifting any gnn with local structure awareness, arXiv preprint arXiv:2110.03755, 2021.
[38] Lingxiao Zhao, Neil Shah, and Leman Akoglu. A practical, segrrssivvel-exppeeseeee gnn. Advances Neural Information Processing Systems, 35:34106-341 2022.
13

---

Appendix
A Proof of theorems
A.1 Proof of Theorem
To prove theorem D we leverage the following lemma from [35,, and we restate it here.
Lemma 1. For two graphs G1 and G2 that are uniformly independently sampled from all n-node r-regular graphs, where S V2log ni, we pick any nodes, each from one graph, denoted by on and 12 respectively, and do [(2 + +.) log(r-1-c) 1-height rooted subgraph extraction around ii and log log(r-1-e) many layers, a proper message passing GNN will generate different
logn
e2- With most representations for the extracted two subgraphs with probability at least 1 - o(n
For completeness, recall Theorem
Theorem 1. Let G and G2 be two graphs uniformly sampled from all n-node, r-regular graphs where 3 S r /2 n. Given an injective pooling function R(O) and an MPNN gp of I-WL expressivity, with a probability of at least 1 - o(1), there exists node (tuple) v E ( whose corresponding node marked graph's 's embedding, fr1 v). is different from any node marked graph" embedding in G2.
Proof. Since the pooling function R(G) is injective and the GNN g has 1-WL expressiveness, which is the maximal expressiveness for message message-passing GNNs, fo satisfies the proper GNN requirement
in Lemma
We also let g have < [ox[r-1-e] + d layers, where d [(} + <) loR(r-T-a) 1o="- Let
log
(v1) on G1, since 21 is marked, we can have the first d layers of perform greadth-first-saare on the graph rooted from 11, and afterward, each node's 's representation embeds whether it's within d
distance from v1. We then let the last <
log(r-I-<) layers of only to perform message passing on the nodes that within d distance from U1. This can be easily done by implementing an injective COMBINE function. We do the same operation for v' = (22) on Gz. According to lemmal] [ since fr proper, the probability that (-(G1, # (G,, v) is at least o(n-1) Then, by union bound, the probability that (f & (fr(G2, v o')l E (VG2) is at least o(1)
log
A.2 Proof of Theorem
We first introduce two non-isomorphic graphs in Figure 6
B
Figure 6: Two CSL graphs.
They are CSL graphs both with nodes, one with skip factor of 2 and the other with a skip factor of They have the following properties,
MPNN cannot differentiate them since they are regular graphs.
They are vertex-transitive graphs. Meaning that all node-rooted subgraphs from one graph are the same.
14

---

The rooted subgraphs from the two graphs are different and can be differentiated by an MPNN.
The properties imply that we cannot use an MPNN to differentiate the two graphs, but we can apply MPNN once on one subgraph from each graph to differentiate them since all subgraphs for one graph are the same.
We name the CSL graphs A and B and them to build super graph. Let d 2n be super graph's 's diameter where n E zt. We randomly permute type A graphs and n type B graphs to form graph sequence S. An example is Se (A, B, B, A, A, B).
We then define S" as the inverse of S, that
A if [S]
[s)]
(11)
if [S] =
The inverse Se, S (B, A, A, B,
We can then use concatenation to form two new sequences.
St = S =Se
(12)
and S (A, B, B, A, A, B, B, A, A, B, B, A). We then use full connections to connect the graphs in sequence and create a super graph G(S), specifically,
For example, S. =
G(S) (VG)),
V(G) = (.S()))), B(G) = ((((()))))))) ( v([S)) x (((S))-
(13)
The nodes in graph S[i] connect to every node in graph Sli) if the two graphs are adjacent in the sequence or are the head and tail of the sequence. Figure shows the example super graphs G(S and G(S,)
G(S
G(S
B
B
Se
S
S
Full connection between graphs
B
B
CSL graphs
A
Figure 7: Two example super graphs.
The properties of G(S) is.
The diameter of the super graph is 2n.
G(S) has an equal number of graphs and B.
Because adjacent graphs in the super graph are fully connected, adding node marking to graph [S]; in the super graph does not help us identify any other graphs [5]; without node marking. To identify [S] must add node marking to [S]5-
For positive super graphs, [S]; [5]+ that is, for any graph in the super graph. the graph
on its direct opposite shares its type. Whereas for negative super graphs, [S]; * [5]
Let the task be to discriminate between positive super graphs and negative super graphs. We then prove that MAG-GNN superior in this case to show strictness.
Proof. We consider the case where k: = 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 the graph that is farthest away from the other node. This is easy to be implemented by MAG--NN since
15

---

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 GSS) 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 4n-1- Consequently, the expected number of random actions required to identify the super graph is
1
4n
A.3 Proof of Theorem
We first describe the PF-GNN algorithm; we refer the readers to Dupty et al. 7 for more details.
PF-GNN starts with set of particles for a graph G,
((G, HT") " Wi/mml:MM
(14)
where H is the node embedding of G using MPNN g. For particle m in step t. PF-GNN uses policy function P(VIH.") to compute 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 ~ P(VIH(") [HT"), MLP((HH))))) In), HHt = g ((HT))
(15)
Then PF-GNN updates the particle weights as follows,
fobo (H1+1) w.
(16)
c+1
Σn fobe (H wx""
where Joha 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 (u]42)m=l::h and reset weights to 1/M. Eventually, all H 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 repre- sentation 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 M-particle, T -tepp 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- where [W], m, = 0, i + T+1 indicates the i-th node in node tuple m has not been individualized. [W]m, m, = j. + T + 1 indicates the i-th node in node tuple m is individualized at time step j. and [W] M is the weight in PF-GNN
Let (U1, ...: Ut) be the sequence of nodes that PF-GNN used individualize. We first note that there exists an MPNN applied onto the node marked graph G(v) with a sufficient number of layers that output the same value the final output, H,, of PF-GNN. Specifically, let T be the number of layers of in the MPNN, and then we can build a &T' + t-layers MPNN gp 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 fw to log the step when a node vú in the tuple m is marked that is [W1m,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. fw also update [W) mm to be equal the weight after the resampling process. Then, by using the IT + f-layers MPNN consists of as 9p and another iT + 1-layers MPNN consists of the MPNN of H()) as grt, MAG-GNN can replicate the output of PF-GNN.
L
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, 7).. where:
16

---

S is set of states.
is a set of actions.
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 particular state. 7 is discount factor determining future rewards' importance.
Agents in an MDP aim to learn 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 Qs,, a), representing the expected cumulative rewards of taking action in state During training the Q-values are updated by,
Q(s, a) Q(s a) + a(r mma))=mmmsmsmmmmmc ,d)) Q(s,
(17)
where a is the learning rate, r is the instant reward of action 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 (s,, for all possible actions A loss function to train the Q-network is,
L [r(r
max Q(s', a') Q(s,a))
(18)
widely adopted method to train the Q-network the Experience Replay. In each step, the RL agent interacts with the environment to store experience tuples E = (s, o, s'), and sample a set of previous experiences to train the Q-network. This technique effectively improves the sample efficiency. We adopt e-greedy training strategy. where the agent takes random actions with probability of e. We start with e = and gradually decrease zero. This ensures that the agent explore the graph and learns the good policy.
C More discussion on the design of MAG-GNN
C.1 Action space
In Section we propose an action to substitute one node in the node tuple with another node. Such action requires a O(m[V 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 becomes
(19) The action space becomes A [k] x V. In such case, we need to enforce an order to update each node tuple sequentially simultaneously update them. In the experiments, we choose update all node tuples for better parallelizability simultancously.
o 0,,1
-.,,,,,-., uI [w:ls+1, ....t...../// (b:lk)
A seemingly
natural design choice is directly replacing the current node tuple with another one. However, it has O(V)*)) possible actions for one node tuple. Its corresponding Q-table size is also exponential to k. The cost of computing such table may equal that of running the complete subgraph GNN, which violates our motivation for complexity reduction.
The state matrix W serves tracker of previous states. The choice of update function fw can vary. can simply
mean pooling function on the embeddings marked nodes. can also be an LSTM pooling function. While using matrix W has good theoretical properties 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 each step the agent obtains experiences from the environment, stores the experiences into memory, and retrieves other experiences from memory to optimize the
17

---

Algorithm RL-Exxperience
Input: Node tuples U, batch of data B, buffer D, RL agent and GNN Drt
Output: New node tuples.
1: Use and to generate a new node tuples U' according to equation (9) and (10)
2: Obtain rewards r from U", U, and B according equation (8)
3: Store the experience tuple (U", UU, Brr) D
4: return U
Algorithm 2 ORD-Trrii
Input: A target dataset of (9,3). A graph prediction model fp. RL agent with RL GNN gri- Output: A trained RI. agent that output the most discriminative node tuples and corresponding prediction GNN fz-
1: while fp is not converged do
2: B < batch from (G.3)
3: Sample random node tuples U for graphs in B
4: Optimize fy according equation (4) w.r.t B and V.
5: end while
6: D +- memory buffer
7: Call RL-Experience until D has enough memory
8: while 9r1 not converged do
9: B <-a batch from (9.3))
Sample random node tuples U for graphs in B
11: RL-Exxerience(U,, B, D, a, 9r!)
12: Sample experience tuples E from
13: Optimize 9r1 using the E
14:
end while
15: return Sp. a, 9rl
Q-learning agent. Since 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 train the 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 environment training step 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-code for the paradigms.
C,3 Other sampling- based methods
The advantage of MAG-GNN over other subgraph sampling methods, including PF-GNN [ and DropGNN [24], is that it effectively chooses the most discriminative subgraphs. DropGNN drops nodes with equal probability; PF-GNN uses 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 there is still 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 refine the likelihood of 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 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,
18

---

Algorithm 3 SIMUL-Train
Input: A target dataset of (G.3). A graph prediction model fpr an RL agent with RL GNN 9r!- trained RI. agent that output the most discriminative node tuples and corresponding prediction GNN fp-
Output:
t- memory buffer
2: Call RL-Experience until D has enough memory
3: while fy and grt are not converged do
4: B <-a batch from (G,3)
5: Sample random node tuples U for graphs in B
6:
E'pree-xpeeieence(U B, D, o, gri)
7:
Optimize fp according to equation (4) w.n. B and U
8:
B' + batch from (G. (G)))
9:
Sample random node tuples U for graphs in B'
10:
RI- B, D, a, 9r1)
11:
Sample experience tuples E from D
Optimize grt using the E'
13:
end while
return fp. a, grl
Algorithm PRE-Train
Input: A target dataset of (g, (). A graph prediction model fp. an RL agent a with RL GNN grl trained with either SIMUL ORD on a different dataset.
Output: A trained RL agent a that output the most discriminative node tuples and corresponding prediction GNN fz-
1: while fp is not converged do
2: B < batch from (9,3)
3: Sample random node tuples U for graphs in B
4: ""=RL-Exxerieenee(, B, D, a, gri)
5: Optimize fp according to equation (4) w.rt B and U
6: end while
7: return fp. a, grl
and we can adopt more robust training methods from the large RL toolbox to improve training time. We leave this to future work.
advantage of subgraph GNN its equivariance property. Given the same graph up 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-basee and andomization-base
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 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- tasks. We conduct the evaluation in Section 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 identify such node tuple. We suspect this because the MAG-GNN's reward is defined at graph-level, and taking plain average of node error does not faithfully capture the environment. This leads exciting research paths extending MAG-GNN to nove-level or even edge-level tasks.
19

---

Table 7: BREC dataset results (t)
Basic (60)
Regular (140)
Model
Number
Accuracy
Number
Accuracy
NGNN
98.3
48
34.3
SSWL+
36
60
100
50
35.7
PGNN
15
60
100
100
71.4
OSAN
26
56
93.3
8
5.7
MAG-GNN
49
81.6
61
43.5
D Complexity analysis
The complexity of MAG-GNN is straightforward. It runs MPNN multiple times; hence, the complex- ity is multiple of MPNN (1).. Specifically, for a MAG-GNN with m k-tuples that runs steps, there is total number of O mt) MPNN runs, incurring ( total complexity. The complexity computing the Q-table equal the size of the Q-table, which is Q(m[[]). We compute Q-tables ((t) times, and the overall complexity for computing the Q-table is V() The overall complexity (mmk/[V] + V2)) In practice, ( < O(T)V)) as can be observed in Table 10 and hence the complexity can be reduced to O(mtT) Note that subgraph GNNs have the complexity of V(T which grows exponentially with k, while MAG-GNN's complexity grows linearly with even when is large.
In equation we propose to use graph embedding function fp to represent the current state; however, alternative is use frt describe the current state. This allows embedding for prediction but only compute that for Q-table. This does not change the complexity but reduce the number of MPNN runs in half, improving efficiency.
to compute graph
E More experimental results
0.05
0.00
-0.05
-0.0
10
20
30
40 50o
Training step (k)
Figure 8: Reward increases over time.
Table 7 shows the results on the BREC dataset's basic and regular graph subsets (88 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 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 PCNNN, MAG-GNN is considerably more efficient. We also noticed that MAG-GNN significantly improves the performance regular graphs over OSAN, further validating that RL-selected subgraphs carry more information.
20

---

Move markings apart
Agent stagnant
Common neighbors w/ edge
Move marking apart
Agent keeps moving Common neighbors w/o edge
Figure 9: MAG-GNN example 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 g 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 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 0.05. We see 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 future work.
We conducted the experiments 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 MAG-GNN with 2-node markings to break the 3- W- expressivity bound. After training, the MAG-GNN generates the trajectory given the initial I-distance random node markings as shown in Figure We can observe that MAG-GNN first learns to push the two marked nodes apart in both graphs. However, in the Shrikhande graph, after 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 downstream MPNN.
F
More experiment details
F.1
Datasets statistics
repository. All training and testing are on standardized targets following the previous practice [77 15]. The task is regression, and we use continuous value to approximate the discrete number of cycles. We use the QM9 dataset provided by -rtorch-Geometric [9], and we use a train/validdtestt 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 eytorch-Geomett [9] and use the official split. We take OGBG-MOLHIT dataset from the Open Graph Benchmark package [4] and use their official split. For CSL and EXP, follow
We took the CYCLE dataset from 37
21

---

Table 8: Hyperparameters Summary
Method
QM9
ZINC
ZINC-FULL CYCLES
Learning rate
0.001
Weight decay
0
0
0
0
# GNN Layers
5
6 Sum
6
Jumping Knowledge
Sum
Sum
Sum
(2,3, 4) (3,4,5,6)
(2,3,4) (3,4,5,6)
(2,3,4] 4] {3, (3,4,5,6)
(2,3, 4) (3, 63 2 5000
k
# Epochs
2000
4000
500
Batch size
512
512
512
500
Memory size
3M
300K
300
300K
Sync rate
3000
3000
3000
3000
# runs
4
10
3
10
Table 9: Hyperparameters Summary (cont.)
Method
MOLHIV
EXP
SR25
CSL
I Learning rate Weight decay
0.001
0.001
0
# GNN Layers
6
4
4
4
Jumping Knowledge
Sum
None
None
None
32 4] 63
[1,2,3] [1,2,3]
[1,2,3}
{1, 2|
k
[1,2,3]
[1,2,3]
2
4
4
# Epochs
1000
500
10000
500
Batch size
128
64
32
64
Memory size
300K
30K
30K
30K
Syne 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 only 15 different graphs, we follow previous practices and 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 & and 9
Curly bracket means the range of hyper-parameters we search. Memory size is the size the memory buffer for experience and replay. Syne rate is the number of optimization steps before the target Q-network syne with the train Q-network. For ZINC datasets, we additionally use learning rate decay with rate of 0.5 for every 200 epochs. All models implemented in DGL [29] and Py Porch [25] We only use Pytorch Geometric [9] to get the datasets.
Since there is 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.
22


Reviews:
[{"review_id": "KFrUPf70Et", "reviewer": "Reviewer_egpV", "summary": "The paper introduces MAG-GNN, an RL approach for identifying important subgraphs. This is valuable since models using fewer subgraphs incur lower inference costs. Experimentally, the authors show that MAG-GNN retains the strong performance of sub-graph approaches while reducing the run time.", "questions": "\nCould the authors report the training time of MAG-GNN? Knowing the timeframe compared to baselines would be helpful, given that the training of MAG-GNN is costly.\n\nI am wondering if it is possible to visualize the discarded and remaining subgraph structures resulted from RL. Given the complexity and freedom of tuning of the reported RL approach, such visualizations could aid in understanding which subgraph structures are useful and which are not. In particular, it will be informative to apply this visualization technique to toy-like graphs, similar to those in Figure 1, where we already know which subgraph structures matter. ", "limitations": "See the \"Weaknesses\" section above.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The guiding question of the paper \"Are a small number of subgraphs sufficient for graph problems?\" is well-motivated. The proposed RL-based approach for identifying subgraphs is well-explained. The approach's performance and efficiency are validated on various benchmark datasets. The negative results of the node-level task are helpful, revealing both the current limit and future potential of the paper.\n", "weaknesses": "**Training efficiency**. While MAG-GNN accelerates inference, the training process becomes more time-consuming. On smaller graphs, this training pipeline may dominate the overall computational costs.\n\n**Task scope**. MAG-GNN excels in graph-level tasks but underperforms in node-level tasks, as the authors gracefully acknowledge.\n\n**Baseline comparison**. The authors primarily compare MAG-GNN with methods using full subgraphs in the experiment section. Considering its efficiency goals, comparing it with other efficient alternatives, such as applying node/edge pruning approaches to the full graph and then using standard subgraph methods, would be useful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer egpV for the insightful review, and we address reviewer egpV’s questions and concerns as follows,\n\n- W1: **Training efficiency**. While MAG-GNN accelerates inference, the training process becomes more time-consuming. On smaller graphs, this training pipeline may dominate the overall computational costs.\n\nReviewer egpV is keen to observe that the critical trade-off in MAG-GNN is between training and inference times. MAG-GNN emphasizes its optimal **inference complexity**, which is **more important** practically, as once trained, the model can be reused on potentially infinite data. The training time for MAG-GNN is less correlated to the graph size, as the training step complexity is linear to the graph size but more correlated to the inherent property of the graphs and the prediction targets. We present the convergence time to obtain the reported performance on ZINC and CYCLE-6 datasets of similar graph sizes.\n\n|  | ZINC | CYCLE-6 |\n| --- | :---: |:---: |\n| NGNN | 36min | 14min |\n| I^2GNN | 127min | 27min |\n| MAG-GNN | 182min | 156min |\n\nNGNN converges much faster than I^2GNN and MAG-GNN, while MAG-GNN has better performance than NGNN on both datasets (see Tables 2 and 4 in the main paper). MAG-GNN is roughly 1.5 times slower in training compared to I^2GNN but much faster for inference, as shown in Table 5 of the main the paper. We note that training MAG-GNN on CYCLE-6 is slower, potentially because the cycle-6 task creates a larger search space for MAG-GNN to find appropriate policies.\n\nWe also want to point out that the training time of the MAG-GNN can be significantly reduced when we use the PRE training scheme, where we first train on a synthetic expressivity dataset and transfer the knowledge to a real-world dataset. After pretraining on CYCLE-6, the convergence time of training the prediction network is **32min** which is even faster than that of NGNN. This shows that the trained agent can be reused and efficiently generalized to downstream real-world tasks, alleviating the high training cost problem.\n\n- W2: **Task scope**. MAG-GNN excels in graph-level tasks but underperforms in node-level tasks, as the authors gracefully acknowledge.\n\nAs the reviewer points out, improving node-level performance is a valuable future direction. A potential solution is to design a better node-level reward function. Meanwhile, we note that the inference time of MAG-GNN on these tasks is more optimal compared to methods that achieve better performance, and effort spent on adapting MAG-GNN to node-level tasks can lead to effective and efficient method.\n\n- W3: **Baseline comparison**. The authors primarily compare MAG-GNN with methods using full subgraphs in the experiment section. Considering its efficiency goals, comparing it with other efficient alternatives, such as applying node/edge pruning approaches to the full graph and then using standard subgraph methods, would be useful.\n\nWe agree such a pruning+subgraph GNN method approach can be an interesting comparison. We did some literature search and didn’t find such existing approaches. Instead,  we provide results of the following experiment where we apply random edge drop with a drop probability p on every edge and then apply NGNN onto edge dropped graphs.\n\n|  | ZINC | Inference Time on 10 tests (ms) |\n| --- | :---: | :---: |\n| p=0.2 | 0.164 | 152.7 |\n| p=0.4 | 0.274 | 137.0 |\n| MAG-GNN | 0.104 | 159.1 |\n\nThe experiments first show that, indeed, by pruning the subgraphs, we can reduce the inference of subgraph GNN. However, we also notice the sharp performance drop of NGNN applied to pruned graphs, meaning that the full graph structure is still critical to optimal performance. MAG-GNN reduces running time while effectively utilizing the full graph structure.\n\n- Q2: I am wondering if it is possible to visualize the discarded and remaining subgraph structures resulted from RL. Given the complexity and freedom of tuning of the reported RL approach, such visualizations could aid in understanding which subgraph structures are useful and which are not. In particular, it will be informative to apply this visualization technique to toy-like graphs, similar to those in Figure 1, where we already know which subgraph structures matter.\n\nWe agree with the reviewer and believe a visualization will be beneficial in understanding the mechanism of the methods. We present two examples in Figures 1 and 2 in the general response's pdf. One example is in EXP dataset using a 2-node-mark. Most graphs in the EXP dataset have two components. MAG-GNN’s most prevalent pattern is that whenever the initial random node markings land on the same component, the agent will move one marking to the other, identifying the graph structure. This is not possible if the node markings only land on one component.\n\nThe other interesting example is when we train on the toy example to distinguish the Shrikhande graph and the Rook 4-by-4 graph. The agent always produces 2-node-markings that are 2-distance away on both graphs. Moreover, the agent always chooses to move one node marking of the Rook 4-by-4 graph to its neighbor, while it chooses to be stagnant in the Shrikhande graph when the common neighbors of the two marked nodes have an edge in between."}, {"review_id": "DrHbpQryJh", "reviewer": "Reviewer_FiDZ", "summary": "This paper proposes a reinforcement learning boosted approach to subgraph graph neural networks (GNNs) that achieves high expressivity with constant complexity, instead of the exponential complexity of existing methods. They demonstrate the effectiveness and efficiency of their approach through experiments that compare it with the state-of-the-art methods.", "questions": "It seems no supporting for this statement: This reduces the exponential complexity of subgraph enumeration to the constant complexity of a subgraph search algorithm while keeping good expressivity", "limitations": "The paper does not discuss limitations or potential social impacts.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written.\n2. The proposed MAG-GNN can achieve good expressivity while reducing the time complexity of subgraph enumeration.\n", "weaknesses": "1. The authors report that most methods in Table 1 achieve 100% accuracy on the dataset. This suggests that the dataset may be too simple to adequately evaluate the expressivity of a model. The authors should consider applying more challenging datasets to better assess the performance of their models.\n2. The table captions are vague and do not provide sufficient information for readers to interpret the experimental results properly.\n3. In Section 3, the authors use the same data for training and test, which may lead to overfitting and unrealistic performance.\n4. The notation |V^2| should be |V|^2.\n5. The proof of the theorem is unclear and lacks details. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the feedback from reviewer FiDZ, and would like to introduce our research's background and intuition further. In traditional deep-learning fields (e.g. CV and MLP), expressivity usually is not an issue, and reducing overfitting is more critical. However, in GNN, the expressivity is limited by 1-WL test [1], preventing the model from distinguishing many different graphs. Hence, a downstream predictor cannot correctly classify these graphs regardless of how the predictor generalizes. Subgraph GNN rises as a solution to the limited expressivity issue, and we observe that as the expressivity increases, the test performance on real-world datasets also increases. However, the drawback of subgraph GNN is its worse runtime because it needs to extract multiple subgraphs to run GNN multiple times. This motivates MAG-GNN that finds the most informative subgraphs to reduce runtime while keeping subgraph GNN's expressivity. We also address the reviewer's questions and concerns as follows,\n\n- W1: Are 100% accuracy datasets too simple?\n\nWe agree that a more challenging dataset will provide a meaningful evaluation of our model. However, since a widely-acknowledged expressivity dataset is missing from the literature, we follow the convention [2,3,4] that evaluates the model on these synthetic datasets to validate the model’s claimed expressivity. Generally, models expressive enough will have a perfect prediction on these datasets, while models with insufficient expressivity will collapse to random guesses, like GIN on the EXP dataset and NGNN on the SR25 dataset. We also provide an evaluation of other real-world datasets, including ZINC, OGBG-MOLHIV, etc., in Tables 2, 3, and 6 to comprehensively evaluate our model.\n\n- W2: Vague table captions.\n\nDue to the space limitation, we keep the table description succinct and describe the experiments in the experimental sections. When more space is permitted in the revision, we will add more detailed descriptions to the table captions.\n\n- W3: In Section 3, the authors use the same data for training and test, which may lead to overfitting and unrealistic performance.\n\nWe follow the convention to train and evaluate the model on the same SR25 dataset [2,4]. The evaluation aims to examine the expressivity instead of the generalizability. Hence the experiment protocol tests if a model can discriminate every graph in the dataset rather than test if knowledge learned from one dataset can be generalized to another.  As the background introduction mentioned, expressivity is particularly important in GNN research, and this motivation section is to show only a few subgraphs achieve similar expressivity. Meanwhile, MAG-GNN’s ability to generalize can be seen in its comparable performance to state-of-the-art methods in real-world datasets.\n\n- W4: The notation |V^2| should be |V|^2.\n\nWe appreciate the reviewer’s careful reading of the paper. Here, we use the power of the set of nodes, $V$, to represent the node tuple set as a convention widely adopted in the literature. For example, $V^2$ means the set of all 2-node-tuples, and $|V^2|$ means the cardinality of such set, while $|V|^2$ is the square of the cardinality of the set of nodes. While both numbers are the same, we believe the former might be more appropriate to our context. Nevertheless, we will explain this in detail in our revision.\n\n- W5: The proof of the theorem is unclear and lacks details.\n\nWe thank the reviewer’s in-depth reading into the proof. We will clarify assumptions and add more details to the proof. Meanwhile, would you kindly let us know which part of the proof causes confusion and we are more than happy to clarify. For readers to understand the proof better, we present the intuition of the proofs.\n\nTheorem 1 is based on the lemma from NGNN [3] that the representation of a subgraph in $G_1$ generated from a GNN is very likely to be different from that of another subgraph in $G_2$. We can then use GNN with more layers to simulate a BFS and the GNN in the lemma. Lastly, we use a union bound to complete the proof.\n\nTheorem 2 constructs super-graphs by two base CSL graphs that 1-node-marked subgraphs cannot distinguish. Hence, the super-graphs’ classes can only be determined when the 2-node-marks land on the opposite side of the super-graph. This scenario can be quickly learned by MAG-GNN but not by random approaches, which proves the theorem.\n\nTheorem 3 uses the fact that as long as we correctly label the order and indices of the selected node-tuples, we can injectively map the output of RL GNN to PF-GNN’s output. We construct a GNN with sufficient layers to realize the mapping.\n\n- Q1: It seems no supporting for this statement: This reduces the exponential complexity of subgraph enumeration to the constant complexity of a subgraph search algorithm while keeping good expressivity\n\nThis is, in fact, the paper's main contribution and can be seen by the better complexity and competitive performance of our model. Specifically, to achieve k-WL expressivity, the number of GNN runs required for subgraph methods is $O(|V|^k)$, which grows exponentially with k. In contrast, MAG-GNN achieves very competitive performance with $O(mt|V|^2)$ (Here $O(|V|^2)$ is the complexity for running MPNN on a dense graph, $mt$ is the number of MPNN runs. See Appendix D for more details), which does not grow exponentially with k. This, in turn, supports the statement.\n\n[1] Xu, Keyulu, et al. \"How powerful are graph neural networks?.\" arXiv preprint arXiv:1810.00826 (2018).\n\n[2] Feng, Jiarui, et al. \"How powerful are k-hop message passing graph neural networks.\" *Advances in Neural Information Processing Systems* 35 (2022): 4776-4790.\n\n[3] Zhang, Muhan, and Pan Li. \"Nested graph neural networks.\" *Advances in Neural Information Processing Systems* 34 (2021): 15734-15747.\n\n[4] Huang, Yinan, et al. \"Boosting the Cycle Counting Power of Graph Neural Networks with I $^ 2$-GNNs.\" *arXiv preprint arXiv:2210.13978* (2022)."}, {"review_id": "YaUOZFSjau", "reviewer": "Reviewer_kfK6", "summary": "this paper proposes a strategy to select the optimal subset instead of extensive enumeration for the subgraph GNN. The proposed method can achieve comparative results and reduce the running time.", "questions": "N/A", "limitations": "N/A", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. the paper is well written and easy to follow.\n2. the experiment part is thoroughly detailed。\n3. the experiment shows the model can achieve the comparative result with subgraph GNN while significantly reducing the running time.", "weaknesses": "N/A", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback and refer the reviewer to our other responses in case the reviewer is interested in more details about our work. We are also happy to answer any further questions the reviewer might have."}, {"review_id": "nsMtCctLN1", "reviewer": "Reviewer_SjU6", "summary": "The paper aims to locate expressive subgraphs for graph representation learning. To do so, the author proposed to use deep Q Learning for training an efficient agent which optimizes the combinatorial problem of choosing the optimal subgraph. The proposed agent-aided GNN shows good performance and is able to transfer to real-world datasets.", "questions": "* In table4, how is no RPE evaluated, why is no RPE on par with training agent on 3-CYCLE 6-CYCLE and SR25?", "limitations": "N/A", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Motivation is interesting and the methodology is well explained. The experiment result seems sound. The proposed model is on par with state-of-the-art while using less computational resources.", "weaknesses": "* How to train an effective agent seems to be a critical point. Despite that, the author has a brief discussion in sec 4.1, more quantitative analysis and illustration on this part can better support the author's choice of ORD.\n* The proposed method seems to be on par with the state-of-the-art. While I don't think absolute performance is an important concern given the method is more efficient. I think it's important to investigate in L316 Q3 how much of the RL agent expressiveness benefit the real-world datasets compared to random ones or human-designed ones.\n\n-------------------------------------------------------------------------------------\nI've read the author's response to me and other reviewers which resolves my concern. I would like to remain my rating as weak accept.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank Reviewer SjU6 for acknowledging our contribution and appreciate the reviewer's suggestions to improve our work. We address the reviewer's concerns and questions as follows,\n\n- W1: How to train an effective agent seems to be a critical point. Despite that, the author has a brief discussion in sec 4.1, more quantitative analysis and illustration on this part can better support the author's choice of ORD.\n\nWe want to clarify that we use ORD and SIMUL training paradigms in different scenarios, as described in the experiment section. We also agree with Reviewer SjU6 that a quantitative study of the effect of different training paradigms is beneficial, and we present the model’s performance on ZINC and SR25 with ORD and SIMUL paradigms (PRE is not considered here because the agent is not trained during fine-tuning stage).\n\n|  | ZINC | SR25 |\n| --- | --- | --- |\n| SIMUL | 0.104 | 100 |\n| ORD | 0.131 | 100 |\n\nWe also present the training time on these scenarios:\n\n|  | ZINC | SR25 |\n| --- | --- | --- |\n| SIMUL | 182min | 204min |\n| ORD | 109min | 141min |\n\nORD usually leads to faster convergence because of its stability. However, because the prediction GNN could be ignorant of the node labels if trained without the agent, as discussed in Section 4.1, the agent trained in the ORD paradigm sometimes only minimally improves the prediction GNN, leading to worse performance than SIMUL.\n\n- W2: The proposed method seems to be on par with the state-of-the-art. While I don't think absolute performance is an important concern given the method is more efficient. I think it's important to investigate in L316 Q3 how much of the RL agent expressiveness benefit the real-world datasets compared to random ones or human-designed ones.\n\nWe thank reviewer SjU6 for emphasizing the comparison between MAG-GNN and the random approach, and we believe such evaluation is essential to understand the realistic utility of the model. We refer the reviewer to the general response GR1 where we provide a full comparison between MAG-GNN and random approach varying m and k on real-world datasets.\n\nWe show that MAG-GNN outperforms the random approach under all circumstances, meaning that the node tuples that MAG-GNN finds are more discriminative.\n\n- Q1: In table4, how is no RPE evaluated, why is no RPE on par with training agent on 3-CYCLE 6-CYCLE and SR25?\n\nNO PRE is evaluated without any pretraining on the synthetic datasets. It is essentially the SIMUL training paradigm, and we use NO PRE to emphasize the comparison. NO PRE achieves close performance because the underlying architectures of the NO PRE and pre-trained models are the same. They differ only in the parameter values, meaning a model without pretraining can potentially achieve the same performance. The PRE model aims to help the RL agent recognize better high-level patterns in a noise-free environment, such as the synthetic high expressivity datasets. Meanwhile, the numbers might be small, but the MSE reduction for ZINC and ZINC-FULL is 9.43% and 23.3%, respectively, which is a considerable performance improvement, showing that the PRE paradigm helps generalize high-level graph features to downstream tasks."}, {"review_id": "8Tr04oEoj1", "reviewer": "Reviewer_jrFK", "summary": "The paper considers the problem of increased complexity of subgraph GNNs with respect to standard MPNNs, due to the enumeration of the entire bag of subgraphs for every given graph.\n\nThe authors first study the need of enumerating all subgraphs, and show that there exist graphs where using a small subset of subgraphs have the same discriminative power as the entire bag of subgraphs. Motivated by these cases, the authors propose an RL-based approach, named MAG-GNN, to find the best subset of subgraphs to be employed in the task at hand.\n\nThe method starts with a subset of subgraphs that are randomly selected. Then, it employs a DQN framework, where the Q-network is parametrized by a subgraph GNN, to replace each subgraph in the subset with a another subgraph. The new subgraph is obtained by replacing only one index of the node tuple identifying the subgraph. Upon reaching a maximum number of steps, the obtained subgraph set is passed through a downstream subgraph GNN for the final predictions.", "questions": "1. **What is the performance of _your prediction subgraph GNN_ (which is built on k-OSAN) when trained with random subsets of subgraphs (randomly selected at every forward)**? Note that I am asking to keep your downstream prediction subgraph GNN formulation identical to what you use in MAG-GNN (except the hyperparameter choices) in order to truly understand the impact of the RL agent. Please use $m$ and $k$ up to the values you consider for MAG-GNN.\n2. Please include comparison with subgraphs GNNs (Bevilacqua et al 2022, Zhang et al 2023) and with k-OSAN with IMLE.\n3. Table 7 shows that MAG-GNN with $m=1$ (and ensemble 3) achieves practically the same performance of the MAG-GNN that uses much higher $m$. Please investigate this aspect. Does it imply that MAG-GNN is effectively only learning to select one subgraph, and the other ones are unimportant (as it can get similar performances when choosing one or many)?\n4. I think $f_p$ in Equation 9 should be $f_{rl, p}$. Indeed, using $f_p$ means that you take the representation of graph obtained through the prediction subgraph GNN. However, from the code I think you take the representation of graph obtained through the subgraph GNN that parametrizes the Q-network (that is based on $g_{rl}$). Please clarify.\n5. In the SIMUL case, do you reset the replay buffer every time you train your downstream prediction subgraph GNN $g_p$?\n6. What is the performance of _your prediction subgraph GNN_ when using the full bag and $k>1$ (that is, for various values of $k>1$)?", "limitations": "n/a", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The problem of reducing the complexity of subgraph GNNs without trading off performances is significant. The authors identify graphs where a small subset of subgraphs has the same discriminative power as the full bag of subgraphs, which motivates the search for such subsets.\n2. The combination of RL with subgraph GNNs is original, and represent a natural way of learning to select subgraphs.", "weaknesses": "1. **The paper misses a thorough comparison with a random policy**, where a random subset of subgraphs is randomly selected at every forward pass. The authors consider tangentially the question in the experimental evaluation (Q1) but fail to answer in details. For every dataset you consider, there should be a table entry reporting the performance of the random baseline. **Such baseline should have the same $k$ and the same $m$ that you consider**. Furthermore, it should be based on the same architecture that MAG-GNN uses as downstream prediction model.  Only with such comparison we can conclude that MAG-GNN performs better than random. Note that:\n    - The answer of Q1 is insufficient, both because you consider only one dataset, and also because you focus only on one node tuple ($m=1$). Please include a table entry for every dataset. \n   - The experiment in Appendix E reports a comparison with a random baseline that uses ensemble predictions. However, it would make more sense to aggregate the representations of the subgraphs and _then_ perform a prediction, an approach that is indeed followed by subgraph GNNs (as reported in Equation 3). The ensemble method instead predicts using each subgraph separately and then aggregates the predictions, which is clearly less powerful. Please include a comparison with random that does not use the ensemble approach, but that aggregates the representations of the $m$ subgraphs to obtain a graph representation, and then performs the prediction. Note that you should evaluate the same $m$ and $k$ values that you sweep on for MAG-GNN (up to 4 and 6 respectively).\n\n2. **The paper misses comparisons with the I-MLE approach proposed in $k$-OSAN** to select subgraphs, which has exactly the same goal.\n\n3. **The paper misses comparison with many subgraph GNNs**, such as Bevilacqua et al 2022, Zhang et al 2023. These methods obtain better performances than MAG-GNN for example on ZINC, using node-based subgraphs (where the bag of subgraph is of size $\\vert V \\vert$, and therefore are tractable). It is unclear what is the advantage of training an RL agent, which has much higher training time, and using $k$-order policies with $k>1$, if the performances are worse than node-based subgraph GNNs and the inference time is not better.\n\n4. The comparison on the ZINC dataset is unfair because the model does not respect the parameter budget imposed by the dataset (500k). Currently he budget is twice the one that is allowed (three times if we consider the target q-network). Please fix and ensure you remain within the budget.\n\n5. **The paper misses a thorough comparison with the model using the full bag and $k>1$**. Again, it should be based on the same architecture that MAG-GNN uses as downstream prediction model. Only in this way we can know what is the performance achievable with the full bag, and understand how distant MAG-GNN is from it. \n\n\n\nBevilacqua et al 2022. Equivariant Subgraph Aggregation Networks. ICLR 2022\n\nZhang et al 2023. A Complete Expressiveness Hierarchy for Subgraph GNNs via Subgraph Weisfeiler-Lehman Tests. ICLR 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are very grateful for reviewer jrFK’s in-depth review and insightful comments. We address the reviewer’s concerns and questions as follows:\n\n- W1: Full comparison with a random approach\n\nAs mentioned in Q1 in the paper, we also believe that outperforming random subgraphs is critical to the actual utility of the model. We agree with reviewer jkF that a full experiment of MAG-GNN and random approach over $m$ and  $k$ provides a more comprehensive and fair comparison of the two methods. We refer the reviewer to the general response GR1, including the full comparison between MAG-GNN and the random approach.\nThe random approach follows reviewer jrFK’s suggestion. We use sum pooling on the GNN graph embeddings of all subgraphs for downstream MLP.\n\n- W2: Comparison to I-MLE\n\nk-OSAN is undeniably a great inspiration for our work. I-MLE also strives for an efficient and effective subgraph selection policy, and it is an important baseline, as we also discussed from lines 236 to 243 and compared to it in Table 6. The key difference between I-MLE and MAG-GNN is that I-MLE is data-driven, while MAG-GNN changes the sampled subgraphs based on high-level graph features. Consequently, I-MLE fails to sample more representative subgraphs when the desired graph feature is beyond the 1-WL test's capability since the model that parametrizes its sampling distribution is an MPNN. Meanwhile, because MAG-GNN operates on labeled graphs which introduces more expressivity to the model, it can change the selected subgraphs according to graph features undetectable by MPNN/I-MLE. In Table 6, Appendix E of our paper, we already compared MAG-GNN’s performance with k-OSAN IMLE’s *best* ZINC performance reported in its paper. However, its performance is even worse than the random approach. Hence, we implemented IMLE using our codebase (adopting IMLE code module to our $f_p$) and reported the results.\n\n|  | ZINC |\n| --- | --- |\n| k-OSAN I-MLE original m=10, k=1 | 0.155 |\n| k-OSAN I-MLE ours m=10, k=1 | 0.121 |\n| Random m=10, k=1 | 0.127 |\n| MAG-GNN m=10, k=1 | 0.110 |\n\nWe can see that I-MLE only minimally improves over random while MAG-GNN effectively reduces the test MAE.\n\n- W3: Comparison to node-based GNNs.\n\nWe agree that comparing the state-of-the-art node-based methods is important, and we refer the reviewer to the general response GR2 for their comparison to MAG-GNN.\n\n- W4: Fair ZINC comparison.\n\nFollowing the reviewer’s comments, we conducted all ZINC follow-up experiments using a hidden dimension of 64 (previously 100) to respect the 500K parameter constraint.\n\n- W5: Full bag experiments.\n\nWe agree with the reviewer the full bag setting is a meaningful comparison to MAG-GNN. Such a setting implies that a GNN is applied to all k-node-rooted subgraphs, and a pooling function is applied to get the final embedding. Unfortunately, this is generally intractable to train for real-world datasets, even for k>1. In fact, this drawback motivates many works trying to realize k-WL expressivity using less computation resources, including our work. Under our specific implementation, using a full bag is essentially a k-OSAN algorithm, as we mentioned in the preliminary. **Note that the original k-OSAN paper did not include full bag experiments (only k-OSAN with I-MLE reported), most likely due to its intractability.** Note that the k-OSAN result we reported in table 6 is the best performance reported in the original k-OSAN paper with I-MLE. Nevertheless, we still managed to perform $k\\leq 2$ full bag training on SR25 and ZINC datasets.\n\n|  | ZINC | SR25 |\n| --- | --- | --- |\n| k=1 | 0.118 | 6.7 |\n| k=2 | 0.261 | 100 |\n| 3-WLGNN | 0.256 | 6.7 |\n| MAG-GNN | 0.104 | 100 |\n\nWe can see that MPNN even outperforms k=2 full bag training on ZINC dataset. This phenomenon aligns with 3-WLGNN results [1], which also realizes high-dimensional WL with neural networks. While they are theoretically more powerful, properly implementing them on real-world datasets is very challenging, potentially due to the implied larger hidden size required for injective mapping.\n\n- Q3: Ensemble and m>1\n\nWe want to clarify that the result that the reviewer mentioned is observed when we train a MAG-GNN that performs 6 RL steps, and we take the last generated node tuple. The ensemble is on the node tuples over timesteps but not m. However, the same best performance (without pre-training) is obtained with only two RL steps but more node tuples (faster inference). We would also like to highlight the best ZINC performance with pre-training, obtained when m>1, meaning that m>1 is still essential for a more comprehensive graph representation. Here we present a list of $m$ used during hyper-parameter searching, and mark chosen $m$ in **bold**\n\n|  | m |\n| --- | --- |\n| ZINC | {2,**3**,4} |\n| ZINC-FULL | {2,**3**,4} |\n| CYCLES | {2,3,**4**} |\n| MOLHIV | {**2**,3,4} |\n| EXP | {**1**,2,3} |\n| SR25 | {1,2,**3**} |\n| CSL | {1,**2**,3} |\n\n- Q4: Reviewer is jrFK is correct on this notation problem. We will reflect on and correct this in our revision.\n- Q5: We do not reset the buffer every time we update the prediction GNN. As mentioned in Section 4, this causes inconsistency during training, but the models gradually become stable.\n\n[1] Morris, Christopher, et al. \"Weisfeiler and leman go neural: Higher-order graph neural networks.\" *Proceedings of the AAAI conference on artificial intelligence*. Vol. 33. No. 01. 2019."}, {"review_id": "LFoTaJcrdQ", "reviewer": "Reviewer_9Mgd", "summary": "The authors present MAG-GNN, a subgraph-based Graph Neural Network (GNN) method that utilizes reinforcement learning to effectively identify discriminative subgraphs. They empirically assess the performance of MAG-GNN using both synthetic and real-world datasets.", "questions": "1. Please provide clarification on what \"this property\" refers to in line 60.\n\n\n2. Can you confirm if \\mathbf{v} represents a set of k nodes in line 88? It seems so, but the notation is not defined.\n\n3. In line 88, it would be helpful to provide context on why a subgraph GNN creates |V^k(G)| copies of the graph and an intuition of what this quantity represents.\n\n4. In Equation (2), could you specify what j represents?\n\n5. Could you explain what N refers to in Equation (3)?\n\n6. In line 136, when you mention a \"2-node-marked graph,\" could you provide a general definition of a node-marked graph?\n\n7. What is meant by a \"GNN run\" or \"MPNN run\"?\n8. The explanation in the text of Figure 2 is quite complex. Additional information from the authors would be appreciated to improve comprehension.\n\n9. Is \\mathcal{V}^k the same as V^k(G) mentioned in line 172? If so, it would be helpful to use standardized notation to help in understanding.\n\n10. In section 4.1 on training MAG-GNN, it is challenging to grasp all the training procedure details solely from the text. It would be beneficial to include an algorithm that illustrates the training procedure.\n\n11. The captions of Table 1 and Table 2 do not indicate which metrics are being presented.\n12. What do the yellow cells indicate in Table 2?\n13. All the tables and figures in the experimental section lack error bars.\n14. How does the runtime of MAG-GNN vary with m and T in section 6.4 on runtime comparison?\n", "limitations": "The authors have acknowledged the limitations of their proposed approach.", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "- The main elements of the RL framework are properly explained.\n- The empirical section includes clear and valuable research questions.", "weaknesses": "1. The writing is significantly lacking in clarity and context.\n2. The presentation of results lacks clarity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful to reviewer 9Mgd’s effort in helping us improve our work and paper. We address the reviewer’s concerns and questions as follows.\n\n- Q1: “The property” refers to many graphs’ properties that a discriminative small subset, rather than the whole set, of a graph’s subgraphs is sufficient in identifying the graph. More detailed explanation in the previous paragraph.\n\n- Q2: Here, $\\boldsymbol{v}$ represents a node tuple of k nodes. It’s different from a node set, as it allows repetition and is ordered. For example, when k=2, a node tuple $(\\boldsymbol{v}_1, \\boldsymbol{v}_1)$ is allowed, and node tuples $(\\boldsymbol{v}_1, \\boldsymbol{v}_2)$ and $(\\boldsymbol{v}_2, \\boldsymbol{v}_1)$ are different. We follow the convention of using the power of  $V$ to represent the set of node tuples.\n\n- Q3: The power of **node-rooted subgraph** GNNs over regular GNNs is that they do not directly run MPNN on the graph. Instead, they create $|V|$ copies of the graph and apply MPNN onto each copy, where $|V|$ is the number of nodes. Such an approach requires each copy to be associated with a node in the graph, making each copy unique. A widely-adopted way is to use the subgraph marking technique introduced in lines 90-92. The unique marking to each copy allows the corresponding nodes to learn local structures that plain GNNs cannot. Ultimately, the information of all copied graphs is pooled to form the final representation. Since GNN with $|V|$ graphs can be more expressive than regular GNNs, later work proposes to use **2-node-rooted subgraph** GNNs that create $|V^2|$ graph copies. This approach results in an even stronger GNN than node-rooted subgraph GNN. Naturally, we can keep extending subgraph GNNs to use more subgraphs, which generalize to **k-node-rooted subgraph** GNNs. Theoretical work finds out that subgraph GNNs expressivity increases as k increases. *Because there are $|V^k|$ unique node tuples, we can thus have $|V^k|$ graph copies with unique markings.*\n\n- Q4: We will clarify using an example:\n\nLet G be a graph consisting of nodes $v_1$, $v_2$, $v_3$, $v_4$, and $k=2$, meaning we assign unique labels to 2-node-rooted graphs. Let the node tuple be $(v_2, v_4)$ , and let $c^+=1$ and $c^-=0$. The resulting unique node labeling matrix $X$ is\n\n| COL 1 | COL 2 |\n| --- | --- |\n| 0 | 0 |\n| 1 | 0 |\n| 0 | 0 |\n| 0 | 1 |\n\nBecause $v_2$ is the first node ($j=1$) in the node tuple, we have $X_{2,1}=1$, and $v_4$ is the second node ($j=2$) in the node tuple, we have $X_{4,2}=1$. $j$ represents the index in the node tuple.\n\n- Q5: N is not for the natural number set. It represents (**N**)ode-level pooling function to distinguish it from (**G**)raph-level pooling function, as described in line 98.\n\n- Q6: Random 2-node-marked graphs represent the graph copies labeled by random 2-node tuples under the labeling described in equation (2). A 2-node-marked graph creates a copy of the original graph, draws a random 2-node tuple from $V^2$ , and labels the graph according to equation (2). In the paper, we use the uniquely labeled/marked graph as a general form of subgraph because previous work points out that a labeled graph is no less expressive than the subgraph [1].\n\n- Q7: One GNN/MPNN run means using GNN or MPNN to process a graph/subgraph once to get its node/graph embedding. If we run GNN/MPNN on a set of $|V|$ graphs, we say we perform $|V|$ GNN runs on the set of graphs.\n\n- Q8: Figure 2 shows the prediction score of individual 2-node-marked graphs in the SR25 dataset. Recall that a graph $G$ in the SR25 dataset has 25 nodes; hence, there are $25^2=625$ unique 2-node-marked graphs. Following equation (4), we can compute the prediction score of each 2-node-marked graph to the class of $G$, resulting in 625 prediction scores. We sort these scores and plot them in Figure 2. Because SR25 has 15 graphs, Figure 2 has 15 lines. The figure shows that among all possible 2-node-marked graphs of graph $G$, many individually identify $G$ from the rest of the graphs in the dataset.  Such an experiment shows that one of the most challenging datasets has the desirable property that a small set of subgraphs is discriminative enough, which motivates our RL framework to find such subgraphs.\n\n- Q9: They are not the same. $V$ refers to the set of nodes in a specific graph $G$. However, graphs have different node sets. Hence, we use $\\mathcal{G}$ to represent the set of all graphs and $\\mathcal{V}$ to represent the set of all possible nodes appearing in $\\mathcal{G}$. Consequently, $\\mathcal{V}^k$ represents all possible k-node tuples of all graphs.\n\n- Q10: We thank the reviewer’s suggestion and include the pseudocode in the pdf attached to the general response.\n\n- Q11: We thank the reviewer for pointing out the unclarity. While we already included the evaluation metric description in the experimental results section, we will include this when more space is permitted in the revision.\n\n- Q12: Yellow cells indicate MAE smaller than 0.01, as mentioned in line 350. We will include this information in the caption in the revision.\n\n- Q13: We omit the error bar following previous works as the variance is minimal [2]. The results for molecular data (Table 6) included in the appendix reports standard deviations. We will move it to the main paper to account for model's randomness.\n\n- Q14: We present the inference time of MAG-GNN on the QM dataset following the setup in section 6.4.\n\n| Time(ms) | m=2 | m=3 | m=4 |\n| --- | --- | --- | --- |\n| t=2 | 704.9 | 847.9 | 959.2 |\n| t=3 | 925.5 | 1275.1 | 1445.5 |\n| t=4 | 1489.0 | 1743.6 | 1875.4 |\n\nWe can see that under all these circumstances, MAG-GNN is still faster than I^2GNN and PPGN.\n\n[1] Zhang, Bohang, et al. \"Rethinking the expressive power of gnns via graph biconnectivity.\" arXiv preprint arXiv:2301.09505 (2023).\n\n[2] Huang, Yinan, et al. \"Boosting the Cycle Counting Power of Graph Neural Networks with I $^ 2$-GNNs.\" *arXiv preprint arXiv:2210.13978* (2022)."}]

Author rebuttal:
We thank all reviewers for their constructive comments that help us improve our work. Here we address some common questions and concerns that reviewers raise, and we also respond to their comments individually.

- GR1: MAG-GNN is a random method that uses an RL agent to refine the selected subgraph, a full comparison with a pure random approach is required to show the necessity and utility of such agent.

This is also the question that we try to answer in Q1 in the experiment section, and following reviewer SjU6 and jrFK’s suggestion, we also believe a more comprehensive study showing that MAG-GNN outperforms random subgraphs is critical to the actual utility of the model. We compare MAG-GNN and the random subgraph approach over $m$ and  $k$ on real-world datasets.

The random approach follows reviewer jrFK’s suggestion. We use sum pooling on the GNN graph embeddings of all subgraphs for a fair comparison. The results to the left of the bar in a cell are the results of the random approach, and the results to the left are MAG-GNN results.

| SR25 |  | m=1 |  | m=2 |  | m=3 |
| --- | --- | --- | --- | --- | --- | --- |
| k=1 |  | 6.7  \|  6.7 |  | 6.7   \|   6.7 |  | 6.7  \|  6.7 |
| k=2 |  | 65.4    \|    100 |  | 82.3    \|    100 |  | 89.6    \|   100 |
| k=3 |  | 67.4    \|     100 |  | 83.9    \|    100 |  | 93.8    \|    100 |

Compared to random approach that does not achieve that perfect score even when m=3 and k=3, MAG-GNN achieves perfect score with k=2 and m=1, meaning that MAG-GNN effectively finds more representative node-tuples.

We conducted a similar comparison on the ZINC dataset.

| ZINC |  | m=2 |  | m=3 |  | m=4 |
| --- | --- | --- | --- | --- | --- | --- |
| k=3 |  | 0.131  \|  0.115 |  | 0.124   \|   0.111 |  | 0.125  \|  0.110 |
| k=4 |  | 0.135    \|    0.106 |  | 0.131   \|    0.104 |  | 0.134    \|   0.101 |
| k=5 |  | 0.132    \|     0.112 |  | 0.136    \|    0.114 |  | 0.131    \|    0.109 |
| k=6 |  | 0.144    \|     0.124 |  | 0.140    \|     0.127 |  | 0.136   \|    0.116 |

The best validation score for MAG-GNN is obtained at m=3, k=4, leading to a test score of 0.104. We note that MAG-GNN outperforms the random approach under all configurations, showing that MAG-GNN also performs better on real-world datasets.

A similar pattern can be found in the OGBG-MOLHIV dataset:

| MOLHIV | | m=2 | | m=3 | | m=4 |
| ---|--- | --- | --- | --- |---|---|
| k=3| | 75.21  \|  77.12 || 76.06   \|   77.26 || 76.29  \|  76.52 |
| k=4 || 74.69   \|   76.53 || 75.92   \|    78.10 || 76.51    \|   76.28 |
| k=5 || 75.17    \|     74.95 || 74.83    \|    76.81 || 75.37    \|    76.30 |
| k=6 || 74.28    \|     74.64 || 75.20    \|     76.14 || 75.68   \|   76.19 |

MAG-GNN also outperforms the random approach in most cases.

Because the QM9 dataset has many labels, due to the space limitation, we pick $U_0$ where usually higher expressivity GNN achieves a better score:

| QM9-$U_0$ |  | m=2 |  | m=3 |  | m=4 |
| --- | --- | --- | --- | --- | --- | --- |
| k=3 |  | 0.306 \|  0.241 |  | 0.318   \|   0.205 |  | 0.285  \|  0.153 |
| k=4 |  | 0.327    \|    0.144 |  | 0.309   \|    0.111 |  | 0.291    \|   0.119 |
| k=5 |  | 0.297    \|     0.137 |  | 0.288    \|    0.125 |  | 0.294    \|    0.121 |
| k=6 |  | 0.336    \|     0.168 |  | 0.314    \|     0.154 |  | 0.296   \|    0.146 |

QM9-$U_0$ is more sensitive to the hyperparameters. We can still see that the random approach brings minimal performance as $m$ increases, while MAG-GNN obtains good performance with appropriate m and k.

- GR2: While MAG-GNN approximates the expressivity of intractable higher-order WL, how does MAG-GNN compare to tractable note-based methods?

We agree that comparing the state-of-the-art node-based methods is important, and we present MAG-GNN's results against them on ZINC, ZINC-FULL and SR25 datasets.

|  | ZINC | ZINC-FULL | SR25 |
| --- | --- | --- | --- |
| SUN[1] | 0.083 | 0.024 | 6.7 |
| SSWL+ [2] | 0.070 | 0.022 | 6.7 |
| MAG-GNN | 0.104 | 0.029 | 100 |
| MAG-GNN PRE | 0.098 | 0.023 | - |

State-of-the-art node subgraph methods on the ZINC dataset outperform MAG-GNN. However, it is also important to note that they do not acquire >3-WL expressivity (random guesses on the SR25 dataset) as MAG-GNN does, meaning that they will potentially fail on datasets requiring higher expressivity. The higher expressivity of MAG-GNN requires more data. We also compared pre-trained MAG-GNN with these models. We find that after pre-training, MAG-GNN is very close to these models on ZINC-FULL when more data is available. We also would like to emphasize that MAG-GNN is more optimal in terms of inference time and show an inference time comparison between MAG-GNN and these models.

| time (ms) | ZINC-FULL | CYCLE |
| --- | --- | --- |
| SUN | 2590.2 | 1410.5 |
| SSWL+ | 2682.6 | 1473.8 |
| MAG-GNN m=3 k=2 | 321.6 | 155.1 |

Note that MAG-GNN, even when k=2, feeds much fewer subgraphs to the MPNN compared to SSWL+ and SUN, which results in better inference time. MAG-GNN is designed to reduce the $O(n^k)$ subgraphs for $k$-tuple-based subgraph GNNs to a constant while still maintaining the maximum potential expressive power the subgraph GNNs can achieve.

- GR3: More details on different training paradigms.

Per reviewer 9Mgd's request, we added pseudocode of our different training paradigms in the attached pdf.

[1] Frasca, Fabrizio, et al. "Understanding and extending subgraph gnns by rethinking their symmetries." Advances in Neural Information Processing Systems 35 (2022): 31376-31390.

[2] Zhang, Bohang, et al. "Rethinking the expressive power of gnns via graph biconnectivity." arXiv preprint arXiv:2301.09505 (2023).

Meta-review: This paper proposes a reinforcement learning boosted approach to subgraph graph neural networks (GNNs). This addresses the exponential complexity of exact subgraph GNNs in comparison to regular MPNNs (since they require to enumerate all subgraphs of a graph). The authors evaluated the need of this exhaustive enumeration, finding that often a select subset of subgraphs offers the same discriminative ability as all subgraphs. Inspired by this observation, an RL-based approach (MAG-GNN) is introduced by the authors to determine the most effective subset of subgraphs for the specific task being addressed.

After reading the paper, all reviews, and rebuttal, I believe there is reasonable consensus that the original submission should have considered a few other scenarios (e.g., random policy, better comparison with existing work), as aptly pointed out by reviewer jrFK (which provided a high-quality review and was highly engaged during the rebuttal). The rebuttal, however, seems to have answered most of these concerns. I am happy to recommend acceptance.
Decision: Accept (poster)