Paper: Efficient Subgraph GNNs by Learning Effective Selection Policies (ICLR-2024)

Reviewer comment (Experiments):
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?

Author response:
**A1**: We thank the Reviewer for their comment, this gives us the opportunity to further clarify why we believe Policy-Learn is significantly different from OSAN. Our work is primarily motivated by scenarios where a subset of subgraphs is sufficient for maximal expressive power, while the I-MLE method in OSAN aims to make higher-order generation policies (where each subgraph is generated from tuples of nodes) practical. Furthermore, our method generates the bag sequentially, compared to their one-shot generation, and we parameterize the selection network using a Subgraph GNN, that is more expressive than the MPNN used in OSAN. Finally, we do not rely on I-MLE but instead use the simpler implementation of a Gumbel-Softmax trick to enable gradient backpropagation through the discrete sampling process. We proved that these choices lead to a framework that can learn subgraph distributions that cannot be expressed by OSAN.

**A2**: We thank the Reviewer for asking this important question. In certain cases, Policy-Learn can achieve the same expressive power of the full-bag, but with a significantly fewer number of subgraphs.

This is the case for CSL graphs, where only one subgraph is sufficient to achieve the same expressive power of the full bag, as we showed in the paper (Figure 1). This result can be extended to other exemplary non-isomorphic WL-indistinguishable pairs, for example, the pair where one graph consists of two triangles connected by an edge and the other of two squares sharing an edge (Figure 1 in Bevilacqua et al., 2022). Additionally, one subgraph is also sufficient to distinguish strongly-regular graphs of different parameters, which are distinguishable by the full-bag but not by 1-WL.

More interestingly, in certain cases only a _careful selection_ of subgraphs can lead to the same expressive power of the full bag, and our method can provably implement these selections. An example of this case is the family of (n, $\ell$)-CSL graphs where only $\ell$ subgraphs are required to attain the expressive power of the full-bag version, which instead employs $n \cdot \ell$ subgraphs. This example can further be generalized to other families of graphs, obtained from any collection of WL-indistinguishable graphs that become distinguishable when marking any node, e.g., strongly regular graphs of different parameters. Each graph in these families is created by considering $\ell$ disconnected, non-isomorphic instances in the corresponding collection.

Reviewer comment (Experiments):
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.

Author response:
**A3**: The Reviewer is raising an interesting point of discussion. The ability of Subgraph GNNs to count substructures can be attributed to the marking of the root nodes, which breaks the anonymity of the marked nodes and allows counting around each marked node (Huang et al., 2023). Interestingly, however, it is possible to show that it is not necessary for a node to be marked in order to count substructures around it: it is sufficient instead that its embedding is unique in its neighborhood. This might happen not only when the node is marked, but also when it receives marking information in the message passing, coming from a root node, and this information is still unique in its surroundings. This implies that our method can exhibit the same counting power of the full-bag whenever each node has an embedding that is unique in its neighborhood, due to the marking information coming from a few selected subgraphs. Note, however, that this might not always be the case, depending on the structure of the graph. For example, it is possible that certain parts of the graph do not see any marking information if we select a small number of subgraphs. This might happen to nodes that are distant from the marked roots of the selected subgraphs, around which a Subgraph GNN necessarily behaves as an MPNN.

We believe that this marks a first step towards the understanding of the counting ability of Subgraph GNNs that rely on a smaller number of subgraphs. However, properly enquiring about this aspect and designing Subgraph GNNs with counting guarantees when considering a subset of subgraphs fall outside the scope of the present work.

Reviewer comment (Experiments):
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.

Author response:
**A4**: We welcome the Reviewer's suggestion and additionally experimented with the CSL and the EXP datasets, as shown in the following.


| Method                         |             CSL (ACC)            |             EXP (ACC)           | 
|--------------------------------|:----------------------------------:|:-----------------------------------:|
|   GIN                             |          10.00 +- 0.0              |            51.2 +- 2.1              | 
|                                      |                                           |                                            |
|   FULL                          |        100.00 +- 0.0              |         100.00 +- 0.0             | 
|   RANDOM $T = 2$      |        100.00 +- 0.0              |           89.92 +- 2.5             | 
|   Policy-Learn $T = 2$  |        100.00 +- 0.0              |         100.00 +- 0.0             | 

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.

