A Flexible, Equivariant Framework for Subgraph GNNs via Graph Products and Graph Coarsening

We appreciate your positive feedback! We are pleased to hear that you find our method novel with solid theoretical foundations. We're also glad that you find the empirical results promising. Below, we address your specific comments and questions.

Q1: The size of the coarsened graph, though sparse, is still very large with $2^n$ nodes, which might result in high computation complexity for large input graphs. Do you have an estimate of the complexity and runtime of the proposed model compared with others?

A1: Let us clarify our approach. In our method, the size of the coarsened graph is always less than or equal to that of the original graph. To further emphasize, the space complexity of storing this sparse graph is upper bounded by the space complexity of storing the original graph $G$ (we do not store $2^n$ nodes). As shown in Figure 1 (left), the coarsened graph $\mathcal{T}(G)$ has fewer nodes than the original graph $G$. Considering the feature dimension to be constant, and given that $V$ is the number of nodes in the original graph, $T$ is the bag size, and $\Delta_{\text{max}}$ is the maximal degree of the original graph, our time complexity is upper-bounded by $\mathcal{O}(T \cdot V \cdot (\Delta_{\text{max}} + T))$. When considering $T=\mathcal{O}(1)$, which is typically the case, the complexity simplifies to $\mathcal{O}(V \cdot \Delta_{\text{max}})$. The space complexity is $\mathcal{O}(T \cdot E + V \cdot T^2)$, where $E$ is the number of edges in the original graph, which reduces to $\mathcal{O}(V +E)$, when $T = \mathcal{O}(1)$. This complexity is comparable to that of other methods in the field [1, 2]. Additionally, we refer the reader to Table 10 in Appendix F.5 (page 28) for a detailed comparison of the runtime of our method with other baselines. Our method's runtime is comparable to theirs.

Q2: Eq. (3) indicates that the input graph is unweighted and edge weights might not be considered. What if the input graph is weighted, for example, with edge weights 1,2,3, and 4.

A2: Thank you for highlighting this point. Our method is capable of handling weighted graphs. We employ a GINE base encoder [3], which can process edge features when applicable. For further details regarding our implementation, please refer to Appendix F (page 24) -- the update mechanism of the GINE base encoder, which supports edge features, is detailed in Equation (92) on page 24.

References:

[1] Bevilacqua et al. Efficient subgraph gnns by learning effective selection policies. ICLR 2023

[2] Kong et al. Mag-gnn: Reinforcement learning boosted graph neural network. NeurIPS 2024

[3] Hu et al. Strategies for pre-training graph neural networks. ICLR 2020

Thank you for your positive feedback. We are pleased you appreciated the novelty of our idea and the identification of the basis of equivariant layers for our product graphs. Below, we address each of your points.

Q1: Key definition

A1: We emphasize that the coarsened graph is not exponentially large in practice (see also lines 288-292 in our paper). While nodes of the coarsened graph are in $P([n])$, only a small number are used. Using spectral clustering, we partition the graph $G$ and construct the coarsened graph $\mathcal{T}(G)$ with nodes representing these partitions and edges induced by the original graph's connectivity (see Figure 1, left). This results in a coarsened graph with fewer nodes and edges than $G$. Although we define the graph for mathematical rigor, in practice, we use a much sparser graph by making $A^{\mathcal{T}}$ and $X^\mathcal{T}$ extremely sparse. The space complexity of this sparse graph is upper bounded by that of the original graph $G$ (we do not store $2^n$ nodes). We thank the reviewer for pointing out this confusion and will clarify this in the final version of the paper.

Q2: Theoretical analysis

A2: We have justified the marking strategy from a theoretical perspective, demonstrating its superiority over previously suggested strategies [1,2,3]. While we have only briefly summarized our results in the main body due to space constraints, App. C (p.16) provides a detailed expansion. We will make efforts to add more information to the main paper.

Q3: Presentation

A3: Thank you for highlighting these issues. Due to space limitations, Section 4.2.3 provides a brief summary of our use of the equivariant layers (see Section 4.2.2). For a detailed discussion, refer to App. F.4, p. 26). We will improve presentation and optimize space in the final version.

'More than a coarsening function’: We wanted to convey the fact that our method does not simply reduce to the joint processing of the graph and its coarsened counterpart. In fact, as we show in App. D.2 in detail, our approach strictly improves over running msg-passing on the graph obtained by combining its original topology with information from the coarsening. We term this new graph the “sum graph” and illustrate it on p. 19. Our method gains in expressiveness over the “sum-graph”, demonstrating the importance of our additional technical contributions, i.e., the use of the graph product and the related equivariant operations.

