Probabilistically Rewired Message-Passing Neural Networks

We thank the reviewer for their helpful comments and observations. We will respond to their raised issues in the following:

W1: More comparisons with deterministic rewiring methods.
RW1: Thank you for your suggestion. In our original manuscript, comparisons with deterministic rewiring methods such as DIGL and SDRF on heterophilic datasets were presented in Supplementary Table 11, with Drew on the QM9 dataset in Supplementary Table 10, and with DIGL and Drew on the Peptides datasets in Supplementary Table 12. Following the reviewer's recommendation, we have now included additional comparisons with FoSR and DropGNN on the TUDataset datasets in Table 4 and have expanded our analysis to cover all variants of SDRF and DIGL on the Heterophilic datasets in Supplementary Table 12.

W2: Over-squashing claim could be strengthened.
RW2: We thank the reviewer for their suggestion. Our empirical analysis of the widely-used Trees-NeighborsMatch dataset [1,2,3,4] demonstrates how PR-MPNNs alleviate over-squashing. As illustrated in Figure 4, PR-MPNNs outperform standard MPNNs and match Graph Transformers' performance, which is believed to not suffer from over-squashing issues [2].

We appreciate the suggestion and agree that further exploring the over-squashing phenomena with PR-MPNNs could be insightful. However, our experience indicates that traditional analyses using topological properties, as in references [3, 5], are not entirely adequate for black-box rewiring methods like ours. Despite logging various topological features such as edge curvatures and spectral gaps, no significant patterns emerged in the rewired graphs.

To illustrate why we believe that topological and spectral properties may be limited when employing learnable rewiring methods, we give an intuition via the following example:

image1.png

Suppose that, to solve some task, nodes A and B from the initial graph a) need to pass messages, but there’s an information bottleneck between C and D. Increasing curvature in the rewiring process (as in example b) might improve information flow. However, a trivial solution, that could be obtained with PR-MPNNs, is to connect nodes A and B, bypassing the bottleneck directly. Moreover, if we delete the edge connecting nodes C and D, we obtain a graph that has the same curvature as the initial graph, as exemplified in example c).

This simplified example may not fully capture the complexity encountered in real-world scenarios. However, we want to point out that PR-MPNNs are indeed able to find perfect solutions in this kind of toy scenarios, as illustrated in Figure 3, where the PR-MPNN is recovering the trivial solution to the task for the Trees-LeafCount synthetic dataset.

[1]: Alon, Uri, and Eran Yahav. "On the bottleneck of graph neural networks and its practical implications." arXiv preprint arXiv:2006.05205 (2020).
[2]: Müller, Luis, et al. "Attending to graph transformers." arXiv preprint arXiv:2302.04181 (2023).
[3]: Karhadkar, Kedar, Pradeep Kr Banerjee, and Guido Montúfar. "FoSR: First-order spectral rewiring for addressing oversquashing in GNNs." arXiv preprint arXiv:2210.11790 (2022).
[4]: Banerjee, Pradeep Kr, et al. "Oversquashing in GNNs through the lens of information contraction and graph expansion." 2022 58th Annual Allerton Conference on Communication, Control, and Computing (Allerton). IEEE, 2022.
[5]: Topping, Jake, et al. "Understanding over-squashing and bottlenecks on graphs via curvature." arXiv preprint arXiv:2111.14522 (2021).

Q1: Are new edges being added when rewiring?
RQ1: Thank you for highlighting this aspect. In our paper, while we focus on edge removal in Figure 5 concerning theoretical discussion in relation to randomized methods like DropEdge and DropGNN, in practical applications, we perform both edge addition and deletion. Detailed statistics on the number of edges added and removed are provided in Supplementary Table 8. To enhance clarity, we have revised our manuscript to explicitly state that we are both adding and removing edges.

Q2: More comparisons with deterministic rewiring methods.
RQ2: Please see our response to Weakness 1.

Q3: Evidence that we reduce over-squashing.
RQ3: As noted in our response to Question 1, our method involves both edge addition and deletion. For a detailed discussion on why spectral and topological analysis may not adequately address over-squashing in the context of learnable rewiring, please refer to our response to Weakness 2. We posit this is also the case for methods that treat graphs as fully-connected, such as Graph Transformers.

I would like to thank the authors for the reply. I think this is a nice piece of work and recommend acceptance to the conference.

We appreciate your kind reply. We would kindly ask you to adjust your score if we have addressed the concerns. Please also let us know what we can further improve.