Reviewer comment (Experiments):
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.

Author response:
**A5**: Following the Reviewer’s suggestion, we have additionally tested Policy-Learn on synthetic datasets that measure the ability of models to count substructures. We chose the dataset in Zhengdao et al., 2020, on 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 previous work and reported the average across the seeds in the Table below.




| Method                         |       Triangle       |         Tailed Tri.    |         Star       |        4-Cycle          | 
|--------------------------------|:---------------------:|:---------------------:|:------------------:|:------------------------:|
|   GIN                             |       0.3569        |         0.2373         |        0.0224    |       0.2185             | 
|                                      |                          |                             |                       |                               |
|   FULL                           |      0.0183         |        0.0170         |         0.0106   |        0.0210            |
|                                      |                          |                             |                       |                               |
|   RANDOM $T = 5$      |       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 = 8$  |       0.1349        |         0.0801         |        0.0100    |        0.0793            |


While Policy-Learn cannot achieve the same performance of the full-bag, it _significantly_ surpasses the random baseline. This supports our previous discussion on the counting power of our method: although counting all substructures might not be possible with a significantly reduced set of subgraphs, there exist subgraph selections that enable more effective countings. Furthermore, increasing the number of subgraphs proves advantageous in terms of the counting power.

Reviewer comment (Experiments):
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$?

Author response:
**A6**: We believe that to adequately answer this question one should look at the time and space complexity of these methods, which we discuss in the following and we will include in the next manuscript revision. In short, our method scales linearly with $T$. When small values of $T$ are used, as in our paper, this results in a speedup by a factor of $n$ compared to the full-bag Subgraph GNN.

We analyze an equivalent efficient implementation of our method, which differs from Algorithm 1 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. We denote $\Delta_\text{max}$ the maximum node degree of the input graph.

**Time Complexity.** The forward-pass asymptotic time complexity of the selection network $f$ amounts to $\mathcal{O}(T \cdot n \cdot \Delta_\text{max})$. This arises from performing $T$ iterations, where each iteration involves selecting a new subgraph by passing the last subgraph added to the bag through the MPNN, which has time complexity $\mathcal{O}(n \cdot \Delta_\text{max})$ and re-using stored representations of previously-selected subgraphs. Similarly, the forward-pass asymptotic time complexity of the prediction network $g$ is $\mathcal{O}((T+1) \cdot n \cdot \Delta_\text{max})$, as each of the $T+1$ subgraphs is processed through the MPNN. Note that, differently from the selection network, the prediction network considers $T+1$ subgraphs as it also utilizes the last selected subgraph.
Therefore, Policy-Learn has an overall time complexity of $\mathcal{O}((T + (T+1)) \cdot n \cdot \Delta_\text{max})$, i.e., $\mathcal{O}(T \cdot n \cdot \Delta_\text{max})$. The time complexity of full-bag node-based Subgraph GNNs amounts instead to $\mathcal{O}(n^2 \cdot \Delta_\text{max})$, as there are $n$ subgraphs in total.

**Space complexity.** The forward-pass asymptotic space complexity of the selection network $f$ is $\mathcal{O}(T \cdot n + (n +  n \cdot \Delta_\text{max}))$, because we need to 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 $\mathcal{O}(n +  n \cdot \Delta_\text{max})$ to store its node features and connectivity in memory. 
Similarly, the forward-pass asymptotic space complexity of the prediction network $g$ is $\mathcal{O}((T+1) \cdot (n +  n \cdot \Delta_\text{max}))$. Therefore, Policy-Learn has an overall space complexity of $\mathcal{O}(T \cdot (n +  n \cdot \Delta_\text{max}))$. The space complexity of the full-bag approach instead amounts to $\mathcal{O}(n \cdot (n +  n \cdot \Delta_\text{max}))$, since the number of subgraphs in the full bag is exactly $n$.

To grasp the impact of this reduction, consider 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 are drastically reduced as $T \ll n$, and indeed Policy-Learn can be effectively run while the full bag Subgraph GNN is computationally infeasible.

