Paper: Probabilistically Rewired Message-Passing Neural Networks (ICLR-2024)

Reviewer comment (Experiments):
I believe experimentally there could be a larger breadth of rewiring benchmarks, for instance comparing against "deterministic" rewiring techniques such as FOSR [1] and SDRF [2] would be valuable. Furthermore, adding a "random rewiring" GNN, i.e. DropGNN, would also be useful for the real-world tasks.

Author response:
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.

Reviewer comment (Reproducibility):
From Figure 1 and the paper overall, it seems like the rewiring technique is only removing edges and not adding any new edges not present in the original graph. Is this the case?

Author response:
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.

Reviewer comment (Experiments):
Regarding (W1), would it be possible to show results for existing deterministic rewiring techniques?

Author response:
Please see our response to Weakness 1.

Reviewer comment (Theory):
The motivation for using $k$-subset constraint when sampling the adjacency matrix is not very clear. What are the advantages of using such constraints?

Author response:
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.

Reviewer comment (Novelty):
The difference between the proposed PR-MPNN and previous works is not explicitly discussed in the related work section.

Author response:
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.

Reviewer comment (Theory):
The theoretical results indicate that PR-MPNN is more effective in probabilistically separating graphs than randomized approaches under certain conditions. But how does this help PR-MPNN address the over-squashing problem?

Author response:
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.

Reviewer comment (Experiments):
In the introduction section, the paper states that PR-MPNNs make MPNNs less vulnerable to potential noise and missing information. However, there are no empirical or theoretical results to validate such statements.

Author response:
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:

| Corruption\\Model | GINE (acc 77.9) | PR-MPNN-RM25 (acc 80.1) | PR-MPNN-RM50 (acc 78.8) |
| ----------------- | --------------- | ----------------------- | ------------------------ |
| RM 10%            | \-7.68%         | +0.04%                  | __+0.34%__                   |
| RM 30%            | \-11.87%        | \-2.56%                 | __\-0.56%__                  |
| RM 50%            | \-9.18%         | \-3.50%                 | __\-0.21%__               |
| ADD 10%           | \-5.93%         | +0.28%                  | __+0.80%__                   |
| ADD 30%           | \-11.14%        | \-1.27%                 | __+0.28%__                   |
| ADD 50%           | \-21.75%        | \-1.99%                 | __\-0.43%__                  |

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.

Reviewer comment (Writing):
PR-MPNN uses SIMPLE to sample adjacency matrices. It seems to be an important component of the proposed model and the paper should provide an introduction to the SIMPLE method to make the paper self-contained.

Author response:
We thank the reviewer for their remark. We have added an explanation of SIMPLE to the appendix of the updated submission.

Reviewer comment (Reproducibility):
In the “Baseline and model configurations” paragraph, the paper states that there are two ways to leverage the sampled adjacency matrices when using multiple priors. However, it is unclear which method is used in their experiments.

Author response:
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.

Reviewer comment (Evaluation):
When answering Q2, the paper compares PR-MPNN with different baselines on different datasets. For example, the paper compares PR-MPNN with OSAN, GIN+POSENC, and DropGNN on EXP, CSL, and 4-CYCLES datasets respectively. Also, the statement “Concerning Q2, on the 4-CYCLES dataset, our probabilistic rewiring method consistently outperforms DropGNN” is inaccurate since PR-MPNN only achieves comparable performance with DropGNN in some cases.

Author response:
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).  
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.

Reviewer comment (Evaluation):
The results in Table 1 are quite confusing. What evaluation metric is used in these results? Why is it that on some datasets, such as OGBG-MOLHIV, a higher metric indicates better performance, while on others, such as ZINC and ALCHEMY, a lower value is preferable?

Author response:
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.

Reviewer comment (Experiments):
There are no ablation studies to validate the effectiveness of each component of the proposed model.

Author response:
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.

Reviewer comment (Experiments):
In the conclusion section, the paper states that PR-MPNN “is competitive or superior to conventional MPNN models and graph transformer architectures regarding predictive performance and computational efficiency”. However, there is no comparison between PR-MPNN and baselines regarding the computational efficiency in the main text.

Author response:
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. 

| Model           | #Params | Total sampled edges | Train time/ep (s) | Val time/ep (s) |
|-----------------|---------|---------------------|-------------------|-----------------|
| GINE            | 502k    | -                   | 3.19 ± 0.03       | 0.20 ± 0.01     |
| K-ST SAT (GINE) | 506k    | -                   | 86.54 ± 0.13      | 4.78 ± 0.01     |
| K-SG SAT (GINE) | 481k    | -                   | 97.94 ± 0.31      | 5.57 ± 0.01     |
| K-ST SAT (PNA)  | 534k    | -                   | 90.34 ± 0.29      | 4.85 ± 0.01     |
| K-SG SAT (PNA)  | 509k    | -                   | 118.75 ± 0.50     | 5.84 ± 0.04     |
| PR-MPNN (Gmb)   | 582k    | 20                  | 15.20 ± 0.08      | 1.01 ± 0.01     |
| PR-MPNN (Gmb)   | 582k    | 100                 | 18.18 ± 0.08      | 1.08 ± 0.01     |
| PR-MPNN (IMLE)  | 582k    | 20                  | 15.01 ± 0.22      | 1.08 ± 0.06     |
| PR-MPNN (IMLE)  | 582k    | 100                 | 15.13 ± 0.17      | 1.13 ± 0.06     |
| PR-MPNN (Sim)   | 582k    | 20                  | 15.98 ± 0.13      | 1.07 ± 0.01     |
| PR-MPNN (Sim)   | 582k    | 100                 | 17.23 ± 0.15      | 1.15 ± 0.01     |

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

Reviewer comment (Evaluation):
The computation overhead is also needed.

Author response:
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. 

| Model           | #Params | Total sampled edges | Train time/ep (s) | Val time/ep (s) |
|-----------------|---------|---------------------|-------------------|-----------------|
| GINE            | 502k    | -                   | 3.19 ± 0.03       | 0.20 ± 0.01     |
| K-ST SAT (GINE) | 506k    | -                   | 86.54 ± 0.13      | 4.78 ± 0.01     |
| K-SG SAT (GINE) | 481k    | -                   | 97.94 ± 0.31      | 5.57 ± 0.01     |
| K-ST SAT (PNA)  | 534k    | -                   | 90.34 ± 0.29      | 4.85 ± 0.01     |
| K-SG SAT (PNA)  | 509k    | -                   | 118.75 ± 0.50     | 5.84 ± 0.04     |
| PR-MPNN (Gmb)   | 582k    | 20                  | 15.20 ± 0.08      | 1.01 ± 0.01     |
| PR-MPNN (Gmb)   | 582k    | 100                 | 18.18 ± 0.08      | 1.08 ± 0.01     |
| PR-MPNN (IMLE)  | 582k    | 20                  | 15.01 ± 0.22      | 1.08 ± 0.06     |
| PR-MPNN (IMLE)  | 582k    | 100                 | 15.13 ± 0.17      | 1.13 ± 0.06     |
| PR-MPNN (Sim)   | 582k    | 20                  | 15.98 ± 0.13      | 1.07 ± 0.01     |
| PR-MPNN (Sim)   | 582k    | 100                 | 17.23 ± 0.15      | 1.15 ± 0.01     |

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

Reviewer comment (Theory):
How does the expressiveness guarantee translate to practical performances?

Author response:
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.