Efficient Subgraph GNNs by Learning Effective Selection Policies
Authors: Beatrice Bevilacqua
Moshe Eliasof
Eli Meirom
Bruno Ribeiro
Haggai Maron
Venue: ICLR 2024 poster (2024)
Keywords: Graph Neural Networks
Subgraphs
Expressive power
Sampling
Primary area: learning on graphs and other geometries & topologies
TL;DR: We propose a novel framework that learns to select subgraphs sequentially in order to reduce the computational cost of Subgraph GNNs.

Abstract: Subgraph GNNs are provably expressive neural architectures that learn graph representations from sets of subgraphs. Unfortunately, their applicability is hampered by the computational complexity associated with performing message passing on many subgraphs. In this paper, we consider the problem of learning to select a small subset of the large set of possible subgraphs in a data-driven fashion. We first motivate the problem by proving that there are families of WL-indistinguishable graphs for which there exist efficient subgraph selection policies: small subsets of subgraphs that can already identify all the graphs within the family. We then propose a new approach, called _Policy-Learn_, that learns how to select subgraphs in an iterative manner. We prove that, unlike popular random policies and prior work addressing the same problem, our architecture is able to learn the efficient policies mentioned above. Our experimental results demonstrate that _Policy-Learn_ outperforms existing baselines across a wide range of datasets.

Published as a conference paper at ICLR 2024
EFFICIENT SUBGRAPH GNNS BY LEARNING EFFEC- TIVE SELECTION POLICIES
Beatrice Bevilacqua
Moshe Eliasof University of Cambridge me532 camm.cmm.. ac. uk
Eli Meirom NVIDIA Research eme www@nviiia..ommmm com
Purdue University
bbevi 1 accpurrdee edu
Bruno Ribeiro
Haggai Maron
Purdue University ribeiro@cs, purdue. edu edu
Technion & NVIDIA Research
hmaron@nvidia com
ABSTRACT
Subgraph GNNs are provably expressive neural architectures that learn graph rep- resentations from sets of subgraphs. Unfortunately, their applicability hampered by the computational complexity associated with performing message passing on many subgraphs. In this paper, we consider the problem of learning to select a small subset of the large set of possible subgraphs in data-driven fashion. We first motivate the problem by proving that there are families of WL-indistinguishable graphs for which there exist efficient subgraph selection policies: small subsets of subgraphs that can already identify all the graphs within the family. We then pro- pose new approach, called POLICY-I.EARN, that learns how to select subgraphs in an iterative manner. We prove that, unlike popular random policies and prior work addressing the same problem, our architecture is able to learn the efficient poli- cies mentioned above. Our experimental results demonstrate that POLICY-LEARN outperforms existing baselines across a wide range of datasets.
INTRODUCTION
Subgraph GNNs (Zhang& & Li, 2021; Cotta et al., 2021; Papp et al., 2021; Bevilacqua et al., 2022;
Zhao et al., 2022; Papp Wattenhofer, 2022; Frasca et al., 2022; Qian al., 2022: Zhang et al., 2023) have recently emerged as promising research direction to overcome the peressive-power limitations of Message Passing Neural Networks (MPNNs) (Morris et al., 2019; Xu et al., 2019). In essence, Subgraph GNN first transforms input graph into bag of subgraphs, obtained according to predefined generation policy. For instance, each subgraph might be generated by deleting exactly one node in the original graph, or, while leaving the connectivity unaltered (Papp & Wattenhofer, 2022). Then, it applies an equivariant architecture process the bag of subgraphs, and aggregates the representations to obtain graph- or node-level predictions. The popularity of Subgraph GNNs can be attributed not only to their increased expressive power compared MPNNs but also their remarkable empirical performance, exemplified by their success on the ZINC molecular dataset (Frasca et al., 2022; Zhang et al., 2023).
generally, by marking exactly one node in the original graph,
Unfortunately, Subgraph GNNs are hampered by their computational cost, since they perform message-passineg operations on all subgraphs within the bag. In existing subgraph generation policies, there are at least as many subgraphs in the bag as there are nodes the graph (n). leading computational complexity of O(n² d) contrasted to just O(n d) of a standard MPNN, where d represents the maximal node degree. This computational overhead becomes intractable large graphs, preventing the applicability of Subgraph GNNs commonly-used datascts. To address this limitation, previous works (Cotta et al., 2021; I; Bevilacqua et al., 2022; Zhao et al., 2022) have considered randomly sampling small subset of subgraphs from the bag in every training epoch, showing that this approach not only significantly reduces the running time and computational cost, but also empirically retains good performance.
This raises the following question: is it really necessary to consider all the subgraphs in the bag? To answer this question, consider the two non-isomorphic graphs in Figure 1, which are indistinguishable

---

Published as a conference paper at ICLR 2024
WL-indistinguussaall non-isomorphic graphs
Bags Wi.dissinguiisaalee subgraph
Original graph
WL histogram
Subgraph
WL histogram
l..
Original graph
WL histogram
Subgraph
WL histogram
Figure I: Sufficiency of small bags, Two non-isomorphic graphs (Circulant Skip Links graphs with 13 nodes and skip length and respectively) that be distinguished using bag containing only a single subgraph, generated by marking one node.
by MPNNs. As demonstrated by prior work (Cotta et al., 2021), encoding such graphs as bags of subgraphs, each obtained by marking single node, makes them distinguishable. However, can be seen from the Figure, the usage of all subgraphs is unnecessary, and distinguishability can be achieved by considering only one subgraph per graph. Thus, from the standpoint of expressive power, Subgraph GNNs can distinguish between these graphs using constant-sized bags, rather than the n-sized bags typically employed by existing methods.
The example in Figure 1 shows the potential of subgraph selection policies returning small number of subgraphs. However, as all markedd subgraphs these two graphs are isomorphic, randomly selecting one subgraph in each graph is sufficient for disambiguation. In Section we generalize this example and prove that there exist families of non-isomorphic graphs where small number of carefully chosen subgraphs is sufficient and necessary for identifying otherwise WL-indistinguishabl graphs. these families, the probability that a random selection baseline returns the required subgraphs tends to zero, motivating the problem of learning to select subgraphs instead of randomly choosing them.
In this paper, aim to learn subgraph selection policies, returning significantly smaller number of subgraphs compared to the total available in the bag, task- and graph-adaptive fashion. To that end, propose an architecture, dubbed POLICY-LEARN, composed of two Subgraph GNNs: (i) a selection network, and (ii) a downstream prediction network. The selection network learns a subgraph selection policy and iteratively constructs bag of subgraphs. This bag is then passed to the prediction network to solve the downstream task of interest. The complete architecture is trained in end-to-end fashion according to the downstream task's loss function, by leveraging the straight-through -ofbeel-Sfftmx estimator (Jang et al., 2016; Maddison et al., 2017) to allow differentiable sampling from the selection policy.
As we shall see, our choice of using Subgraph GNN as the selection network, as well as the sequential nature of our policy generation, provide additional expressive power compared random selection and previous approaches (Qian et al., 2022) and allow us to learn more diverse set of policies. Specifically, we prove that our framework can learn to select the subgraphs required for disambiguation in the families of graphs mentioned above, while previous approaches provably cannot do better than randomly selecting subset of subgraphs.
Experimentally, we show that POLICY-LEARN is competitive with full-bag Subgraph GNNs, even when learning to select a dramatically smaller number of subgraphs. Furthermore, we demonstrate its advantages over random selection baselines and previous methods wide range of tasks.
Contributions. This paper provides the following contributions: (1) A new framework for learning subgraph selection policies; (2) A theoretical analysis motivating the policy learning task well as our architecture; and (3) An experimental evaluation of the new approach demonstrating its advantages in terms of predictive performance and reduced runtime.
2

---

Published as a conference paper at ICLR 2024
2 RELATED WORK
Expressive power limitations of GNNs. Multiple works have been proposed to overcome the expressive power limitations of MPNNs, which are constrained by the WL isomorphism test (Xu et al., 2019; Morris et al., 2019). Other than subgraph-based methods, which are the focus of this work, these approaches include: (1) Architectures aligned the k-WL hierarchy (Morris et al., 2019; 2020b; Maron et al., 2019b;a); (2) Augmentations of the node features with node identifiers (Abboud et al., 2020; Sato, 2020; Dasoulas et al., 2021; Murphy et al., 2019): (3) Models leveraging additional information, such as homomorphism and subgraph counts (Barceló et al., 2021; Bouritsas et al., 2022), graph polynomials (Puny et al., 2023), and simplicial and cellular complexes (Bodnar et al., 2021b:a). We refer the reader to the survey by Morris et al. (2021) for review of different methods.
Subgraph GNNs Despite the differences in the specific layer operations, the idea of constructing bags of subgraphs to be processed through weessage-passing layers has been recently proposed by several concurrent works. Cotta et al. (2021) proposed to first obtain subgraphs by deleting nodes in the original graph, and then apply an MPNN to each subgraph separately, whose representations are finally aggregated through a set function. Zhang & Li (2021) considered subgraphs obtained by constructing the local ego-network around each node in the graph, an approach also followed by Zhao et al. (2022), which further incorporates feature aggregation modules that aggregate node representations across different subgraphs. Bevilacqua al. (2022) developed two classes of Subgraph GNNs, which differ in the symmetry group they are equivariant to: (i) DS-GNN, which processes each subgraph independently, and (ii) DSS- DSS-GNN, which includes feature aggregation modules of the same node in different subgraphs. Huang et al. (2023); Yan et al. (2023); Papp & Wattenhofer (2022) further studied the expressive power of Subgraph GNNs, while Zhou et al. (2023) proposed framework that encompasses several existing methods. Recently, Frasca et al. (2022); Qian et al. (2022) provided a theoretical analysis of node-based Subgraph GNNs, demonstrating that they are upper-bounded by 3-WL, while Zhang et al. (2023) presented complete hicrarchy of existing architectures with growing expressive power. Finally, other works such as Rong et al. (2019); Vignac et al. (2020); Papp et al. (2021); You al. (2021) can be interpreted as Subgraph GNNs.
Most related to our work is the recent paper of Qian et al. (2022), which, to the best of our knowledge, presented the first framework that learns a subgraph selection policy rather than employing a full bag or a random selection policy. Our work differs from theirs in four main ways: (i) Bag of subgraphs construction: Our method generates the bag sequentially, compared their one-shot generation; (ii) Expressive selection GNN architecture: We parameterize the selection network using a Subgraph GNN, that is more expressive than the MPNN used by Qian et al. (2022); (iii) Differentiable sampling mechanism: We use the well-known Gumbel-Softmax trick to enable gradient backpropagation through the discrete sampling process, while Qian et al. (2022) use I-MLE (Niepert et al., 2021); ); and (iv) Motivation: our work is primarily motivated by scenarios where a subset of subgraphs is sufficient for maximal expressive power, while Qian et al. (2022) aim to make higher-order generation policies (where each subgraph is generated from tuples of nodes) practical. We prove that our choices lead to a framework that can learn subgraph distributions that cannot be expressed with Qian et al. (2022), and show that our framework performs better than Qian et al. (2022) on real-world datasets.
3 PROBLEM FORMULATION
Let G (A, be a node-attributed. undirected, simple graph with nodes, where A R represents the adjacency matrix of G and X E g= is the node feature matrix. Given a graph G, Subgraph GNN first transforms G' into bag (multiset) of subgraphs, Bg SS S2..., ... S, B. obtained from a predefined generation policy, where S, = (A,, X,) denotes subgraph i. Then, it processes Be through stacking of (node- and subgraph-) permutation equivariant layers, followed by a pooling function to obtain graph or node representations to be used for the final predictions.
In this paper, we focus the node marking (NM) subgraph generation policy (Papp & Wattenhofer, 2022), where each subgraph is obtained by marking a single node in the graph, without altering the connectivity of the original graph' More specifically, for each subgraph S: = (A,, X2). there exists a node eú such that A; = A and X, X @ X where @ denotes channel-wise concatenation and
Althoughh technically, NM does not generate subgraphs in the traditional sense, the marked versions of the original graph are commonly regarded as subgraphs (Papp & Wattenhofer, 2022).
3

---

Published as a conference paper at ICLR 2024
Xe; is one-hot indicator vector for the node u,. We refer to the node v as the root of S,, which will equivalently denote as St.- Notably, since each subgraph is obtained by marking exactly one node, there is a bijection between nodes and subgraphs, making the total number of subgraphs under the NM policy equal to the number of nodes m n. Finally, we remark that we chose NM as the generation policy as it generalizes several other policies (Zhang et al., 2023). Any other bose- policy, preserving this bijection between nodes and subgraphs, can be equivalently considered.
Objective. Motivated by their recent success, would like to employ Subgraph GNN solve given graph learning task (e.g., graph classification). Aiming to mitigate the computational overhead of using the full bag Ba. wish to propose method for learning a subgraph selection policy (GG) that, given the input graph G. returns subset of Bg to be used as an input to the Subgraph GNN in order solve the downstream task at hand.
We denote by B% C Bg the output of π (G), consisting of the original graph and I selected subgraphs. That is, B = [G, Su2y S. B. where node in G is the root node of Sey-
4 INSIGHTS FOR THE SUBGRAPH SELECTION LEARNING PROBLEM
In this section, we motivate the problem of learning select subgraphs in the context of Subgraph GNNs. Specifically, we seek to address two research questions: (i) Are there subgraph selection policies that relurn small yel effective bags of subgraphs? and (ii) How do these policies compare to strategies that uniformly sample a random subset of all possible subgraphs in the bag?. In the following, we present our main results, while details and proofs are provided in Appendix C.
Powerful policies containing only single subgraph. To address the first question, we consider Circulant Skip Links (CSL) graphs, an instantiation of which was shown in Figure 1 (CSL(13, 5) and CSL(13, 3)).
Definition 1 (CSL graph (Murphy et al., 2019)). Let n, k, be such that n is co-prime with & and k < - 1. A CSL graph, CSL(n, k) is undirected graph with nodes labeled as {0, , }}. whose edges form a cycle and have skip links. That is, the edge set E is defined by a cycle formed by (j.5 + 1)
E for j E (O.....7, 2}, and (n 1, 0) E, along with skip links defined recursively by the sequence ($1, S1+1) E E. with $1 = and $;+1 = (s; + k) mod n.
While CSL graphs are WL-indistinguishable (Murphy et al., 2019). they become distinguishable when considering the full bag of node-marked subgraphs (Cotta et al., 2021). Furthermore, since all subgraphs in the bag are isomorphic, it is easy to see that distinguishability is obtained even when considering a single subgraph per graph. This observation motivates the usage of small bags, as there exist cases where the expressive power of the full bag can be obtained with very small subset of subgraphs.
The extreme case of CSL graphs, however, is too simplistic: randomly selecting one subgraph in each graph is sufficient for disambiguation. This leads to our second question, and, in the following, we build upon CSL graphs define a family of graphs where random policies not effective. As shall see next, in these familics only specific small subgraph selections lead to complete identification of the isomorphism type of each graph.
Families of graphs that require specific selections of & subgraphs. We obtain our family of graphs starting from CSL graphs, which we use as building blocks to define an (n, E)-CSL graph.
Definition 2 (n, C)L graph). A (n, E)- CSL graph, denoted as (n, C)-CSL(R, (ki,...,ke)) is a graph with é-n nodes ob- tained from disconnected con-isomorphic CSL (sub-))se huub-ggraphs with n nodes, CSL(n, ke), i (1.....f). Note that the maximal value of depends
Our interest in these graphs lies in the fact that the family of non-isomoiphic (n, f)- CSL graphs contains Wi-iniistiguuishable graphs, as we show in the following.
Figure 2: An (n, C)LC graph is obtained from dis- connected, con-isomoep CSL graphs with nodes, where 13 in this case.
4