Finally, we additionally report empirical runtimes on the ZINC dataset in the following. For all methods, we estimated the training time on the entire training set as well as the inference time on the entire test set using a batch size of 128 on an NVIDIA RTX A6000 GPU. 

| Method  |  ZINC (Train time (ms))  |  ZINC (Test time (ms)) |  ZINC (MAE) | 
|--|:--:|:--:|:--:|
|   GIN                             |     1370.10 +- 10.79      |          84.81 +- 0.26         |    0.163 +- 0.004  |
|                                      |                                      |                                         |                             |
|   OSAN   $T = 2$          |      2964.46 +- 30.36    |           227.93 +- 0.21      |  0.177 +- 0.016    | 
|                            |                                      |                                         |                             |
|   FULL                          |      4872.79 +- 14.30     |          197.38 +- 0.30        |  0.087 +- 0.003    | 
|   RANDOM $T = 2$      |     2114.00 +- 27.88      |         107.02 +- 0.22        |  0.136 +- 0.005    |
|   Policy-Learn $T = 2$  |     2489.25 +- 9.42       |         150.38 +- 0.33         |  0.120 +- 0.003    | 

During training, the runtime of Policy-Learn is 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 is significantly faster than OSAN while also achieving better results.

Reviewer comment (Theory):
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?

Author response:
**A7**:  Yes, we can extend our approach to these cases. As for edge-based policies, we can employ the same architecture and output a probability distribution of the edges at the last layer, by using any GNN layer coupled with a final predictor to output edge features. Regarding node $k$-tuples, one option would be to output a distribution on the space of $k$-tuples. However, for $k>2$ this may become impractical in large graphs. Alternatively, nodes in each tuple can be sampled sequentially based on a probability distribution over the nodes, conditioned on previous selections. This strategy can be easily implemented within our framework and it is scalable to large $k$ values even on large graphs.

Reviewer comment (Theory):
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.

Author response:
Our architecture is invariant as a function. However, the use of non-deterministic sampling makes the input to the network stochastic and this may result in different representations for the same graph, similar to other approaches that involve randomness in the GNN computation. In practice, however, the non-deterministic nature of the sampling does not seem to hurt the performance, as can be seen in the experimental section.

Reviewer comment (Evaluation):
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.

Author response:
We welcome the Reviewer’s suggestion and we report the runtime of these methods on the ZINC dataset in the following. For all methods, we estimated the inference time on the entire test set using a batch size of 128 on an NVIDIA RTX A6000 GPU. 


| Method                         |      ZINC (Time (ms))          |        ZINC (MAE)            | 
|--------------------------------|:-----------------------------------:|:-------------------------------:|
|   GIN                             |      84.81 +- 0.26                |       0.163 +- 0.004         |
|                                      |                                           |                                       |
|   OSAN   $T = 2$          |       227.93 +- 0.21             |      0.177 +- 0.016          | 
|                                      |                                           |                                       |
|   FULL                          |        197.38 +- 0.30            |      0.087 +- 0.003          | 
|   RANDOM $T = 2$      |       107.02 +- 0.22            |      0.136 +- 0.005           |
|   Policy-Learn $T = 2$  |       150.38 +- 0.33            |      0.120 +- 0.003           | 


Notably, Policy-Learn is significantly faster than OSAN, while also obtaining better results. Importantly, Policy-Learn places in between the RANDOM approach and the full-bag Subgraph GNN FULL. Furthermore, Policy-Learn takes around 2x the time of the corresponding GIN baseline, but significantly outperforms it.

Reviewer comment (Theory):
**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?

Author response:
We thank the Reviewer for bringing this point up to our attention. First, we would like to remark that the number of nodes in a (n,$\ell$)-CSL graph is $n \cdot \ell$, which is much larger than $n$, the value we used as a loose bound on the number of layers in the proof. However, we understand the Reviewer’s concern, and we have therefore now proved that the bound can be tightened to $n/2$ by considering the circular structure of the connected components. We believe that this bound can be further reduced by taking into account the skip connections.

Reviewer comment (Theory):
**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?