W7: No computational efficiency in the main text.
RW7: We compute the computational complexity of PR-MPNNs in Section 3, noting that we obtain a better worst-case complexity than full self-attention-based graph transformers in the downstream module during test time, i.e., $\mathcal{O}(n^2)$ for a graph transformer versus $\mathcal{O}(L)$ for PR-MPNNs, where $n$ is the number of nodes in the graph, and $L$ is the number of edges in the rewired graph. Moreover, we also indicate that we can further reduce the complexity during training by restricting the edge candidates set, therefore achieving a lower overall complexity than a graph transformer. In the following table, we compare our runtime with a simple GINE model and the SAT Graph Transformer. PR-MPNN is approximately 5 times slower than a GINE model with a similar parameter count, while SAT is approximately 30 times slower than the GINE, and 6 times slower than the PR-MPNN models.

Furthermore, we have also revised Table 9 in the Appendix.

We thank the reviewer for their thorough feedback. In the following, we address the raised issues one by one.

W1: The motivation for k-subset sampling is not clear.
RW1: The choice of sampling $k$-subsets with $k>1$ is required for introducing complex dependencies between the edge random variables. On the other hand, even more complex constraints on these random variables would render the computation of marginals, as needed for the gradient estimation algorithm SIMPLE, intractable. Moreover, the $k$-subset sampling also enables fine-grained control of the sampled graphs’ sparsity. In the final paper, we will add a detailed discussion.

W2: The difference with prior work is not discussed.
RW2: We appreciate the reviewer's remark and have made improvements to the related work section, clarifying our method's difference from existing heuristic and probabilistic rewiring techniques. Our manuscript's "Related Work" section extensively contrasts PR-MPNNs with heuristic-based methods such as DIGL, SDRF [1], SP-MPNN [2], and Drew [3]. We emphasize the novelty of our learning-based rewiring approach and its robust theoretical link with WL expressivity. Additionally, and as noted by Reviewer j89x, we want to point to the extended section in the Appendix, which addressed the relation to Graph Structure Learning (GSL) methods, which bear conceptual similarities to our work but typically employ $k$-NN algorithms for sparsification or model edges with independent Bernoulli variables. Moreover, the GSL methods mainly focus on node classification tasks, while we’re more interested in graph-level data. Nevertheless, our theoretical PR-MPNN framework extends to various GSL methods, offering broader applicability.

[1]: Topping, Jake, et al. "Understanding over-squashing and bottlenecks on graphs via curvature." arXiv preprint arXiv:2111.14522 (2021).
[2]: Abboud, Ralph, Radoslav Dimitrov, and Ismail Ilkan Ceylan. "Shortest path networks for graph property prediction." Learning on Graphs Conference. PMLR, 2022.
[3]: Gutteridge, Benjamin, et al. "Drew: Dynamically rewired message passing with delay." International Conference on Machine Learning. PMLR, 2023.

W3: How does PR-MPNN help alleviate over-squashing?
RW3: Our empirical analysis on the widely-used Trees-NeighborsMatch dataset [1,2,3,4] demonstrates how PR-MPNNs alleviate over-squashing. As illustrated in Figure 4, PR-MPNNs outperform standard MPNNs and match Graph Transformers' performance, which is believed to suffer less from over-squashing issues [2]. Using these benchmark datasets is the de facto community standard for assessing the over-squashing behavior of MPNNs.

[1]: Alon, Uri, and Eran Yahav. "On the bottleneck of graph neural networks and its practical implications." arXiv preprint arXiv:2006.05205 (2020).
[2]: Müller, Luis, et al. "Attending to graph transformers." arXiv preprint arXiv:2302.04181 (2023).
[3]: Karhadkar, Kedar, Pradeep Kr Banerjee, and Guido Montúfar. "FoSR: First-order spectral rewiring for addressing oversquashing in GNNs." arXiv preprint arXiv:2210.11790 (2022).
[4]: Banerjee, Pradeep Kr, et al. "Oversquashing in GNNs through the lens of information contraction and graph expansion." 2022 58th Annual Allerton Conference on Communication, Control, and Computing (Allerton). IEEE, 2022.

W4: No experiments to validate claims about robustness.
RW4: We thank the reviewer for this observation. A beneficial side-effect of training PR-MPNNs is the enhanced robustness of the downstream model to graph structure perturbations. This is because PR-MPNNs generate multiple adjacency matrices for the same graph during training, akin to augmenting training data by randomly dropping or adding edges, but with a parametrized "drop/add" distribution rather than a uniform one.

Motivated by the reviewer’s feedback, we conducted a dedicated experiment on the PROTEINS dataset. After training our models on clean data, we compared the models’ test accuracy on the clean and corrupted graphs. Corrupted graphs were generated by either deleting or adding a certain percentage of their edges. We report the change in the average test accuracy over 5 runs, comparing the base model with variants of PR-MPNN in the following table:

PR-MPNNs consistently obtain the best results for both removing k=25 and k=50 edges. The three models achieve test accuracies on the unperturbed dataset as follows: GINE 77.9%, PR-MPNN-RM25 80.1%, and PR-MPNN-RM50 78.8%.
We included these initial results in the revised manuscript; see Table 13 in the appendix.

W5: SIMPLE does not have a detailed explanation.
RW5: We thank the reviewer for their remark. We have added an explanation of SIMPLE to the appendix of the updated submission.

W6: Concerns about experiments.
(1) Confusion about how we leverage the adjacency matrices.
R: We thank the reviewer for the remark. To clarify, we always use a downstream ensemble to obtain the results when using multiple rewired adjacency matrices. We have updated our manuscript with a clarification.

(2) a) Different baselines for different datasets.
R: We selected different baselines for each dataset based primarily on the availability of the authors’ code and published results. Our evaluation of the PR-MPNN method covers 20 datasets, detailed in Table 5 of the Appendix. Due to time and resource limitations, it wasn't feasible to comprehensively evaluate all baselines across these datasets. We relied on publicly available results for comparisons where the direct evaluation was impractical. However, we conducted additional experiments for some baselines on datasets that were not initially reported (e.g., SAT on OGBG-Molhiv, GIN, and GIN+PE in Tables 1, 3, 10, and 11). Furthermore, we have introduced two new baselines on the TUDataset and have run new experiments with FoSR (PTC_MR, NCI1, NCI109) and DropGNN (NCI1, NCI109).
(2) b) Comparable performance with DropGNN on 4-Cycles.
R: We appreciate the reviewer pointing this out. We intended to compare PR-MPNNs with DropGNN on 4-Cycles to emphasize the significance of utilizing multiple priors in sampling. This key aspect is detailed in the caption of Figure 2, where we illustrate the impact of our sampling approach. We have updated the text in the revised manuscript to reflect this.

(3) Table 1 is confusing.
R: We appreciate the reviewer's feedback on Table 1. The evaluation metric for the ZINC and Alchemy datasets is the Mean Absolute Error (MAE), with a lower score indicating better performance. For the OGBG-Molhiv dataset, the metric is ROCAUC, where a higher score denotes superior performance. In our initial submission, we used symbols $\uparrow$ and $\downarrow$ to signify this. In our revised version, we explicitly state the evaluation metric for each dataset in Supplementary Table 5 to improve clarity.

(4) There are no ablation studies.
R: We have conducted ablation studies focusing on the performance impact of different gradient estimators, specifically Gumbel, I-MLE, and SIMPLE, as detailed in Table 1. Additionally, our comparison with DropGNN, presented in Figure 2, demonstrates the benefits of utilizing multiple priors and samples. Furthermore, we have revised the computation efficiency in Table 9 in the Appendix, where we compare the runtime of various PR-MPNN models when using different gradient estimators and the total number of edges sampled. These sections collectively serve as our approach's ablation analysis.

We thank the reviewer for their helpful comments and the positive evaluation of our work.

W1: The computation overhead is also needed. RW1: We compute the computational complexity of PR-MPNNs in Section 3, noting that we obtain a better worst-case complexity than full self-attention-based graph transformers in the downstream module during test time, i.e., $\mathcal{O}(n^2)$ for a graph transformer versus $\mathcal{O}(L)$ for PR-MPNNs, where $n$ is the number of nodes in the graph, and $L$ is the number of edges in the rewired graph. Moreover, we also indicate that we can further reduce the complexity during training by restricting the edge candidates set, therefore achieving a lower overall complexity than a graph transformer. In the following table, we compare our runtime with a simple GINE model and the SAT Graph Transformer. PR-MPNN is approximately 5 times slower than a GINE model with a similar parameter count, while SAT is approximately 30 times slower than the GINE, and 6 times slower than the PR-MPNN models.

Furthermore, we have also revised Table 9 in the Appendix.

W2: How does the expressiveness translate to practical performance?
RW2: Our empirical results validate the expressivity claims using the EXP and CSL datasets. These datasets are designed so GNNs limited to 1-WL expressivity will not perform well. As detailed in Tables 2 and 3, our rewiring technique achieves perfect performance on EXP and nearly perfect on CSL, significantly outperforming the base GIN model's random performance, clearly indicating that our expressivity guarantees translate into practice.

In addition, as shown in Table 1, our SIMPLE-based rewiring method consistently surpasses the base model regarding predictive performance on real-world datasets. Further, Table 10 highlights a substantial improvement, up to 14.13x, over the base GIN model on QM9, demonstrating significant real-world performance gains.