---

Published as a conference paper at ICLR 2024
Sampling
Selection Network
Current Bag
Distribution over the nodes
Sampled subgraph
If (8 7+1
Node/Graph representation
Prediction Network
Figure 3: An overview of POLICY-LEARN. POLICY-LEARN consists of two Subgraph GNNs; selection network and prediction network. The selection network generates the bag of subgraphs by iteratively parameterizing probability distribution over the nodes of the original graph. When the hag size reaches its maximal size T ++1 (the original graph plus 7 selections), it is passed to the prediction network for the downstream task.
Theorem 1 ((n, ()-CSL graphs are WL-indistinguishabl)) Let Gn,e e be the family of non-isomorphic (n, graphs (Definition 2). All graphs in Gn.E WL-indistinguishable
While these graphs are WL-indistinguishable we can distinguish all pairs when utilizing the full bag of marked subgraphs. More importantly, the full bag allows to fully identify the isomorphism type of all graphs within the family, which implies the distinguishability of pairs of graphs². As we show next, it is however unnecessary consider the entire bag of all the f subgraphs. Instead, there exists a subgraph selection policy that returns only & subgraphs, significantly smaller number than the total in the full bag, which are provably sufficient for identifying the isomorphism type of all graphs within the family. This identification is attained when the root nodes of all l subgraphs belong different CSL guup)s. Furthermore, we prove that this condition is not only sufficient but also necessary. In other words, the isomorphism-tyee type identification of all graphs within the family only occurs when there is at least one mode-marked subgraph obtained from each CSL. (sub-)er )graph.
Proposition 2 (There exists an efficient π that fully identifies (n, f)-CSL graphs). Let Gn.E be the family of non-ioomore -oniisooo (n, E)-CSL graphs (Definition 2). node marking bused subgraph selection policy n can identify the isomorphism type of any graph G in Gn, if and only if its bag has marking each CSL (sub-) surbph.raph.
The above results demonstrate the existence of efficient and more sophisticated subgraph selection policies that enable the identification of isomorphism types. Remarkably, these can return bags as small containing subgraphs, as long each of them is obtained by marking node in a different CSL Jgrah. The insight derived from the necessary condition is equally profound: a random selection approach, which uniformly samples é subgraphs per graph, does not ensure identification of the isomorphism type, because the subgraphs might not be drawn from distinct CSL formalize below.
This finding highlights the sub-optimality of random selection strategies,
Proposition 3 (A random policy cannot efficiently identify (n, L)-CLL graphs). Let be any graph in the family of non-ist isomorphic (n, L-CCLL graphs, namely Gu,E (Definition 2). The probability that random policy, which uniformly samples E subgraphs from full bag, can successfully identify the isomorphism type of G is fl/ef, which tends to 0 as & increases.
The above result has important implications for Subgraph GNNs when coupled with random policies. It indicates that Subgraph GNNs combined with random subgraph selection policy, are unlikely to
Note that we consider the family of non-isomorphic (n, E)-CSL graphs, thus excluding cases where the choices of ki -.. be lead isomrrphic connected components.
5

---

Published as a conference paper at ICLR 2024
Algorithm 1 POLICY-LEARN: Feedforward with learnable subgraph selection policy
Initialize bag with original graph: B == [G]
2: for = ....., T
1 do
3: Feed current bag B% to selection network f, obtain a distribution over the nodes of G,
Pt-1 f
4: Sample node Ut+1 ~ Pt+1
5: Add subgraph bag B B%
USS
6: end for
7: Feed final bag By to the downstream network g and output result g
identify the isomorphism type of G, assuming the number of sampled subgraphs is l. In Lemma 6 in Appendix C, we further show that the expected number of subgraphs that a random policy must uniformly draw before identification is T = O In f). Notably, this number is significantly larger than the minimum number of subgraphs required by æ. For instance, when f = 100, a random policy would have to sample around 460 subgraphs per graph. This is more than four times the number of subgraphs returned by т, making it clear that there exists a substantially more effective approach than random sampling to the problem.
METHOD
The following section describes the proposed framework: POLICY-LEARN. We start with overview of the method, before discussing in detail the specific components, namely the selection and the prediction networks, as well the learning procedure used differentiate through the discrete sampling process.
5.1 OVERVIEW
Our method illustrated in Figure and described in Algorithm 1. It is composed of two main Subgraph GNNs: subgraph selection network, denoted as f. that generates the subgraph selection policy n, and a downstream prediction network, denoted as g. that solves the learning task of interest. More concretely, given input graph G, we initialize the bag of subgraphs with the original graph, B2 = [G Then, at each step t + 1, the selection network takes the current bag of subgraphs B input and outputs nodes of the original graph G. Our algorithm samples from Pt+1 to select a node 0++1 that will serve the root of the next subgraph S This subgraph is added to the bag, i.e. we define B B% { B. After 7 subgraph selections, where T is hyperparameter representing the number of subgraphs in the bag (T' + 1 considering the original graph), the final bag π (G) == B is passed to the prediction network g. Subgraph GNN designed to solve the downstream task at hand. trick (Jang al., 2016; Maddison et al., 2017) to backpropagate through the discrete sampling process. In what follows, we elaborate on the exact implementation and details of each component of our approach.
un-normalized distribution P++1 = f fB E R", defined over the
Both f and g are trained jointly in an end-to-end manner, and we use the Gumbel-Softma
5.2 SUBGRAPH SELECTION NETWORK
The goal of f output a distribution over the nodes of the original graph, determining which node should serve as the root of the next subgraph that will be added to the bag. As this network takes bag of subgraphs employ the most efficient Subgraph GNN variant, namely DS-GNN (Cotta et al., 2021; Bevilacqua al., 2022), that performs message-passing operations on each subgraph independently using an MPNN, referred to as the subgraph encoder. Subsequently, it applics pooling operation across the subgraphs the bag, to obtain single representation for each node in the input graph G.
its input, it is natural implement
Subgraph GNN. Specifically, we
Since we have both the original node features and the markings, the subgraph encoder in our DS-GNN architecture is implemented using the model proposed by Dwivedi et al. (2021), which decouples the two representations to improve learning. This encoder consists of two cesssage-passing functions: fn. responsible for updating the markings, and fa. responsible for updating the original features. Formally, at the (-th iteration, we provided with bag of subgraphs B of size 1. For each

---

Published as a conference paper at ICLR 2024
subgraph O,....t, ., we compute the following for 1 1, 1,.... L1 where L1 represents the total number of layers in the selection network:
p!0 1 ( AA): p RRxx
(la)
h = f/ (n x A.) h.o E ER
(1b)
with fo and f( the 1-th message passing layers in fp and fni respectively. At 0, we initialize the marking as (0) Xv, 2 and the input node features as h22 X E paxeii..
After the L1-th layer, we combine the marking and node features to obtain the updated node features, as follows:
hg = -!-p)- + h(4"), hg E R"
(2)
Note that in this way we obtain node features for every subgraph in the bag. In order to unify them into a single output per node, we pool the node features along the subgraph dimension. We then apply MLP that reduces the channel dimension one, to obtain the following an-noraalized node probabilities from which we sample the next root node:
Po+1
hig = MLP(ppool he Pu+1 E R..
(3)
Backpropagating through the sampling process. A naive sampling strategy not differentiable and therefore will prevent us from backpropagating gradients to update the weights of the selection network f. Therefore, we use the trriiht-throoulh
Humbel-Softman estimator (Jang et al., 2016; Maddison et al., 2017) to sample from Pr+1 and obtain the one-hot encoding marking of the next subgraph to be added to the current bag of subgraphs:
Xe GumbelSoftmax (p++1), X E E{0,1}"
(4)
We denote the subgraph associated with the node marking X0e11 by Smgr Therefore, the bag of subgraphs at the (t + 1)-th iteration reads B = B& U & B.
5.3 DOWNSTREAM PREDICTION NETWORK
Similarly to the selection network f. our prediction network g is parameterized as a Subgraph GNN. It takes input the final bag of subgraphs By that was sampled using f. and produces graph- or node-level predictions depending on the task of interest. Like the selection network, we employ a DS-GNN architecture and implement the subgraph encoder using two message- paagge--psinngggg functions, Ip and gh. with layers / = l,..., Le and updates identical to Equations (la) and (lb). Note that, while and y are architecturally similar, they are two different networks with their own set of learnable weights. In the first layer, i.e., / = O, we set the initial marking as ps Xo. E R" and input node features h X E p=xc.,, where
the bar, i.c., (!) h to denote the representations in the downstream prediction model and distinguish them from the ones in the selection model.
After the L2-th th layer, combine the marking and node features to obtain the updated node features Th = (24 T(( and pool the node features along the subgraph and potentially the node dimension to obtain final node or graph representation. This representation is then further processed with an MLP to obtain the final prediction.
THEORETICAL ANALYSIS
In this section, we study the theoretical capabilities of our framework, focusing on its ability to provably identify the isomorphism type of the (n, E)-CSL graphs introduced in Section 4. Due to space constraints, all proofs are relegated Appendix C. We start by showing that, differently than random selection approaches, POLICY-LEARN can identify the isomorphism type of any (n, E)- CSL graph (Definition 2) when learning to select only subgraphs. Remarkably, this number significantly smaller than the total & subgraphs in the full bag, as well as the expected number of subgraphs @ In f) required by a random policy.
Theorem 4 (POLICY-LEARN can identify Gn.t)- Let G be any graph in the family of non-isomorphic (n, e)-CSL graphs, namely Gm.F (Definition 2). There exist weights of f and such that POLICY- LEARN with T can identify the isomorphism type of G within the family.
7

---

Published as a conference paper at ICLR 2024
The idea of the proof lies in the ability of the selection network within POLICY-LEARN to implement a policy π that returns subgraphs per graph where: (i) each subgraph is obtained by marking exactly one node and (ii) each marked node belongs to a different CSL (sub-)graph. Given that this policy serves as a sufficient condition for the identification of the isomorphism types (Proposition 2). our proposed framework successfully inherits this ability. More specifically, the proof (in Appendix C) proceeds by showing that the selection network can always distinguish all the nodes belonging to already-marked CSL (sub-)graph from the rest, and therefore can assign them a zero probability of being sampled, while maintaining uniform probability on all the remaining nodes.
The previous result further sets our method apart from the OSAN method proposed in Qian al. (2022). Indeed, as its selection network is parameterized as an MPNN, and given that Gn.e contains Wh-innistinuuishable graphs composed by WL-iinisttiguushabbl CSL. subb))graphy then OSAN cannot differentiate between all the nodes in the graph. Therefore, since all the subgraphs are selected simultaneously rather than sequentially, the resulting sampling process is the same as uniform sampling. Consequently, it cannot ensure that marked nodes belong to different CSL (sub-)graphs. Since this condition is not only sufficient but also necessary for the identification of the isomorphism type (Proposition 2), OSAN effectively behaves like a random policy for indistinguishable graphs.
Theorem 5 (Qian et al. (2022) cannot efficiently identify Gn,)). Let G be any graph in the family of non-isomorphic (n, E)-CSL graphs, namely In.e (Definition 2). Then, the probability that OSAN (Qian et al., 2022) with = subgraphs can successfully identify the isomorphism type of G is C,/e, which tends to 0 as l increases.
Up until this point, our focus has been on demonstrating that POLICY-LEARN has the capability to implement the policy ı for identification within Gn.e. However, it is important to note that is not the sole policy that POLICY-LEARN can implement. Indeed, POLICY-LEARN can also learn subgraph selection policies that depend on the local graph structure, such as the policy returning all subgraphs with root nodes having degree greater than specified threshold (see Proposition 7 in Appendix C).
EXPERIMENTS
In this section, we empirically investigate the performance of POLICY-LEARN. In particular, we seek to answer the following questions: (1) Can our approach consistently outperform MPNNs, as achieve better performance
demonstrated by full-bag Subgraph GNNs? (2) Does POLICY-L.AAN than random baseline that uniformly samples a subset of subgraphs at each training epoch? (3) How does our method compare to the existing learnable selection strategy OSAN (Qian et al., 2022)? (4) Is our approach competitive to full-bag Subgraph GNNs on real-world datasets?
In the following, we present our main results and refer to Appendix E for details. 3 For each task, we include two baselines: RANDOM, which represents the random subgraph selection, and FULL, which corresponds to the ball-bag approach. These basclines utilize the same downstream prediction network (appropriately tuned) POI.CY-LEAAR, with RANDOM selecting random subsets of subgraphs equal in size to those used by POLICY-LEARN (denoted by 7"), and FULL using the entire bag.
Table I:
Test results on the ZINC molecular dataset under the 500k parameter budget.
We experimented with the ZINC-12K molecular dataset (Sterling & Irwin, 2015;
ZINC.
Method
ZINC (MAF 4)
Gömez-Bombarelli et al., 2018; Dwivedi et al., 2020), where as prescribed we maintain 500k parameter budget. As can be seen from Ta- ble 1, POLICY-LEARY significantly improves over OSAN, and surpasses the random baseline. Notably, OSAN performs worse than our ran- dom baseline due to differences in the implemen- tation of the prediction network. As expected, the gap between POLICY-LEARN and RANDOM is larger when the number of subgraphs T is smaller, where selecting the most informative subgraphs is more crucial.
GCN (Kipf & Welling, 2017) GIN (Xu et al., 2019)
0.321
0.163
PNA (Corso et al., 2020)
0.133++0.011
GSN (Bouritsas et al., 2022)
0.101=00010
CIN (Bodnar el al., 202
0.079-10.006
OSAN (Qian et al., 2022) T' 2
0.177-10.016
FULL RANDOM T 2 POLICY-LEARN T
0.087=00003
0.036
0.120-10.003
RANDOM T =
0.113-10.006
POLICY-LEARN T' =
0.109+10.005
*Our code is available at wwtpthttp//// http:///ttttttcco///////e///////
8

---

Published as a conference paper at ICLR 2024
Table Test results on the OGB datasets. POLICY-LEARN outperforms both the learnable selection approach OSAN and the random selection by large margin, while approaching surpassing the corresponding full-bag architecture. N/A indicates the result was not reported.
Method 4 Dataset +
MOLESOL. RMSE L
MOLTOX21 ROC-AUC t
MOLBACE ROC-AUC T
MOLHIV ROC-AUC 1
GCN (Kipf & Welling, 2017)
1.114+00.036
75.29+:0.69 74.91+0.51
79.15+1.44
76.06+
GIN (Xu et al., 2019)
1.173+0.057
72.97+4.00
7..58+1.40
OSAN (Qian et al., 2022) =
0.984
N/A
72.30-16.00
N/A
FULL
0.847=000155
76.25+1.12
78.41+1.94
1..5
RANDOM T =
0.951+0.039
76.65+0..99
75.36=4.22
1..5
POLICY-LEARN T
0.877=0.029
77.47+0.82
78.40=2.85
79.13+0.60
RANDOM T 5
0.900-10.032
76.62-10.63
78.14
77.30
POLICY-LEARN T =
0.883110.032
77.36-10.0
78.29
78.49:11011
Reddit. To showcase the capabilities of POLICY-
Table 2: Results on the REDDIT-BINARY dataset demonstrate the efficacy large graphs
LEARN large graphs, experimented with the REDDIT-BINARY dataset (Morris et al., 2020a) where full-bag Subgraph GNNs cannot be otherwise applied due to the high average number of nodes, and consequently subgraphs, graph (429.63). POLICY-LEARY exhibits the best performance among all efficient baselines (Table 2), while maintaining competitive infer- ence runtimes (Table in Appendix E), opening up the applicability of Subgraph GNNs unex- plored avenues. Table 6 in Appendix further compares to random baselines coupled with dif- ferent subgraph-
Method
RDT-B (ACC
GIN (Xu et al., 2019)
92.4:12.55
SIN (Bodnar et al., 2021b)
92.2+1.0
CIN (Bodnar al., 2021a)
92.4+2.1
FULL
00M
RANDOM ' 20
92
RANDOM POLICY-LEARN
92.4#10
93.010.9
downstream models.
OGB. We tested our framework several datasets from the OGB benchmark collection (Hu et al., 2020). Table shows the performances of our method when compared to MPNNs and efficient subgraph selection policies. For OSAN, we report the best result across all node-based policies, which is obtained with a larger number of subgraphs than what we consider (T = 10 compared to T E {2, 5)). Notably, POLICY-LEARN improves over MPNNS while retaining similar asymptotic complexity. and consistently outperforms OSAN and RANDOM. The performance gap between POLICY-LEARN and RANDOM more significant with MOLBACE. Remarkably, POLICY-LEARN approaches the full-bag baseline FULLI on MOLESOL. and MOLBACE, and even outperforms it on MOLTOX21 and MOLHIV, achieving the highest score. These results can be attributed the generalization challenges posed by the dataset splits, which are further confirmed by the observation that increasing the number of subgraphs does not improve results. Table 4 in Appendix E additionally reports the results of existing full-bag Subgraph GNNs.
2. particularly on MOLESOL and
CONCLUSIONS
In this paper we introduced a novel framework, POLICY-LEARN, for learning to select a small subset of the large of possible subgraphs in order to reduce the computational overhead of Subgraph GNNs. We addressed the pivotal question of whether a small, carefully selected, subset of subgraphs can be used identify the isomorphism type of graphs that would otherwise WI.- indistinguishable. We proved that, unlike random selections and previous work, our method can provably learn to select these critical subgraphs. Empirically, we demonstrated that POLICY-LEARN consistently outperforms random selection strategies and approaches the performance of full-bag methods. Furthermore, it significantly surpasses prior work addressing the same problem, across all the datasets considered. This underscores the effectiveness of our framework in focusing on the necessary subgraphs, mitigating computational complexity, and achieving competitive results. We believe approach opens the door to broader applications of Subgraph GNNs.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
BR acknowledges support from the National Science Foundation (NSF) awards CCF-1918483, CA- REER IIS-1943364 and CNS-2212160, Amazon Research Award, AnalytiXIN, the Wabash Heartland Innovation Network (WHIN). Ford, NVidia, CISCO, and Amazon. Computing infrastructure was supported in part by CNS- 192-01 (CloudBank). HM is the Robert J. Shillman Fellow, and is supported by the Israel Science Foundation through a personal grant (ISF 264/23) and an equipment grant (ISF 532/23). Any opinions, findings and conclusions or recommendations are those of the authors and do not necessarily reflect the views of the sponsors.
REFERENCES
Ralph Abboud, ismail Îlkan Ceylan, Martin Grohe, and Thomas Lukasiewicz. The surprising power of graph neural networks with random node initialization. In Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence (IJCAI), 2020.
Pablo Barceló, Floris Geerts, Juan Reutter, and Maksimilian Ryschkov. Graph neural networks with local graph parameters. Advances in Neural Information Processing Systems, 34:25280-25295, 2021.
Beatrice Bevilacqua, Fabrizio Frasca, Derek Lim, Balasubramaniam Srinivasan, Chen Cai, Gopinath Balamurugan, Michael M Bronstein, and Haggai Maron. Equivariant subgraph aggregation networks. In International Conference on Learning Representations, 2022.
Lukas Biewald. Experiment tracking with weights and biases, 2020. Software available from wandb.com.
Gunnar Blom, Lars Holst, and Dennis Sandell. Problems and Snapshots from the World of Probability. Springer Science Business Media, 1993.
Cristian Bodnar, Fabrizio Frasca, Nina Otter, Yuguang Wang, Pictro Liō, Guido F Montúfar, and Michael Bronstein. Weisfeiler and lehman go cellular: Cw networks. In Advances in Neural Information Processing Systems, volume 34, 2021a.
Cristian Bodnar, Fabrizio Frasca, Yuguang Wang, Nina Otter, Guido F Montúfar, Pictro Liø, and Michael Bronstein. Weisfeiler and lehman go topological: Message passing simplicial networks. In International Conference on Machine Learning, 2021b.
Giorgos Bouritsas, Fabrizio Frasca, Stefanos P Zafeiriou, and Michael Bronstein. Improving graph neural network expressivity via subgraph isomorphism counting. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022.
Guangyong Chen, Pengfei Chen, Chang Yu Hsich, Chee-Kong Lee, Benben Liao, Renjie Liao, Weiwen Liu, Jiezhong Qiu, Qiming Sun, Jie Tang, Richard Zemel, and Shengyu Zhang. Alchemy: A quantum chemistry dataset for benchmarking ai models. arXiv preprint arXiv: 1019900099777 2019,
Zhengdao Chen, Lei Chen, Soledad Villar, and Joan Bruna. Can graph neural networks count substructures? Advances in neural information processing systems, 33:1 13:10383-10395, 2020.
Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Liò, and Petar Veličkovié. Principal neighbourhood aggregation for graph nets. In Advances in Neural Information Processing Systems, volume 33, 2020.
Leonardo Cotta, Christopher Morris, and Bruno Ribeiro. Reconstruction for powerful graph repre- sentations. In Advances in Neural Information Processing Systems, volume 34, 2021.
George Dasoulas, Ludovic Dos Santos, Kevin Scaman, and Aladin Virmaux. Coloring graph neural networks for node disambiguation. In Proceedings of the Twenty-Ninth International Conference on International Joint Conferences on Artificial Intelligence, pp. 2126-2132, 2021.
Vijay Prakash Dwivedi, Chaitanya K Joshi, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks. arXiv preprint arX9v:2003.00982, 2020.
10

---

Published as a conference paper at ICLR 2024
Vijay Prakash Dwivedi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Graph neural networks with learnable structural and positional representations. In International Conference Learning Representations, 2021.
Matthias Fey and Jan Eric Lenssen. Fast graph representation learning with pytorch geometric. arXiv preprint arXiv: 024428, 2019.
Fabrizio Frasca, Beatrice Bevilacqua, Michael Bronstein, and Haggai Maron. Understanding and extending subgraph gnns by rethinking their symmetries. In Advances in Neural Information Processing Systems, 2022.
Rafael Jomez-Bombarrlli, Jennifer N. Wei. David Duvenaud, José Miguel Hernández-Lobatd Benjamín Sánchez-I Sánchez-Lengelling Dennis Sheberla, Jorge eguilera-Iparraguire Timothy D. Hirzel, Ryan P. Adams, and Alán Aspuru-Guzs Guzik, Automatic chemical design using data-driven continu- ous representation of molecules. ACS Central Science, 4(2):268-211 Jan 2018.
Shengding Hu, Zheng Xiong, Meng Qu, Xingdi Yuan, Marc-Alexandre Côté, Zhiyuan Liu, and Jian Tang. Graph policy network for transferable active learning on graphs. Advances in Neural Information Processing Systems, 33: 13:104-4--0888 2020.
Yinan Huang, Xingang Peng, Jianzhu Ma, and Muhan Zhang. Boosting the eycle counting power of graph neural networks with GNNNs. International Conference on Learning Representations, 2023.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparamcterization with gumbel-softmax. In International Conference on Learning Representations, 2016.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations, 2017.
Chris Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: continuous
relaxation of discrete random variables. In Proceedings of the international conference learning Representations, 2017.
Haggai Maron, Heli Ben-Hamu, Hadar Serviansky, and Yaron Lipman. Provably powerful graph networks. In Advances in Neural Information Processing Systems, volume 32, 2019a.
Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. In International Conference on Learning Representations, 2019b.
Karolis Martinkus, Pál András Papp, Benedikt Schesch, and Roger Wattenhofer. Agent-based graph neural networks. In International Conference Learning Representations (ICLR). arXiv, 2023.
Eli Meirom, Haggai Maron, Shie Mannor, and Gal Chechik. Controlling graph dynamics with reinforcement learning and graph neural networks. In International Conference on Machine Learning, pp. 7565-7577. PMLR, 2021.
Michael Mitzenmacher and Eli Upfal. Probability and computing: Randomization and probabilistic techniques algorithms and data analysis. Cambridge university press, 2017.
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. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pp. 4602-4600 2602-4609,00 2019.
Christopher Morris, Nils M. Kriege, Franka Bause, Kristian Kersting, Petra Mutzel, and Marion Neumann. Tudataset: A collection of benchmark datasets for learning with graphs. In ICML 2020 Workshop on Graph Representation Learning and Beyond (GRL+ 2020), 2020a.
Christopher Morris, Gauray Rattan, and Petra Mutzel. Weisfeiler and leman go sparse: Towards scalable higher-order graph embeddings. In Advances in Neural Information Processing Systems, volume 33, 2020b.
Christopher Morris, Yaron Lipman, Haggai Maron, Bastian Rieck. Nils M Kriege, Martin Grohe, Matthias Fey, and Karsten Borgwardt. Weisfeiler and leman go machine learning: The story so far. arXiv preprint 1 2021.
11

---

Published as a conference paper at ICLR 2024
Ryan Murphy, Balasubramaniam Srinivasan, Vinayak Rao, and Bruno Ribeiro. Relational pooling for graph representations. In International Conference on Machine Learning. pp. 466346634669 4673. PMLR, 2019.
Mathias Niepert, Pasquale Minervini, and Luca Franceschi. Implicit mle: backpropagating through discrete exponential family distributions. Advances in Neural Information Processing Systems, 34: 14567-14579, 2021.
Pál András Papp and Roger Wattenhofer. A theoretical comparison of graph neural network extensions. In International Conference on Machine Learning, pp. 17323-17345. PMLR, 2022.
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.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performant deep learning library. In Advances in Neural Information Processing Systems, volume 32, 2019.
Omri Puny, Derek Lim, Bobak Kiani, Haggai Maron, and Yaron Lipman. Equivariant polynomials for graph neural networks. In International Conference on Machine Learning, pp. 28191-28222 PMLR, 2023.
Chendi Qian, Gaurav Rattan, Floris Geerts, Christopher Morris, and Mathias Niepert. Ordered subgraph aggregation networks. In Advances in Neural Information Processing Systems, volume 35, 2022.
Chendi Qian, Andrei Manolache, Kareem Ahmed, Zhe Zeng, Guy Van den Broeck, Mathias Niepert, and Christopher Morris. Probabilistically rewired messaag-passing neural networks, 2023.
Yu Rong, Wenbing Huang, Tingyang Xu, and Junzhou Huang. Dropedge: Towards deep graph convo- lutional networks on node classification. In International Conference on Learning Representations, 2019.
Ryoma Sato. A survey on the expressive power of graph neural networks. arXiv preprint arXiv:2003.04078, 2020.
Teague Sterling and John J. Irwin. ZINC 15 - ligand discovery for everyone. Journal of Chemical Information and Modeling, 55(11):232 11 2015. doi: 10.1021/acs.jiim.5bbl 5b00559..
Clement Vignac, Andreas Loukas, and Pascal Frossard. Building powerful and equivariant graph neural networks with structural message-passing Advances in neural information processing systems, 1014144--14555 2020.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? International Conference on Learning Representations, 2019.
Zuoyu Yan, Junru Zhou, Liangcai Gao, Zhi Tang, and Muhan Zhang. Efficiently counting substruc- tures by subgraph gnns without running gnn subgraphs. arXiv preprint arXiv:2303. 10576. 2023.
Gilad Yehudai, Ethan Fetaya, Eli Meirom, Gal Chechik, and Haggai Maron. From local structures to size generalization in graph neural networks. In International Conference on Machine Learning, pp. 11975-11986, PMLR, 2021.
Jiaxuan You, Jonathan M Gomes-Selman, Rex Ying, and Jure Leskovec. Identity-aware graph neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 10737-10745, 2021.
Chulhee Yun, Suvrit Sra. and Ali Jadbabaie. Small relu networks are powerful memorizers: a tight analysis of memorization capacity. Advances in Neural Information Processing Systems, 32, 2019.
12

---

Published as a conference paper at ICLR 2024
Bohang Zhang, Guhao Feng, Yiheng Du, Di He, and Liwei Wang. A complete expressiveness hierarchy for subgraph gnns via subgraph weisfeiler-Ichme tests. In International Conference on Machine Learning, 2023.
Muhan Zhang and Pan Li. Nested graph neural networks. In Advances in Neural Information Processing Systems, volume 34, 2021.
Lingxiao Zhao, Wei Jin, Leman Akoglu, and Neil Shah. From stars to subgraphs: Uplifting any GNN with local structure awareness. International Conference on Learning Representations, 2022.
Cai Zhou, Xiyuan Wang, and Muhan Zhang. From relational pooling to subgraph gnns: A universal framework for more expressive graph neural networks. In International Conference Machine Learning, 2023.
13

---

Published as a conference paper at ICLR 2024
A ADDITIONAL RELATED WORK
Our approach is related not only to Subgraph GNNs and expressive GNN architecture, as discussed in the main paper, but also other sampling strategies the graph domain. In the context of active learning, Hu et al. (2020) propose reinforcement learning approach that learns which subset of nodes to label in order to reduce the annotation cost of training GNNs. For controlling the dynamics of a system using localized interventions, Meirom et al. (2021) employ reinforcement learning agent that selects a subset of nodes at each step and attempts to change their state, with an objective that depends on the number of nodes in each state. More similar to sampling strategy is the work by Martinkus et al. (2023) which proposes a GNN architecture augmented with neural agents, that traverse the graph and then collectively classify it. Notably, each step, an agent chooses neighboring node to transition to using the Gumbel-Softmax trick.
B ILLUSTRATION OF WL ON THE BAG OF SUBGRAPHS
In this section explain in detail Figure 1. For each original graph, consider the application of the WI. test to the subgraph obtained by marking the crossed node. More precisely, at time step O, each non-marked node in the subgraph is assigned the same initial constant color, which differs from the initial color of the marked node. Without loss of generality, denote these colors as color 1 and color O, respectively. Then, the WL test proceeds by refining the color of each node by aggregating the colors of its neighbors (including itself). More precisely, the new color of a node is obtained through a hash function taking as input the multiset of colors of the neighbors and of the node itself. It is easy to see that, at time step I, the marked node has unique color, which we call color 2, its neighbors all have the same colors, color 3, and all the remaining nodes have the same color (different from the marked node and its neighbors), color 4. The refinement is repeated until we reach stable coloring. At this point, we simply collect the number of nodes having the same color in a histogram. For the upper graph in Figure 1, the histogram indicates that in the stable coloring, there are 4 blue nodes, 4 green nodes, 4 yellow nodes, and orange node. Importantly, since the WL. test represents a not
necessary condition for graph isomorphism, different histograms imply that the two graphs isomorphic, and thus the WL test distinguishes them. Figure shows that marking only one node sufficient for distinguishing the two graphs. which instead indistinguishable when no node is marked, as the WL test returns the same histogram for them (13 orange nodes).
C PROOFS
This appendix includes the proofs for the theoretical results presented in Sections 4 and 6. We start by formally proving that the family Gz.e of non-isomorphic (n, f)-CSL graphs, defined in Definition 2, contains WL-indistinguishable graphs.
Theorem 1 (n, C)-CSL graphs WL-indistinguisaable. Let Gn.e e be the family of non-isomorphic (n, e)-CSL graphs (Definition 2). All graphs in Gn.s are WL-indistinguishabb
Proof. Let G1, G2 be any two (n, E)-CSL graphs in Gn.e- We will prove that G1, G2 are WL- indistinguishable by simulating one WL round starting from constant color initialization.
(Iter. 1) The colors after the first iteration represent the node degrees. Since each node has degree 4 (because it belongs to one CSL ssub-)graph), then all nodes will have the same color.
Since the colors after the first iterations are a finer refinement of the initialization colors, the algorithm stops and the graphs are not separated.
Next, we show that there exists an efficient policy π, that can return as few subgraphs as f, providing strong guarantees. Indeed. represents necessary and sufficient condition for identifying the isomorphism type of all graphs in Gn.e, as we prove next.
Proposition 2 (There exists an efficient т that fully identifies (n, E)-CSL graphs). Let Gn.t be the family of non-isomorphic (n. CCLC graphs (Definition 2). A
based subgraph selection
14

---

Published as a conference paper at ICLR 2024
policy π can identify the tsomorphism rype of any graph G in Gn,E if and only if its bag has marking in each CSL (sub-)graph.
Proof. We will prove the two cases separately.
Sufficiency. Consider any graph G E Gn,e, and let S, be subgraph returned by obtained by marking any node in the i-th CSL (sub-)graph of graph G. Note that St is sufficient to identify the isomorphism type of the i-th CSL )gab-, because marking any node is sufficient for disambiguation of any possible pair of CSL graphs (Cotta et al., 2021, Theorem 2). Consider the set H of l histograms, each obtained by marking different CSL (sub-) )graph. Note that if returns more than subgraphs, then identical histograms indicate that the corresponding subgraphs are obtained from the same CSL (sub-)graph, because G is composed by non-isomorphic CSL (sub-) )graphs. Thus, identical histograms can be simply discarded when constructing H. Then, the set H is sufficient to identify the isomorphism type of G. as each element of H identifies one its CSL (sub-)graphh.
Necessity. The proof proceeds by contrapositive. Consider a policy T' and assume that there exists least one CSL (sub-) )graph such that none of the marked nodes belongs to it. Without loss of generality, assume that only one CSL (sub-)graphe is not covered, meaning no marked nodes belong to it. Then, n' cannot identify the isomorphism type of G. because is unaware of the isomorphism type of the CSL (sub-)graph that is not covered. More precisely, consider graph G' E Gn. such that: (1) G' is non-isomorphic to G and (2) G' differ from G only in one CSL (sub-)graph, which the one not covered by n.. Then, n' cannot identify whether the isomorphism type of its input is the one of G or the one of G¹.
The above theorem implies that a random selection baselines might fail to select the subgraphs required for the identification, it might fail to draw at least one subgraph from each CSL (sub- Ograph. In the following we prove that the implications are even more severe as l grows.
Proposition 3 (A random policy cannot efficiently identify (n, f)-CSL graphs). Let G be any graph in the family of non-ison non-isomorphic (n, e)-CSL graphs. namely Go.E (Definition 2). The probability that a random policy, which uniformly samples & subgraphs from the full bag, can successfully identify the isomorphism type of G is 0l/ef, which tends to 0 as & iincresses.
Proof. From Proposition 2, each of the l subgraphs must be generated by marking a node belonging to different CSL (sub-)graph for the identification of the isomorphism type (necessary condition). Recall that the total number of CSL (sub-)graphs is f, and each CSL (sub-)graph has exactly n nodes. Thus, the probability of choosing é nodes from different CSL (sub-)graphe Cl/e Finally, since +om El/e" = 0, the probability goes to zero as l grows.
Importantly, we derive the expected number of subgraphs that a random policy has to draw before identification.
Lemma 6. Let G be any graph in the family of non-isomorphic (n, f)-CSL graphs, namely Gn.E (Definition 2). The expected number of subgraphs that a random policy must uniformly draw before identification is T = (n
Proof. Note that we have an instance of the coupon collector problem. The proof follows the steps in Mitzenmacher & Upfal (2017); Blom et al. (1993). Let t, be the time to collect (any node from) the i-th CSL (sub-) .graphs after i 1 CSL (sub-)graphr have already been collected by drawing at least one of their nodes for marking. Let pt be probability of collecting a new CSL. (sub-)graph (one from which nodes where not drawn before). Then,
- (i
f-i+1
Pt
E.
t
15

---

Published as a conference paper at ICLR 2024
Therefore, t; has geometric distribution with expectation E(t,] 1
i
By linearity of
P:
8-1+1
expectation, then the number of draws needed to collect all CSL (sub-)graphs is:
Eit + t2 + +... + tel =
+ E[t2]
+ E[te]
1
1
+
pt
P2
Pt
L
i
+ 8-1
+
-( (((
2)
1
+
2
=HH,
where Hg is the & harmonic number, which asymptotically grows as O(n E). Thus, we obtain the expectation O In E).
Our next theorem shows that POLICY-LEARN does not suffer from the shortcomings of the random selection. Indeed, POLICY-LEARN provably implement the π in Proposition 2. which in turn implies that it can learn to identify the isomorphism type of any graph in Gn.e-
Theorem 4 (POLICY-LEARN can identify Gn,t). Let G be any graph in the family of non-isomorphic (n, e)-CSL graphs, namely Gn.t (Definition 2). There exist weights of f and such that POLICY- LEARN with T f can identify the isomorphism type of G within the family.
Proof. We will show that POLICY-LEARN can implement the subgraph selection policy which returns l subgraphs such that: (i) each subgraph is obtained by marking exactly one node and (ii) each marked node belongs to different CSL (sub-)graph. Since is sufficient for identification (Proposition 2), then POLICY-LEARN can provably identify the isomorphism type of G.
Recall that we initialize the bag of subgraphs B8 [G]. At step t O, since all nodes are WL-indistinguishable (see proof of Theorem 1), then P1 = f(B)) is a uniform distribution. possible nodes, and the bag is updated as B Bg U (S,,,, B = {G, B. At step L = 1, we obtain a distribution over the nodes as P2 = (B)). Assuming ʃ has enough layers (it is sufficient* that the total number of layers L. is n), then all nodes belonging to the CSL. (sub-)graph of node V1 will have different color than the remaining (e 6--1 1)- n nodes. Therefore the MLP in Equation (3) can map them to have zero probability. On the contrary, the remaining (L 1)- nodes will all have the same color, and the MLP will map them to have the same probability. Thus p2 will be uniform over all the nodes in the CSL (sub-)grre -grapgg other than the one containing node v1. Therefore, U2 ~ P2 is randomly sampled from CSL (sub-)graph different than one containing v1- By repeating the argument, we have that at each iteration POLICY-LEARN samples one node from a CSL (sub-)graph that has not yet been selected. Thus, at ( = l the bag of subgraphs B% will contain all the subgraphs that are sufficient for identification.
Thus, U1 ~ Pl is randomly sampled from the L
On the contrary, since OSAN (Qian et al., 2022) uses an MPNN selection network and samples all subgraphs at once, then it cannot perform better than randomly selecting subgraphs for the graphs in Gr.e-
Theorem 5 (Qian et al. (2022) cannot efficiently identify Gn,)). Let G be any graph in the family of non-isomorphic (n, E)-CSL graphs, namely Gn.e (Definition 2). Then, the probability that OSAN (Qian et al., 2022) with T = subgraphs can successfully identify the isomorphism type of G is 0./e" which tends to 0 as e increases.
Proof. Recall that the selection network in OSAN is parameterized as an MPNN taking as input the original graph and returning as output & probability distributions over the l nodes, one for each of
the l subgraphs to sample. Since all nodes in G are WL--inistinguishable (see proof of Theorem 1). then each of the l probability distribution is uniform over the L. n nodes. Therefore, each node is sampled uniformly at random. Since the sampled nodes must belong to different CSL Jgrahh
But not necessary, as for example the bound can be tightened to n by considering the circular structure of the connected components, and even further by taking into account the skip connections.
16

---

Published as a conference paper at ICLR 2024
for identification (Proposition 2), and because the probability of sampling nodes each belonging to different CSL bub-))raph from the uniform distribution is é1/e" (Proposition 3), then the probability that OSAN can identify the isomorphism type of G is e1/E".
Finally, we prove that POLICY-LEARN can return all and only subgraphs having as root a node with degree greater than predefined number.
Proposition 7 (POL.CY-LFFARN returns all m node-marked subgraphs whose root nodes have degree greater than a predefined d, with d 1. Then, POLICY-LEARN can implement with T m.
implement degree-aware policies). Consider a policy that
Proof. We will prove that we can implement degree-aware policies by showing that exists a set of weights such that POLICY-LEARN gives uniform probability to all nodes with degree d that have not been already selected. Since the selection network is iteratively applied for T = m times, then POLICY-LEARN would have selected all subgraphs after the last iteration.
Recall that we initialize the bag subgraphs as B% := {GG. and, at any step, we sample a node and add the corresponding subgraph the bag. In the following, drop the dependency on both the graph and the step in the notation, and call the current bag simply by B.
only interested in sampling subgraphs having root nodes with degree d we do not assume the presence of any additional initial node feature, and we directly work with a DS- DN-GNN architecture with GraphCony encoders (Morris al., 2019).5 The representation for node in subgraph given by the (1 + 1) layer of our selection network can be written as:
Since
(++1) b = o( hos + W
(
Σ h +6 W
(5)
where k1 = (M...) with id.,; = O if v is not the root of subgraph i, and id..; = 1 if v is the root of subgraph i. Recall that, given the representations for each node in each subgraph in the bag, we then obtain a distribution over the nodes in the original graph by simply pooling the final node representations ( at the final layer L, across all the subgraphs, i.e., the unnormalized node distribution is obtained as:
hpp = pool, (119) (2)
We consider a selection Subgraph GNN with 4 layers (i.e., T. = 4), and ReI.U non-linearities. Each layer will have two channels, one used to compute an indicator variable storing whether node has degree greater than d, and the other one propagating whether a node is the root of the subgraph. In the following describe each layer in detail.
First layer. We set w(() (8 W (18), 6(0)
=(id). then for node we have,
(1 =
h (1).0)
(9 1 ido,3)
(0
0
Σ wm
idu,, ) +
1
d
1 (1),1
0
(di(v)= d) ido,i
where di(v) is the degree of v in subgraph i, and since we are using a node marked policy, it is equal the degree of G, namely d(v). Note that in this way the zero-th entry of (1) namely (1).0 is greater than if and only if d(v) > d, and 0 otherwise.
'The extension to the presence of additional node features, as well as the usage of the separate propagation of them as in Dwivedi et al. (2021) is straightforward, and can be done by considering those additional weights zeroed out.
17

---

Published as a conference paper at ICLR 2024
Second layer. We w(D) ( 1). W22 (88). 60) = (b), then for node we have,
(1)
[2) =
(2).0)
-1
1
0 + :
m + ())
h (2).1
0
1
1
(1)00 +1
(1)
id
Note that in this way the zero-th entry of h (2) namely 1, is 0 d(v) > d, and is equal to 1 if d(v) S d.
Third layer. We set W(2) (o 9). w/(
-088) b(2) =()). then for node we have,
1 =
(3).0)
-1
0
(: 0 Σ 2 + ))
h N2
(3).1
= a
0
+
u
he +1
(2),0
n.a id.,s
Note that in this way the zero-th entry of h h.g namely (31.0 is 0 d(u) S d, and is equal to d(v) d. This means that at this point, (3),,0 contains the binary information indicating whether node has degree greater than d. Therefore, we can use this entry construct our distributions over nodes, ensuring we sample only nodes that have degree greater than d. However, we want to avoid re-sampling nodes that have already been sampled. We will show how this can be done by first using one additional layer that makes use of idv,i, Vo E G, and Vi E B, and then relying on our pooling function.
Fourth layer. We set w(()) = (1 -1). W = =(00, 6(3) = 0. then for node v we have,
h A/d = a((1 -1) (1 + (0
0) Σ
h + 0)
(3)
a (1)00)
h (3),1)
(h id...)
(3).0
Note that in this way (4) is if and only if d(v) > d and id..; = O. Therefore, if we consider all the subgraphs in the current bag, if is the root of one subgraph in the bag, then there exist i E B such
that (1) = 0. In other words, Vi E B, we have 1,44) 1, then means that d(v) > d and e has
(4) not been already selected. On the contrary, if 3i E B, such that a, e.. O, then it means that either
d(v) < or e has been already selected as root. Thus, we can rely on our pooling operation to create the node distribution.
Pooling layer. We set pool, min;, and therefore
hy min (1 (1)) KEB
Note that if h2 O. then either u has degree less than or equal to d or it has been already selected as a root node (i.c., its corresponding subgraph is in B). If instead h, = 1. then v has degree greater than d and has not been selected. Thus, hy represents the unnormalised probability distribution over the nodes.
18

---

Published as a conference paper at ICLR 2024
D ADDITIONAL EXPRESSIVE POWER ANALYSIS
In this section we extend the expressive power analysis reported in the main paper. We start by thoroughly comparing the expressive power of POLICY-LEARN and OSAN (Qian et al., 2022). While we have seen in Theorems and 5 that there exist families of graphs that POLICY-LEARN can efficiently identify while OSAN cannot, we show here that the contrary is not true: there are no families of graphs that OSAN can identify and POLICY-LEARN cannot. This follows from the fact that on graphs of a fixed size n, for any instantiation of OSAN there exists a set of weights for POLICY-LEARN such that outputs exactly the same probability distribution.
Proposition 8 (POLICY-LEARN can implement OSAN). Let n be a natural number and consider all graphs of size n. Then there exists set of weights for POLICY-LEARN such that it can parameterize all the probability distributions that OSAN can, but the contrary is not true.
Proof. We start by showing that on all graphs of size n POLICY-LEARN can parameterize all probability distributions that OSAN can.
Recall that the selection network in OSAN is parameterized an MPNN taking input the original graph and returning as output T probability distributions over the n nodes, one for each of the T subgraphs to sample (simultaneously). This means that the output of the selection network in OSAN can be seen as a matrix of size n x T. Consider now the selection network f in POLICY-LEARN, We will show that there exists a set of weights such that at every iteration with & (l....,T]J outputs the same probability distribution that OSAN outputs for the 1-th subgraph.
Let / be the current iteration, and recall that outputs node features for every subgraph in the bag (Equation (2)). which are then pooled and passed through an MLP (Equation (3)). Note that if f discards the marking information, then it implies that f outputs the same node features for every subgraph in the bag, and can be parameterized exactly as the MPNN used in OSAN. Consider therefore the case where f discards any marking information, and outputs for each subgraph a matrix of size n x (T' + 1), where the first T columns are exactly equal to the ones in the matrix returned by OSAN, and the last column is simply all ones. Therefore, if we pool the node representations across subgraphs using sum pooling, then the output is matrix of size n x (T + 1), where the first T columns are a multiple of the ones in OSAN, and the last column contains the number of subgraphs in the bag, or, equivalently, the iteration t. Finally, we rely the ability of MLPs to memorize finite number of input-output pairs (see Yun et al. (2019) and Yehudai et al. (2021, Lemma B.2)). We use that result to show that there exists a point-wise MLP that transforms, for each row, the value t in the last column into its one-hot representation, thus obtaining matrix n X (T + T), where once again the first T columns are identical to the ones in OSAN and the last T columns are used to maintain for each row a one-hot representation of l. Finally, we only need a final point-wise MLP that computes, for each row, the inner product between the first T entries and the last T entries. Since we have a finite number of graphs, we only need to make sure that the MLP memorizes a finite number of such inner products, and the lemma mentioned above in Yehudai et al. (2021) guarantees there is such an MLP. This MIP effectively implements a selection of the 1-th column of the matrix using the one-hot representation of encoded in the last columns, and the t-th column corresponds to the distribution returned by OSAN for the t subgraph. This means that for every iteration t POLICY-LEARN can output the probability distribution returned by OSAN for the f-th subgraph.
Furthermore, from Theorems 4 and 5 it follows that OSAN cannot parameterize all probability distributions that POLICY-LEARN can.
While the above result compared the expressive power of OSAN and POLICY-LEARN, it is also important to understand how well POL.IC--LFARN compares to the full-bag approach on different families of graphs. In Section 4 we have already discussed that for CSL graphs it is sufficient to mark any single node to obtain the same expressive power of the full-bag. However, this result can be extended to other non-isomorphic WL-indistinguushable families of graphs. Specifically, we show next that marking any node is sufficient to distinguish strongly regular graphs of different parameters (n. k. X, μ) (i.e., number of nodes, degree, number of common neighbors with adjacent and non- adjacent nodes), but, just like 3-WL and node-based full-bag Subgraph GNNs, POLICY-LEARN cannot distinguish strongly regular graphs of the same parameters.
19

---

Published as a conference paper at ICLR 2024
Proposition 9 (POLICYLEEARN can distinguish certain strongly-regular graphs). POLICY-LEARN can distinguish any strongly regular graphs of different parameters, but cannot distinguish any strongly regular graphs of the same parameters.
Proof. The proof extends Bevilacqua et al. (2022, Lemma 14) to the node-marking generation policy, and follows the same steps. We show that the WL. algorithm converges to the same coloring of any subgraph obtained by marking a single node of a strongly regular graph, and this coloring depends only on the parameters of the strongly regular graph. Let G be strongly regular graph with parameters (n, k. A, p)) i.e., G has n nodes, is regular of degree k, any two adjacent nodes have common neighbors, any two non-adjacent nodes have l common neighbors. Consider the subgraph obtained by marking node u. We initialize all nodes the same color 1, except node which has color 2.
(Iter. 1) Nodes adjacent to colored 3, nodes non-adjacent to u are colored 4, u is colored 5,
(Iter. 2) Nodes adjacent to u colored 6, they each had: (i) color 3; (ii) À neighbors of color 3; (iii) (k: 1-1-A- A) neighbors of color 4; (iv) one neighbor of color 5. To see this, recall that these nodes are adjacent to u, so they have À common neighbors with u, and these nodes had color 3 at iteration 1 because they are indeed also neighbors with u. This also implies that the remaining (k
X) neighbors necessarily had color 4 at iteration 2.
Nodes non-adjacent to u are colored 7, as they each had: (i) color 4; (ii) μ neighbors of color 3; (iii) (k- - p)) neighbors of color 4. This is because they are non-adjacent to u, so they have exactly common neighbors with and these nodes had color 3 at iteration 2. Similarly, the remaining (k- vt) neighbors had color 4 at iteration
Node u is colored 8 as it had: (i) color 5; (ii) & neighbors of color 3.
Since the colors at iteration 2 are isomorphic to the colors at iteration 1, WL converged. Since this coloring is the same for any node-marked subgraph of any strongly regular graph of the same parameters, then POLICY-LEARN cannot distinguish strongly regular graph of the same parameters. Similarly, the coloring differs for strongly regular graphs of different parameters: at initialization if n differs, iteration if & differs, at iteration 2 if X or μ differs. Thus POLICY-LEARN can distinguish strongly regular graphs of different parameters.
Proposition 9 has important implications on the relation with higher order WL tests, as strongly regular graphs are instead distinguishable by 4-WL. Since there exist pairs of 4-WL distinguishable graphs that are not distinguishable with our approach, then POLICY-LEARN is not as powerful as 4-WI.
Finally, we remark here that it is also possible to consider other families of graphs, beyond the family of non-isomorppic (n, C)-C graphs (Definition 2), where marking limited number of nodes is sufficient for identifiability and therefore where POLICY-LEARN can attain the same expressive power of the full bag while using significantly reduced set of subgraphs. Indeed, our theoretical results are valid for any family of graphs, as we describe next, obtained from any collection of WL-indistinguishable graphs that become distinguishable when marking any node, c.g., strongly regular graphs of different parameters. Similarly the construction of the (n, C)-CSL family, each graph in these families is created by considering & disconnected, non-isomorphic instances in the corresponding collection.
E
ADDITIONAL EXPERIMENTS AND DETAILS
E.1 ADDITIONAL RESULTS ON THE OGB, ZINC AND REDDIT- BINARY DATASETS
We additionally compared POLICY-LEARN to a wide range of baselines, including existing full-bag Subgraph GNNs, on the OGB, ZINC and REDDIT- BINARY datasets presented in the main paper. We further report the random baseline using the prediction network considered in OSAN, whose results are taken from Qian et al. (2022) and which we denote as RANDOM-OSAN, The results on the various OGB datasets are presented in Table 4. Interestingly, POI.CC--LEARN consistently
20

---

Published as a conference paper at ICLR 2024
Table Comparison of POLICY-LEARN to standard MPNNS and full-bag Subgraph GNNs demon- strate the advantages of our approach. indicates the results was not reported in the original paper.
Method // Dataset
MOLESOL RMSE
MOLTOX21 ROC-AUC
MOLBACE ROC-AUC f
MOLHIV ROC-AUC
MPNNS
GCN (Kipf& & Welling, 2017)
1.114+0.036 1.73+00057
75.29+00.69 74.91
19.1 72.97:44.00
76.06=0.977
GIN (Xu et al., 2019)
75.58 1.40
FULL-BAG SUBGRAPH GNNS
RECONSTR. GNN (Cotta 2021)
1.026+00033
1..15+1.40
76.32
NGNN (Zhang & Li, 2021)
78.34=1.88
DS-GNN (EGO+) (Bevilacqua et al., 2022)
76.39:1118
77.40=2.19
DSS-GNN (EGO+) (Bevilacqua al., 2022)
77.95110.4
76.78=
GNN-AK+ (Zhao al., 2022)
79.61=1.19
SUN (EGO+) (Frasca et al., 2022)
80.03=0.55
GNN-SSWL+ (Zhang al., 2023)
79.58=0.35
FULL
0.847=0.015
76,25+1112
78.41=11.44
76.54=1.37
SAMPLING SUBGRAPH GNNS
RANDOM-OSAL OSAN (Qian et al., 2022) T 10
1.128:0.055
71.90:13.90
OSAN (Qian al., 2022)
0.084
72.30=16.00
RANDOM T
0.951
76.65
75.36:14.28
77.55
POLICY-LEARN T
0.877=0.029
77.47
78.40=
79.13=0.60
RANDOM T
0.900=0.032
76.62-10.63
78.1442..6
77.30
POLICY-LEARN T
0.883=0.032
77.36
78.3942.28
78.49=100
approaches outperforms computationalllyinttnnn -intennsv full-bag methods across these datasets. On the ZINC- 1K dataset (Table 5), POLICY-LEARN with 7 5 performs similarly to NGNN (Zhang & Li, 2021), which uses average 5x more subgraphs. Additional comparisons with other values of T. namely 7' 3 and 7 = 8, align with the observations made for T = 2 and T = 5. In particular, POLICY-LEARN always outperforms the random baseline, and the gap larger when the number of subgraphs T is smaller, where selecting the most informative subgraphs is more crucial. Finally, results on the REDDIT-BINARY cases where full-bag Subgraph GNNs, such our implementation FUL.L, well as DS-GNN and DSS-GNN (Bevilacqua al., 2022), are otherwise inapplicable. Notably POLICY-L.FAR surpasses the random counterparts that uniformly sample the same number of subgraphs, even when coupled with the DS-GNN or DSS-GNN downstream prediction model proposed Bevilacqua et al. (2022).
BINARY dataset (Table 6) demonstrate the applicability of our method to
E.2 TIME COMPARISONS
We performed timing analysis on REDDIT-BINARY and ZINC datasets, which we discuss in the following.
Timings on REDDIT-BINARY. further underscore the effectiveness of POLICY-LEARN on the REDDIT-BINARY dataset, conducted a time comparison, presented in Table 7. Specif- ically, we estimated the inference time on the entire test set using a batch size of 128 on an NVIDIA RTX A6000 GPU. We measure the times in two scenarios: when setting torch. www_uueeeeeeeeeeeeeeeeeeeeeeeeee n st itticclll eeallorrit goliiittmm tims (True) and without setting it (which equiva- lent to setting to se). While the former ensures deterministie results, it comes the cost of an increased running time, and it is not usually set in Subgraph GNNs implementations. Nonetheless, since determinism might be important in practice, we report both results. We compared POLICY- LEARN with the RANDOM baseline, maintaining the same hyperparameters for all methods for fair comparison. As expected, POLICY-LEARN takes longer than RANDOM with better results (see Table 6). To further verify its effectiveness, we ran RANDOM with T = 20, and observed that takes more than twice the time of POLICY-LEARN (and gets worse performance, see Table 6), showcasing how our approach is beneficial in practice.
2, also obtains
Timings ZINC-12K 12k. We report empirical runtimes the ZINC dataset in Table 8. For all methods, we estimated the training time on the entire training set as well as the inference time on the entire test using batch size of 128 NVIDIA RTX A6000 GPU. For all these experiments,
21

---

Published as a conference paper at ICLR 2024
Table 5: Comparison of POLICY-LEARN to existing expressive methods on the ZINC-12K graph dataset. All methods obey to the 500k parameter budget.
Method
ZINC (MAE 4)
MPNNs
GCN (Kipf & Welling, 2017)
0.301 0.163=00.004
GIN (Xu et al., 2019)
EXPRESSIVE GNNS
PNA (Corso al., 2020)
0.133=0.011
GSN (Bouritsas al., 2022)
0.101
CIN (Bodnar et al., 2021a)
0.079=#0.006
FULL--BAG SUBGRAPH GNNS
NGNN (Zhang & Li, 2021)
0.000
DS-GNN (EGO+) (Bevilacqua et al., 2022)
0,105=0,003
DSS-GNT ENS-GNN (EGO+) (Bevilacqua al., 2022)
0.097=0.006
GNN-AK (Zhao et al., 2022)
0.105=0.010
GNN-AK- (Zhao et al., 2022)
0.091+0.011
SUN (EGO+) (Frasca et al., 2022)
0.084+0.002
GNN-SSWI. (Zhang et al., 2023)
0.082+000033
GNN-SSWL+ (Zhang al., 2023)
0.07010.005
FULL
0.087110.003
SAMPLING SUBGRAPH GNNS
RANDOM-OSAL (Qian et al., 2022) T =
0.214+0.007
OSAN (Qian et al., 2022) T =
0.0161
RANDOM 7 =
0.006
POLICY-LEARN T
0. 0.120-10.003
RANDOM 3
0. 0.120:00004
POLICY-LEARN ''
0.116:0.008
RANDOM 7 -5
0.113=10.006
POLICY-LEEAN T=
0. 0.000
RANDOM =
0.102++0.003
POLICY-LEARN
0.097+000005
Table 6: Results on the REDDIT- BINARY dataset. POLICY-LEARN outperforms all baselines, includ- ing Subgraph GNNs coupled with random selection policies.
Method
RDT-B (ACC 1)
GIN (Xu et al., 2019)
92.4+2.5
SIN (Bodnar et al., 2021b)
92,21110
CIN (Bodnar et al.. 2021a)
92.412.1
RANDOM DS-GNN GNN (Bevilacqua et al., 2022) T =
91.3+1.66
RANDOM DSS-GNN (Bevilacqua et al., 2022) T -=
92.7+0088
FULL
00M
RANDOM T 20
92.6.1..5
RANDOM
92.4=1.0
POLICY-LEARN T 2
93.0=109
(False). During training, the runtime of POLICY-LEARN significantly closer to that of the RANDOM approach than to the one of the full-bag Subgraph GNN FULL, and POLICY-LEARN is faster than OSAN. At inference, POLICY-LEARN places in between RANDOM and FULL, and it significantly faster than OSAN while also achieving better results.
we set torch..chhh. dess__ddteeminnsttiie
Finally, we compare the time and the prediction performance on the ZINC dataset with the CIN model (Bodnar et al., 202 in Table 9, For all methods, report the inference time the entire test set using a batch size of 128. The runtime of CIN is taken from the original paper, and it is measured NVIDIA Tesla V100 GPU. To ensure fair comparison, therefore timed POLICY-LEARN
22

---

Published as a conference paper at ICLR 2024
Table 7: Timing comparison on the REDDIT-BINARY dataset on an RTX A6000 GPU (48 GB). Time taken at inference on the test set. Results on the left column are obtained when setting torch. use__ete seee_determinn aetic_algoritte (True), while those on the right column are without it. All values are in milliseconds.
RDT-B (Time (ms))
Method
det=True
det=False
FULL
OOM
0OM
RANDOM
T = 20
13
180.65::5.2
RANDOM
T = 216.7=00.4
39.7:12.3
POLICY-LEARN
T
411.7=0..6
7..41:12
Table 8: Timing comparison on the ZINC- 12K dataset on an RTX A6000 GPU (48 GB). Time taken train for epoch and inference the test All values are in milliseconds.
ZINC
Method
Train time (ms)
Test time (ms)
MAE 4
GIN (Xu el al., 2019)
1370.10110..99
84.81 1.0.26
0.163-10.004
FULL
4872.79: 14.30
197.38
0.087=00003
OSAN (Qian et al., 2022)
T = 2
227.93=#0.21
0.177=0.016
RANDOM
T 2
2114.00=227.88
107.02:+0.22
0.136=
POLICY-LEARN
2489.25+: 9.42
150.38
0. 0.000000.003
the same GPU type. First, we observe that POLICY-LEARN is faster than CIN, especially considering the additional preprocessing time required by CIN for the graph lifting procedures, which is not measured in Table 9. Second, although CIN outperforms POLICY-LEARY, it should be noted that CIN explicitly models cycles and rings, which are obtained in the preprocessing step and serve as domain-specific inductive bias. On the contrary, the subgraph generation policy in POLICY-LEARN (i.e., node-mon (add-marknng) is entirely domain-agnoste and not tailored to the specific application.
E.3 ADDITIONAL DATASETS
Alchemy. We conducted additional experiment considering the Alchemy-125 dataset (Chen et al., 2019), where we compared not only to the direct baselines FULL and RANDOM, but also positional and structural encoding augmented method GIN + RWSE & LAPPE, as reported in Qian et al. (2023). As can be seen from Table 10, these results further demonstrates the capabilities of POLICY-LEARN, which surpasses the random baseline RANDOM, and get results close those obtained by the FULL approach.
Expressivity Datasets. In effort to answer whether POLICY-LEARN maintains the expressive power of full-bag Subgraph GNNs in practice, we experimented on the CSL (Murphy et al., 2019; Dwivedi al., 2020) and EXP (Abboud et al., 2020) synthetic datasets, which are constructed in way that 1-WL. GNNs cannot outperform the performance a random guess (10% accuracy in CSL, 50% accuracy in EXP) We followed the evaluation procedure in Bevilacqua al. (2022), consisting of a k-fold cross validation (k 5 in CSL, = 10 in EXP). Results are reported in Table 11. On the CSL dataset, both RANDOM and POLICY-LEARN achieve 100% test accuracy, aligning with our theoretical analysis at the beginning of Section 4, where we discussed the disambiguation of these graphs. On the EXP dataset, POLICY-LEARN surpasses the performance of RANDOM, and achieves the same performance of the full-bag approach. This numerically confirms the importance of learning which subgraphs to select, rather than randomly sampling them.
Counting Datasets. We tested the abilities counting substructures on the benchmarks from Chen et al. (2020). which full-bag Subgraph GNNs have been extensively evaluated (Zhao et al., 2022: Frasca et al., 2022). We used the dataset splits and evaluation procedure of Zhao et al. (2022); Frasca et al. (2022) and reported the average across the seeds in Table 12. While POLICY-LEARN
23

---

Published as a conference paper at ICLR 2024
Table Timing comparison with CIN (Bodnar al., 2021a) the ZINC- dataset on an RTX A100 GPU. Time taken at inference on the test set. All values are in milliseconds.
ZINC
Method
Time (ms)
MAB +
GIN (Xu et al., 2019)
126.9110.82
0.163:10.004
CIN (Bodnar et al.. 2021a)
47100013.00
0.079-10.006
POLICY-LEARN
T = 2 T = 5
1.5.21
0.120=t0.003
POLICY-LEARN
1990199
0.109=00.005
Table 10: Comparison of POLICY-LEARN to existing methods on the ALCHEMY- 12k graph dataset.
Method
ALCHEMY (MAE L
MPNNS
GIN (Xu al. 2019)
11.12+0..90
GIN RWSE LAPPE (Dwivedi al., 2020)
7.197+00.09
FULL-BAG SUBGRAPH GNNs
FULL
6.65
SAMPLING SUBGRAPH GNNS
RANDOM-OSAL (Qian al., 2022) T =
12.11100.210
OSAN (Qian et al., 2022) T 10
8.87 187.1..120
RANDOM T =
1.9510
POLICY-LEARN T
6.8910.151
RANDOM
6.78
POLICY-LEARN T
6.72+t0100
Table 11: Results on the CSL and EXP expressivity datasets. POLICY-LEARN exhibits the same
disambiguation capabilities of the full-bag.
Method
CSL (ACC 1)
EXP (ACC 1)
GIN (Xu et al., 2019)
10.00=00.0
51.20++2.1
FULL
100.00=0.0
100.00+00.0
RANDOM T
100.00=0.0
89.9242.5
POLICY-LEARN
100.00-10.0
100.00-10.0
cannot achieve the same performance of the full-bag, it significantly surpasses the random baseline. Therefore, although counting substructures might not be possible with a significantly reduced set of subgraphs, there exist subgraph selections that enable more effective counting. Furthermore, increasing number of subgraphs proves advantageous in terms the counting power.
E.4 EXPERIMENTAL DETAILS
We implemented POLICY-LEARN using Pytorch (Paszke et al., 2019) and Pytorch Geometric (Fey & Lenssen, 2019). We ran our experiments on NVIDIA DGX V100, GeForce 2080, NVIDIA RTX A5000, NVIDIA RTX A6000, NVIDIA GeForce RTX 4090 and TITAN V GPUs. We performed hyperparameter tuning using the Weight and Biases framework (Biewald, 2020). We used mean aggregator to aggregate node representations across subgraphs the selection network f. Similarly, in the prediction network we use mean aggregator to obtain subgraph representations given the representations of nodes in each subgraph, and mean aggregator to obtaín graph representations given the subgraph representations. We always employ GIN layers (Xu et al., 2019) to perform message passing, and make use of Batch Normalization, with disabled computed statistics the selection model (as the number of subgraphs changes every t). Our MLPs are composed of two linear layers with ReLU non- con-linearities. Unless otherwise specified, we use residual connections. Each experiment repeated for different seeds. During evaluation, replace sampling with the arg max function and choose the node with the highest probability. Details of hyperparameter grid for each dataset can be found in the following subsections.
24

---

Published as a conference paper at ICLR 2024
Table 12:
Comparison of counting capabilities on the datasets Chen al. (2020). While reducing the set of subgraphs necessarily impacts the counting abilities of Subgraph GNNs, POLICY-LEARN significantly outperforms RANDOM.
Counting Substructures (MAE 4)
Method
Triangle
Tailed Tri.
Star
4-Cycle
GIN (Xu et al., 2019)
0.3569
0.2373
0,0224
0.2185
FULL
0.0183
0.0170
0.0106
0,0210
RANDOM
0.1841
0.1259
0.0105
0.1180
POLICY-LEARN T = 5
0.1658
0.1069
0.0100
0.0996
RANDOM T =8
0.1526
0.0999
0.0119
0.0933
POLICY-LEARN T =
0.1349
0.0801
0.0100
0.0793
E.4.1 OGB DATASETS
We considered the challenging scaffold splits proposed in Hu et al. (2020), and for each dataset we used the loss and evaluation metrie prescribed therein. For all models (FULL, RANDOM, POLICY- LEARN), we used Adam optimizer with initial learning rate of 0.001. We set the batch size to 128, except for the FUI.L method MOLBACE and MOLTOX21 where reduced it to 32 to avoid out-of-memory errors. We decay the learning rate by 0.5 every 300 epochs for all dataset except MOLHIY, where we followed the choices of Frasca et al. (2022), namely constant learning rate, downstream prediction network with 2 layers, embedding dimension 64 and dropout in between layers with probability in 0.5. We used 5 layers with embedding dimension 300 and dropout in between layers with probability in 0.5 in the prediction network for MOLTOX21 I and MOLBACE as prescribed by Hu et al. (2020). For MOLESOL we employed a downstream prediction network with layers in (2 and embedding dimension 64. In all datasets, for POLICY-LEARN we employed selection network architecturally identical to the prediction network, trained with a separate Adam optimizer without learning rate decay and learning rate 0.001. We tuned the temperature parameter T the Gumbel-Softma trick in {0.33, 0.66, 1,2]. To prevent overconfidence in the probability distribution over nodes, added dropout during train, with probability tuned in {0, 0.3, 0.5).
The maximum number of epochs is set to 1000 for all models and datasets except MOLHIV and MOLTOX21, where it is reduced to 500. The test metric is computed at the best validation epoch.
E.4.2 ZINC- 12K
We considered the dataset splits proposed in Dwivedi et (2020), and used Mean Absolute Error (MAE) both as loss and evaluation metric. For all models (FULL,, R RADDMM, POLICY-LEARR), we used a batch size of 128 and Adam optimizer with initial learning rate of 0.001, which is decayed by 0.5 every 300 epochs. The maximum number of epochs is set to 1000 for RANDOM and POLICY- LEARN when = 2, and to 600 for RANDOM and POLICY-LEARN when T 5, as well as for FULL, due the increase in training time. The test metric is computed at the best validation epoch. We used a downstream prediction network composed of 6 layers, embedding dimension 64. For POLICY-LEARN, we employed selection network architecturally identical to the prediction network, trained with a separate Adam optimizer without learning rate decay and learning rate 0.001. We further tuned the temperature parameter 7 of the Gumbel-Softmax trick in (0.33, 0.66, 1, 2). Early experimentation revealed that POLICY-LEARN became overconfident in the probability of specific nodes, thus preventing the exploration of other nodes through sampling. We mitigated the problem by adding dropout on the node probability distribution during train, with probability tuned in {0, 0.3, 0.5). For T 5 we additionally masked out nodes corresponding to subgraphs that had already been selected, to avoid repetitions and encourage exploration especially at training time. To ensure a fair comparison, we applied a similar mechanism to RANDOM with T = 5,, preventing it to sample the same subgraph (i.e., we did not allow replacement when uniformly sampling subset of subgraphs from the bag). For the additional experiments featuring other values of T,, namely T = and T = 8, the same strategy employed for T = 5.
25

---

Published as a conference paper at ICLR 2024
E.4.3 REDDIT-BINARY
We used the evaluation procedure proposed in Xu et al. (2019), consisting of 10-fold cross validation and metric the best averaged validation accuracy across the folds. The downstream prediction network is composed by 4 layers with embedding dimension 32 and no residual connections. We use Adam optimizer with learning rate tuned in {0.01, 0.001). For the selection network we use the same architectural choices adopted for the prediction network, and the same learning rate and optimizer type. We consider a batch size of 128, and trained for 500 epochs. We encouraged exploration during train through dropout on the node probability distribution, with probability tuned in {0, 0.5). Finally, we tuned the temperature r within the Gumbel-Softmal {0.33, 0.66, 1,
E.4.4 ALCHEMY
We used the same dataset splits of Morris et al. (2020b), and adopted the evaluation procedure followed by Qian et al. (2022), which differs from Morris et al. (2020b) only in the additional final denormalization of the predictions and ground-truths. For all models (FULL, RANDOM, POLICY- leArN), we used a batch size of 128 and Adam optimizer with initial learning rate of 0.001, which is decayed by 0.5 every 300 epochs. The maximum number of epochs is set to 1000, and the test metric used a downstream
is computed at the best validation metric. Following Morris et al. (2020b), prediction network composed of 6 layers and embedding dimension 64. For POLICY-LEARN, we employed a selection network architecturally identical to the prediction network, trained with separate Adam optimizer without learning rate decay and learning rate 0.001. We further tuned the temperature parameter T of the Gumbel-Softmax trick in {0.33, 0.66, 2]. and, to prevent overconfidence in the probability distribution over nodes, we added a dropout during train, with probability tuned in {0. 0.3, 0.5}. Similarly to ZINC, when T 5 mask out already selected subgraphs both for POLICY-LEARN and RANDOM, to avoid repeated subgraph selections.
E.4.5 EXPRESSIVITY DATASETS
Following Bevilacqua et al. (2022), we used the same training procedure and evaluation strategy considered for the REDDIT- BINARY datasets. The downstream prediction network is composed by layers with embedding dimension 64. For the EXP dataset, we additionally used residual connections and set to False the batch statistics instead of dataset ones during evaluation. We used Adam optimizer with learning rate tuned in {0.01, 0.001} for CSL and set to 0.0001 for EXP. For the selection network we used the same architectural choices adopted for the prediction network and the same optimizer type, with learning rate set to 0.001 for CSL and 0.0001 for EXP. We consider a batch size of for CSL and 128 for EXP, and trained for 100 epochs. We fixed the temperature within the Gumbel-Softmax to 1. We mask out already selected subgraphs both for POLICY-LEARN and RANDOM, to avoid repeated subgraph selections.
_stats flag in the batch normalization, thus employing
E.4.6 COUNTING DATASETS
We considered the dataset splits of previous work (Zhao al., 2022; Frasca et al., 2022), and used Mean Absolute Error (MAE) both as loss and evaluation metric. For all models (FULL, RANDOM, POLICY-LEARN). we used batch size of 128 and Adam optimizer with initial learning rate of 0.001. which is decayed by 0.5 every 300 epochs. The maximum number of epochs is set to 1000 and the test metric is measured at the best validation epoch. We used a downstream prediction network composed of 6 layers, embedding dimension 32. For POLICY-LEARN, we employed a selection network architecturally identical to the prediction network, trained with a separate Adam optimizer without learning rate decay and learning rate 0.001. We further tuned the temperature parameter of the Gumbel-Soffi -Softma trick in (0.33, 0.66, 1,2), 2}. and, to prevent overconfidence in the probability distribution over nodes, we added a dropout during train, with probability tuned in {0, 0.3, 0.5). For both POLICY-LEARY and RANDOM, we masked out already selected subgraphs, therefore preventing it to sample the same subgraph.
26

---

Published as a conference paper at ICLR 2024
F COMPLEXITY ANALYSIS
In this section we analyze the complexity of POLICY-LEARN and compare it with its full-bag Subgraph GNN counterpart. We analyze here an equivalent efficient implementation of our method, which differs from Algorithm in that it does not feed to the selection network the entire current bag at every iteration. Instead, it processes only the last subgraph added to the bag, and re-uses the representations of the previously selected subgraphs, passed though the selection network in previous iterations and stored. In what follows, we consider the feature dimensionality and the number of layers to be constants. We denote Amax the maximum node degree of the input graph, the number of its nodes and 7 the number of selected subgraphs.
Time complexity. The forward- pass asymptotic time complexity of the selection network f amounts to OT . - Amxx). This arises from performing T iterations, where each iteration involves selecting new subgraph by passing the last subgraph added to the bag through the MPNN, which has time complexity On Amax) and re- re-using stored representations of mreviously-selecte subgraphs. The forward- pass asymptotic time complexity of the prediction network g is O of the T + subgraphs is processed through the MPNN, which is different than the one used in the selection, in (n Amax). Note that, differently from the selection network, the prediction network considers I' + 1 subgraphs as it also utilizes the last selected subgraph. Therefore, POLICY-LEARN has an overall time complexity of O((T' + (T' + 1)) Amax). i.c., O(T' Amax).
Amax). each
Space complexity. The corward-pass asymptotic space complexity of the selection network f is OT +nnnn-+n+-n---n m)x)), because we need store n node features for each of the T subgraphs, and, at each iteration, the space complexity of the MPNN on the subgraph being processed is O Amax) store its node features and connectivity in memory. Similarly, the forward pass asymptotic space complexity of the prediction network g is (((T + +1) (n + n Amx)). Therefore, POLICY-LEARN has an overall space complexity of O(T (n + n Amax))
Notably, POLICY-LEARN is advantageous in both time and space complexity when compared to full-bag Subgraph GNNs, Indeed, as studied in Bevilacqua et al. (2022), the time complexity of full-bag node-based Subgraph GNNs amounts to O(n² Amax), while the space complexity is O(n-(n (2+ since the number of subgraphs in the full bag is exactly n. Since POLICY-LEARN scales linearly with T, when small values of T are used, in our paper, this results in speedup by factor of n compared to full-bag Subgraph GNNs.
To grasp the impact of this reduction, let consider for example the REDDIT-BINARY dataset, where the average number of nodes per graph is approximately n = 429, while the number of selected subgraphs is T = 2. The time and space complexities drastically reduced, because T < and indeed POLICY-LEARN can be effectively run while the full bag Subgraph GNN is computationally infeasible.
27


Reviews:
[{"review_id": "x6CRX897Ko", "reviewer": "Reviewer_VjjR", "summary": "Subgraph GNNs generally refers to GNN methods that run GNNs on several subgraphs obtained from the input graph. Recently, a variety of such methods, differing from each other in the subgraph selection policies, are proposed. Of particular relevance to the current paper is the OSAN framework by Qian et al, in which the subgraph selection policy is learned. In the current paper, a more expressive GNN architecture DS-GNN is used (rather than a classical GNNs) of the subgraph selection policy. This DS-GNN provides a distribution from which is sampled using standard Gumbel softmax trick. It is shown theoretically the approach can be more powerful than the OSAN approach.", "questions": "**Q1** Could you explain the histograms in Figure 1 after labeling one vertex? \n\n**Q2** What is the ingredient of the method which results in more power than say OSAN?\n\n**Q3** Is the proposed method at least as powerful as OSAN or incomparable? What about comparisons with other subgraph formalisms?\n\n**Q4** A number of subgraphs are selected in order to reduce complexity. However, the DS-GNN method used for selection policy relies on all subgraphs? What is the overall complexity of the method?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. Research in how to learn subgraph selection policies is highly relevant in view of the popularity of subgraph GNN approaches.\n\n2. Related work is well described.\n\n3. The policy learn method is well designed.\n\n4. Theoretical guarantees over special classes of CSL graphs are presented. In particular, it is argued that the proposed approach can be stronger than a previous approach.", "weaknesses": "1. It is not clearly described what gives the proposed method more power than e.g., OSAN.\n\n2. The method seems to depend on a subgraph GNN method (DS-GNN) which high computational cost.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SjMiaHv7IP", "reviewer": "Reviewer_gwbR", "summary": "**TLDR**: The paper proposes a learnable subsampling policy to reduce the number of sampled subgraphs in subgraph GNNs.\n\nSubgraph GNNs refer to the family of message-passing graph neural networks which sample subgraphs to improve their expressive power. The paper introduces a new subsampling strategy, _Policy-Learn_ which aims to reduce the number of subgraphs needed for subgraph GNNs. _Policy-Learn_ consists of two subgraph GNNs, one selection network $f$ and one policy network $g$: $f$ learns a distribution over the nodes of the input graph to select which subgraph to sample next, $g$ takes the sampled graphs as input and performs a prediction task. The paper motivates _Policy-Learn_ by considering the graph class of $(n,\\ell)$-CSL graphs, where one instance consists of $\\ell$ disconnected, non-isomorphic CSL graphs on $n$ nodes. For $(n, \\ell)$-CSL graphs, _Policy-Learn_ only needs to sample $\\ell$ subgraphs in comparison to random subsampling and the existing subsampling strategy OSAN. In an experimental evaluation, _Policy-Learn_ outperforms OSAN and is, on average, competitive with the presented baseline methods.", "questions": "1. **Expressiveness**: _Policy-Learn_ can distinguish non-isomorphic instances in the graph class $(n, \\ell)$-CSL, which are indistinguishable by 1-WL. What about the opposite? Can you characterize graph classes whose non-isomorphic instances are provably indistinguishable by _Policy-Learn_? Are there graph classes where OSAN is stronger than _Policy-Learn_? Is _Policy-Learn_ limited by higher-order WL?\n2. **Assumption in proof of Theorem 4**: Is the (necessary) assumption that $f$ has $n$ layers feasible for larger graphs? Do you have experimental results on $(n, \\ell)$-CSL graphs?\n3. **Extension of theoretical results**: Have you thought about extensions of your theoretical results, e.g., other graph classes where marking any node in a graph is sufficient or where marking a limited number of nodes is sufficient?\n4. **Experiments**:\n\n    a. How did you choose the values for $T$ (2 and 5 in Tables 1 and 3, 2 and 20 in Table 2)?\n\n    b. In Section 7, paragraph ZINC: \"Notably, OSAN performs worse than our random baseline due to differences in the implementation of the prediction network\". Could you elaborate on the differences and why this affects the performance?\n5. **Time comparison**: How does _Policy-Learn_ compare with respect to time vs. prediction performance in comparison to more expressive GNNs (e.g., GSN, CIN)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The proposed subsampling strategy _Policy-Learn_ is well-motivated and novel.\n* _Policy-Learn_ can provably sample subgraphs such that non-isomorphic instances can be distinguished for one specific graph class ($(n, \\ell)$-CSL graphs) on which 1-WL fails.\n* In the experimental evaluation, _Policy-Learn_ is competitive with most baseline methods and outperforms OSAN", "weaknesses": "* Theoretical limitations: While the presented theoretical results are novel and interesting, they also appear to be limited. The artificially constructed graph class $(n, \\ell$)-CSL is 1-WL indistinguishable; however, higher-order models and GNN variants are able to distinguish them. A more comprehensive analysis of the expressive power of _Policy-Learn_ could strengthen the contribution significantly.\n* Clarity: Although the paper is generally well-written, more precise language would improve readability:\n     * \"[...] preventing the applicability of Subgraph GNNs on important datasets\" -> What datasets are important?\n    * \"Contributions: [...] An experimental evaluation of the new approach demonstrating its advantages.\" -> It would be more informative if the advantages are specified.\n    * \" [...] which includes feature aggregation in light of the alignment of nodes [...]\" -> Could you specify what that means\n    * \"[...] and demonstrate that our framework performs better on real-word datasets\" -> Better than what?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1hGKGMWkxv", "reviewer": "Reviewer_w3Fw", "summary": "The paper focuses on learning effective subgraph selection policies for subgraph GNNs. In particular, it is inspired by an observation that only a small number of subgraphs are needed to differentiate a family of non-isomorphic graphs called the CSL graph. Based on the observation, it constructs a learning-based subgraph selection policy and surpasses previous works on various benchmarks.", "questions": "Definition 2 and Theorem 1 could possibly lead to some misunderstandings. For example, from my perspective, if the multiset $\\\\{k_i| i\\in \\\\{1,\\dots, l \\\\}\\\\}$ is the same, then $CSL(n,(k_1, …, k_l))$ should be isomorphic to each other. The observation should be pointed out, since the definition now seems that the sequence of k_i might also lead to non-isomorphism. In addition, the fact that $CSL (n, k)$ is isomorphic to $CSL (n, n-k)$ would also influence the isomorphism between graphs, which also need to be mentioned.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "Generally speaking, I like the efforts on subgraph sampling since the efficiency of subgraph GNNs limits them from being applied in real-world scenarios. In addition, the work is generally motivated and well-written.", "weaknesses": "1. Most subgraph sampling strategies face a problem: they cannot guarantee permutation invariance, i.e., generate the same representation for the same graph no matter how the graph is permuted. It seems that the proposed method also cannot guarantee such property as well.\n\n2. I appreciate the efforts in distinguishing the CSL graphs. However, CSL graphs are just a family of regular graphs that cannot be differentiated by 1-WL. Have you analyzed some other families, for example, the strongly regular graphs proposed in (Bodnar et al, 2021b) or some pairs of graphs that are mentioned in (Wang and Zhang 2023)?\n\n3. It seems that the time evaluation is only provided in Table 7, where the time of the full subgraph GNN is not provided due to OOM. I recommend adding time evaluation on more datasets, and reporting the time of \"policy learn\", \"random selection\", the full subgraph GNN, and MPNNs. For example, you can report the time of GIN, OSAN, FULL, RANDOM, and POLICY-LEARN on ZINC.\n\n(Wang and Zhang 2023) Wang Y, Zhang M. Towards Better Evaluation of GNN Expressiveness with BREC Dataset. arXiv, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dJAHUnMCRP", "reviewer": "Reviewer_ew9j", "summary": "This paper is motivated by the high computational cost of subgraph GNNs and the bag-of-subgraph contains many redundant information and proposed a learnable method to efficiently select a fixed number of subgraphs for downstream prediction. The proposed method achieves a good balance between the cost and the expressive power. It obtains reasonable performance across various datasets.", "questions": "1. The current method only works for node-based subgraphs. Could the proposed method be generalized to other policies like edge-based [1] or node-tuple-based [2] subgraphs?\n\n\n[1] Huang et al., Boosting the cycle counting power of graph neural networks with i$^2$-GNNs, ICLR23.\n\n[2] Qian et al., Ordered subgraph aggregation networks, Neurips22.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is overall sound and easy to follow. \n2. The motivated CSL example is clear.\n3. The proposed method achieves reasonable results across various datasets.", "weaknesses": "The main weaknesses of the paper lie in its insufficient theoretical analyses and experiment validation. Specifically:\n1. The proposed method is overall an extension to the 1-MLE method in the OSAN, which is not a big problem to me. However, what I would expect is a more in-depth analysis of the proposed method. The authors only use a single category of graphs (CSL and its $(n, l)$ extension) to show that the proposed method is more powerful than a random policy and OSAN, which is trivial from my perspective. Instead of comparing it with random policy, a more interesting question would be: How well does the proposed method compare to the full-bag version? Can it achieve the same expressive power as the full-bag version and in how much $T$ can it be from a theoretical perspective? \n2. To distinguish non-isomorphic graphs is only the first step towards an expressive GNN. Can the model successfully encode the structure information like counting cycles play a more important role in real-world tasks. Full-bag versions of subgraph GNNs have a much better ability to encode sub-structures [1]. I am wondering can the proposed method maintain its ability to encode sub-structures.\n3. Some commonly used synthetic datasets for comparing expressive power are missing (e.g. EXP [2], CSL [3]). This could be part of the answer to weakness 1.\n4. Some commonly used synthetic datasets for evaluating the counting power of GNNs are missing [1]. This could be part of the answer to weakness 2.\n5. The main contribution of the subgraph sampling is its lower computational cost compared to the full-bag version. However, the authors only show a simple comparison of the inference time using one dataset where the full-bag version is OOM (Table 7). I believe a more comprehensive comparison between the full-bag version and the proposed method is required. What is the time and memory cost of the proposed method compared to the full-bag version in both the train and the test? How does the cost vary if we increase the $T$? \n\n[1] Huang et al., Boosting the cycle counting power of graph neural networks with i$^2$-GNNs, ICLR23.\n\n[2] Abboud et al., The surprising power of graph neural networks with random node initialization. IJCAI21.\n\n[3] Murphy et al., Relational pooling for graph representations. ICML19.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)