Author response:
In the paper, we showed that for CSL graphs it is sufficient to mark any single node for distinguishability (Figure 1). However, this result can be extended to other examples of  non-isomorphic WL-indistinguishable pairs, for example, the pair where one graph consists of two triangles connected by an edge and the other of two squares sharing an edge (Figure 1 in Bevilacqua et al., 2022). Additionally, marking any node is also sufficient to distinguish strongly-regular graphs of different parameters.

Importantly, it is also possible to consider other families of graphs beyond the (n, $\ell$)-CSL graph family where marking a limited number of nodes is sufficient for disambiguation. 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, e.g., strongly regular graphs of different parameters. Similarly to the construction of the (n, $\ell$)-CSL family in our paper, each graph in these families is created by considering $\ell$ disconnected, non-isomorphic instances in the corresponding collection. Finally, we believe that it is possible to extend this analysis to certain families of connected graphs, and we leave this aspect for future work.

Reviewer comment (Experiments):
**Experiments**: How did you choose the values for $T$ (2 and 5 in Tables 1 and 3, 2 and 20 in Table 2)?

Author response:
Experimentally, our goal was to show that it is possible to obtain compelling results on standard benchmarks even when considering a significantly smaller number of subgraphs than the total typically employed by existing methods. Therefore we chose $T=2$ and $T=5$ as two exemplary numbers, and demonstrated that even with only $T=2$ subgraphs Policy-Learn achieves very good performance. Nonetheless, we have further experimented with other values of $T$, namely $T=3$ and $T=8$, which we report in the following. 


| Method                       |          ZINC (MAE)             |
|------------------------------|:-----------------------------------:|
|   GIN                           |        0.163 +- 0.004            |
|                                    |                                            |
|   FULL                        |         0.087 +- 0.003           |
|   RANDOM $T = 3$    |        0.128 +- 0.004           |
|   Policy-Learn $T = 3$|         0.116 +- 0.008          |
|   RANDOM $T = 8$    |         0.102  +- 0.003         |
|   Policy-Learn $T = 8$|         0.097 +- 0.005          |

These results align with the observations made for other values of $T$. In particular, Policy-Learn always outperforms the random baseline, and the gap is larger when the number of subgraphs $T$ is smaller, where selecting the most informative subgraphs is more crucial.

Reviewer comment (Experiments):
**Experiments**: 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?

Author response:
We ran our random baseline on exactly the same prediction network that we used for Policy-Learn to ensure a direct comparison. This comprises architectural choices that are different from those taken in OSAN, including the number of layers, the embedding dimension, and the implementation of the residual connections. In our next revision, we will make this point clearer and further include the performance of the RANDOM baseline using the same prediction network considered in OSAN, as reported by Qian et al., 2022.

Reviewer comment (Novelty):
It is not clearly described what gives the proposed method more power than e.g., OSAN.

Author response:
**Q2**: What is the ingredient of the method that results in more power than say OSAN?

**A2**: The iterative procedure that selects one subgraph at a time based on previous selections, paired with the Subgraph GNN used to implement the selection network $f$, are the ingredients that result in more power than OSAN. 

Indeed, recall that OSAN selects all subgraphs at the same time, by applying an MPNN over the original graph only. Consider again the family of non-isomorphic (n, $\ell$)-CSL graphs we studied in the paper, which contains WL-indistinguishable graphs composed of disconnected copies of WL-indistinguishable (sub-)graphs. Given that OSAN uses the original graph, then it cannot differentiate between the different CSL (sub-)graphs and will return a uniform probability distribution over all the nodes. This implies that it cannot ensure that the marked nodes will belong to different CSL (sub-)graphs, which represents a necessary condition for identification (Proposition 2). On the contrary, Policy-Learn iteratively adds a new subgraph based on the subgraphs already present in the bag. Whenever a subgraph is in the bag, then all nodes in the CSL (sub-)graph of its marked root node will have a different color than the others. Thus it is always possible to distinguish nodes belonging to CSL (sub-)graphs having a marked node, and therefore the final MLP can assign these nodes a zero probability while maintaining a uniform probability over the remaining nodes. This implies that the next node to be marked will be sampled from a CSL (sub-)graph that has no marked node yet, effectively implementing the efficient policy $\pi$ which is sufficient for identification.

Reviewer comment (Evaluation):
The method seems to depend on a subgraph GNN method (DS-GNN) which high computational cost.