Q4: Connection with OSAN

A4: We agree with the reviewer and will add a detailed comparison with OSAN in the next revision.

In OSAN, the authors indeed show that ordered subgraphs are more expressive than unordered ones. In our case, although we choose coarsening functions which output unordered sets of nodes, we additionally leverage a set of additional equivariant operations which we formally derive as part of our contribution. Similarly to OSAN, it would be straightforward to show that our approach can indeed recover Subgraph GNNs operating on unordered subgraphs. However, the question whether OSAN fully captures our approach is of a less trivial answer, due to the aforementioned additional operations. Nevertheless, from an empirical perspective, our model demonstrates significantly better results.

Additionally, we note that unordered tuples are practically preferred in OSAN due to computational complexity reasons. Throughout their experimental section, the authors focus on unordered subgraphs and learn how to sample subsets of them. Our method shares this trait of (practically) considering only unordered subgraphs. Contrary to OSAN, however, we choose coarsening and graph products as an alternative to sampling.

We will refer to these observations in the next revision.

Q5: Experimental details

A5: We will answer each of the bullet points below:

• "Using effective subsets of subgraphs": we refer to how our method naturally selects and uses relevant subgraphs for optimal performance. To clarify, the initial step of our method involves coarsening the original graph $G$. We then use the graph Cartesian product between the coarsened graph $\mathcal{T}(G)$ and $G$ to create the product graph, as illustrated in Figure 1 on p. 2. Even without explicitly extracting subgraphs, elements in $\mathcal{T}(G) \square G$ can be conceptually associated with them. The "rows" in the product structure (Figure 1, right) represent implicit subgraphs, a common association in previous works [2,3]. We will improve our wording in the final version of the paper.
• To ensure a fair comparison, we compared with the best-performing version of OSAN, as shown in Table 1(a) of the OSAN paper [4], p. 9.
• Different papers use various datasets. We used common datasets where possible for comparison. Using different baselines for different datasets aligns with field practices due to the high computational cost of running all baselines on all datasets, as shown, for example, in Tables 2,3 in [2].

Q6: Maximally expressive

A6: We agree with the reviewer's observation, and indeed certain restrictions exist. To clarify, [1] studies the expressiveness of Subgraph GNNs using node-based policies. Thus, "maximal expressiveness" in our paper refers to Subgraph GNNs with node-based subgraph selection [1,3]. We will add an app. section to clarify this point.

Q7: Minor comments

A7:

• $\mathcal{X}$ is the node feature matrix of the graph $\mathcal{T}(G) \square G$.
• Thanks for the notes on L218, L225, and the figures. We'll revise them in the final version.

Refs.:

[1] Zhang et al. A Complete Expressiveness Hierarchy for Subgraph GNNs via Subgraph Weisfeiler-Lehman Tests. ICML 2023

[2] Bevi. et al. Equivariant Subgraph Aggregation Networks. ICLR 2022

[3] Frasca. et al. Understanding and Extending Subgraph GNNs by Rethinking Their Symmetries. NeurIPS 2022

[4] Qian et al. Ordered subgraph aggregation networks. NeurIPS 2022

Thank you for your detailed feedback. We appreciate your positive remarks on our method's novelty, strong mathematical formulation, and clear writing. Below, we address your specific questions.

Q1: The goal of our work and baselines

A1: Let us clarify. The goal of this work is to reduce the time and space complexity of Subgraph GNNs. As a result, some model degradation is expected. Subgraph GNNs are provably expressive models with theoretical guarantees [3, 4, 5], which may offer certain advantages over MPNNs and graph transformers. However, Subgraph GNNs typically have high computational complexity. Similarly to works like OSAN, PL, and Mag-GNN [6, 7, 8], our objective is to balance the high performance/accuracy of Subgraph GNNs with reduced computational expense. Unlike these other efficient Subgraph GNNs [6, 7, 8], our method offers a much finer control over this trade-off between computational complexity and accuracy in graph learning tasks, this is clearly illustrated in Figure 2, where bag size (x-axis) correlates with computational complexity. Within this family of more efficient Subgraph GNNs, we achieve SOTA results on most of the datasets these baselines are typically benchmarked on (see Tables 1, 4).

Q2: large scale datasets

A2: We refer the reviewer's attention to Table 8 in the appendix (p. 27), which presents results on the larger-scale ZINC-full dataset (250,000 graphs). We achieve top results in the efficient Subgraph GNN regime ($T=4$), outperforming MAG-GNN, the only other efficient subgraph GNN that reported results on this large-scale dataset. Furthermore, when using our model in the farthest end of the trade-off (Full bag - $T="Full"$), we achieve results comparable to the top-performing traditional "non-efficient" Subgraph GNNs, and outperform leading Graph Transformers. For instance, GRIT [9] reported an MAE of 0.023 on ZINC-FULL, while we obtained an MAE of 0.021.