Author response:
**Q4**: The DS-GNN method used for selection policy relies on all subgraphs? What is the overall complexity of the method?

**A4**: The DS-GNN method used for the selection policy *does not* rely on all subgraphs. On the contrary, at each timestep $t$, it uses only the current bag, which contains $t$ subgraphs, to select the next subgraph to be added to the bag, starting from the original graph. This means that the number of subgraphs passed through the selection network is significantly smaller than the total number.

Nonetheless, we acknowledge the fact that including the computation complexity of the method would improve the quality of the manuscript. Thus, we present it below and we will include it in the next revision. 


We analyze an equivalent efficient implementation of our method, which differs from Algorithm 1 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 $\Delta_\text{max}$ the maximum node degree of the input graph, $n$ the number of its nodes and $T$ the number of selected subgraphs.

The forward-pass asymptotic time complexity of the selection network $f$ amounts to $\mathcal{O}(T \cdot n \cdot \Delta_\text{max})$. This arises from performing $T$ iterations, where each iteration involves selecting a new subgraph by passing the last subgraph added to the bag through the MPNN, which has time complexity $\mathcal{O}(n \cdot \Delta_\text{max})$ and re-using stored representations of previously-selected subgraphs. Similarly, the forward-pass asymptotic time complexity of the prediction network $g$ is $\mathcal{O}((T+1) \cdot n \cdot \Delta_\text{max})$, as each of the $T+1$ subgraphs is processed through the MPNN. Note that, differently from the selection network, the prediction network considers $T+1$ subgraphs as it also utilizes the last selected subgraph.

Therefore, Policy-Learn has an overall time complexity of $\mathcal{O}((T + (T+1)) \cdot n \cdot \Delta_\text{max})$, i.e., $\mathcal{O}(T \cdot n \cdot \Delta_\text{max})$. The time complexity of full-bag node-based Subgraph GNNs amounts instead to $\mathcal{O}(n^2 \cdot \Delta_\text{max})$, as there are $n$ subgraphs in total.
 
To grasp the impact of this reduction, consider the REDDIT-BINARY dataset, which we have experimented on to make this point clear (Tables 2 and 7). In this dataset the average number of nodes (and thus subgraphs) per graph is approximately $n = 429$, making the full-bag Subgraph GNN infeasible. In our experiment the number of selected subgraphs is $T=2$. The time complexity is drastically reduced as $T \ll n$, and, indeed, Policy-Learn can be effectively and successfully run.

Reviewer comment (Presentation):
Could you explain the histograms in Figure 1 after labeling one vertex?

Author response:
**Q1**: Could you explain the histograms in Figure 1 after labeling one vertex?

**A1**: For each original graph, we consider the application of the WL test to the subgraph obtained by marking the crossed node. More precisely, at time step 0, 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 0, 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 1, the marked node has a 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 a 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 1 orange node. Importantly, since the WL test represents a necessary condition for graph isomorphism, different histograms imply that the two graphs are not isomorphic, and thus the WL test distinguishes them. Figure 1 shows that marking only one node is sufficient for distinguishing the two graphs, which are instead indistinguishable when no node is marked, as the WL test returns the same histogram for them (13 orange nodes).

Reviewer comment (Experiments):
What is the ingredient of the method which results in more power than say OSAN?

Author response:
**Q2**: What is the ingredient of the method that results in more power than say OSAN?

**A2**: The iterative procedure that selects one subgraph at a time based on previous selections, paired with the Subgraph GNN used to implement the selection network $f$, are the ingredients that result in more power than OSAN. 

Indeed, recall that OSAN selects all subgraphs at the same time, by applying an MPNN over the original graph only. Consider again the family of non-isomorphic (n, $\ell$)-CSL graphs we studied in the paper, which contains WL-indistinguishable graphs composed of disconnected copies of WL-indistinguishable (sub-)graphs. Given that OSAN uses the original graph, then it cannot differentiate between the different CSL (sub-)graphs and will return a uniform probability distribution over all the nodes. This implies that it cannot ensure that the marked nodes will belong to different CSL (sub-)graphs, which represents a necessary condition for identification (Proposition 2). On the contrary, Policy-Learn iteratively adds a new subgraph based on the subgraphs already present in the bag. Whenever a subgraph is in the bag, then all nodes in the CSL (sub-)graph of its marked root node will have a different color than the others. Thus it is always possible to distinguish nodes belonging to CSL (sub-)graphs having a marked node, and therefore the final MLP can assign these nodes a zero probability while maintaining a uniform probability over the remaining nodes. This implies that the next node to be marked will be sampled from a CSL (sub-)graph that has no marked node yet, effectively implementing the efficient policy $\pi$ which is sufficient for identification.

Reviewer comment (Theory):
Is the proposed method at least as powerful as OSAN or incomparable? What about comparisons with other subgraph formalisms?

Author response:
**Q3**: Is the proposed method at least as powerful as OSAN or incomparable? What about comparisons with other subgraph formalisms?

**A3**: The proposed method Policy-Learn is _strictly more powerful_ than OSAN, although the term is more convoluted in this context due to the probabilistic nature of the sampling procedure.

More formally, for any instantiation of OSAN, there exists a set of weights for Policy-Learn such that it outputs exactly the same probability distribution. In other words, Policy-Learn can parameterize all the probability distributions that OSAN can. The contrary, however, is not true: there exist probability distributions that Policy-Learn can parameterize but OSAN cannot. An exemplary case is the one that yields to the policy $\pi$ considered for the family of non-isomorphic (n, $\ell$)-CSL graphs, allowing to sample all nodes from different CSL (sub-)graphs. 

More broadly, it is possible to show that Policy-Learn is more expressive than 1-WL, and not as powerful as 4-WL as it cannot distinguish strongly-regular graphs of the same parameters. We will discuss these aspects more thoroughly in the next version of our manuscript.

Reviewer comment (Theory):
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?

Author response:
**Q4**: The DS-GNN method used for selection policy relies on all subgraphs? What is the overall complexity of the method?

**A4**: The DS-GNN method used for the selection policy *does not* rely on all subgraphs. On the contrary, at each timestep $t$, it uses only the current bag, which contains $t$ subgraphs, to select the next subgraph to be added to the bag, starting from the original graph. This means that the number of subgraphs passed through the selection network is significantly smaller than the total number.

Nonetheless, we acknowledge the fact that including the computation complexity of the method would improve the quality of the manuscript. Thus, we present it below and we will include it in the next revision. 


We analyze an equivalent efficient implementation of our method, which differs from Algorithm 1 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 $\Delta_\text{max}$ the maximum node degree of the input graph, $n$ the number of its nodes and $T$ the number of selected subgraphs.

The forward-pass asymptotic time complexity of the selection network $f$ amounts to $\mathcal{O}(T \cdot n \cdot \Delta_\text{max})$. This arises from performing $T$ iterations, where each iteration involves selecting a new subgraph by passing the last subgraph added to the bag through the MPNN, which has time complexity $\mathcal{O}(n \cdot \Delta_\text{max})$ and re-using stored representations of previously-selected subgraphs. Similarly, the forward-pass asymptotic time complexity of the prediction network $g$ is $\mathcal{O}((T+1) \cdot n \cdot \Delta_\text{max})$, as each of the $T+1$ subgraphs is processed through the MPNN. Note that, differently from the selection network, the prediction network considers $T+1$ subgraphs as it also utilizes the last selected subgraph.

Therefore, Policy-Learn has an overall time complexity of $\mathcal{O}((T + (T+1)) \cdot n \cdot \Delta_\text{max})$, i.e., $\mathcal{O}(T \cdot n \cdot \Delta_\text{max})$. The time complexity of full-bag node-based Subgraph GNNs amounts instead to $\mathcal{O}(n^2 \cdot \Delta_\text{max})$, as there are $n$ subgraphs in total.
 
To grasp the impact of this reduction, consider the REDDIT-BINARY dataset, which we have experimented on to make this point clear (Tables 2 and 7). In this dataset the average number of nodes (and thus subgraphs) per graph is approximately $n = 429$, making the full-bag Subgraph GNN infeasible. In our experiment the number of selected subgraphs is $T=2$. The time complexity is drastically reduced as $T \ll n$, and, indeed, Policy-Learn can be effectively and successfully run.