As the reviewer requested, we also conducted a preliminary experiment using the large-scale molpcba dataset (437,929 graphs), which is also considered in the paper mentioned by the reviewer [2]. In particular, we train one instance of our model with a 1M parameter budget and report its performance against the methods in [2], see Table 1 in the attached PDF (in the global response).

Notably, our method achieves comparable results to SOTA on this dataset while using the least number of parameters. Compared to GatedGCN-LSPE, which uses a similar number of parameters, our method demonstrates better performance. This improvement is also evident when compared to the GraphTrans baseline, which utilizes approximately four times more parameters than our approach. We believe that the results over ZINC-FULL, as well as the preliminary results over molpcba, demonstrates the potential of our method to handle large-scale datasets well. If the reviewer considers it important, we will extend our experiments on molpcba and include the results in the final version of the paper.

Due to resource constraints, we couldn't conduct an experiment with the PCQM4Mv2 dataset (approx. 3,500,000 graphs) within the rebuttal period.

Q3: SOTA results on large scale datasets

A3: Recalling A1, A2, we achieve SOTA results on the large-scale ZINC-FULL dataset (Table 8, appendix, p. 27) in both efficient ($T=4$) and full-bag Subgraph GNN scenarios. Our approach outperforms top transformers like GRIT [9]. Preliminary experiments on the large-scale molpcba also indicate our model's competitive SOTA performance (Table 1 in the pdf).

Q4: Weak motivation

A4: We want to emphasize that our primary motivation, similar to previous works [6, 7, 8], is to reduce the computational cost of Subgraph GNNs. As mentioned in A1, we believe that within the family of algorithms that tackle the same setup (OSAN, PL, and MAG-GNN) [6, 7, 8], our method provides significant improvements in most cases – Table 1, 4, 8.

Q5: Could the experimental results improved if the proposed approach was combined with an approach similar to [2]. Improving the expressibility of the proposed model even further by combining with other classes of GNNs to get SOTA results would significantly improve the chances of the paper.

A5: We appreciate the reviewer's suggestion and have incorporated a version of the LCB block from [2] into our framework, which we term Ours + LCB. We ran experiments on three datasets (ZINC12k, Molbace, and Molesol) for three different bag sizes ($T=2$, $T=5$, $T="Full"$). The results are in Table 2 of the attached PDF in the global response.

We positively note that, in the small-bag regime, the Ours + LCB method improves the performance of our model (which already outperformed most competitors on these datasets), sometimes significantly. We will do our best to extend and include these results in the final version of the paper, should the reviewer find them relevant.

Summary: We believe we have thoroughly addressed your concerns about our method's capability with large-scale datasets and its integration with approaches from [2]. If you find our response satisfactory, we kindly request that you consider raising your score.

References:

[1] https://ogb.stanford.edu/docs/lsc/pcqm4mv2/

[2] https://openreview.net/pdf?id=duLr8BIzro

[3] Zhang et al. A Complete Expressiveness Hierarchy for Subgraph GNNs via Subgraph Weisfeiler-Lehman Tests. ICML 2023

[4] Bevi. et al. Equivariant Subgraph Aggregation Networks. ICLR 2022

[5] Frasca. et al. Understanding and Extending Subgraph GNNs by Rethinking Their Symmetries. NeurIPS 2022

[6] Qian et al. Ordered subgraph aggregation networks. NeurIPS 2022

[7] Bevi. et al. Efficient subgraph gnns by learning effective selection policies. ICLR 2023

[8] Kong et al. Mag-gnn: Reinforcement learning boosted graph neural network. NeurIPS 2024

[9] Ma et al. Graph Inductive Biases in Transformers without Message Passing. ICML 2023

Dear Reviewer hoHT,

We appreciate your thoughtful review and hope that we have adequately addressed your concerns. We made an effort to address your feedback during the rebuttal period. Specifically, we shared new results on new experiments:

• We conducted new experiments using a large-scale dataset OGBG-MOLPCBA (437K graphs). Additionally, we referred you to another large-scale experiment we already included in the original submission using ZINC FULL (250K graphs). Both experiments show solid results for our method compared to SoTA.

• We successfully incorporated components from [2] to enhance our architecture.

These new results can be found in our original response and PDF. We would be grateful if you could review these additions and consider raising your score if they satisfactorily address your concerns.

Best regards,

The Authors