Adaptive Message Passing: A General Framework to Mitigate Oversmoothing, Oversquashing, and Underreaching

We thank the reviewer for recognizing that the paper is well written and the findings are convincing. We will comment below on the questions raised.

Other Strengths And Weaknesses:

1. While oversmoothing, oversquashing, and underreaching are topics of great interest in the community, which are far from being fully understood or solved, our paper’s unique motivation is that, by combining the ability to learn the depth with the adaptive filtering of messages, we want to determine i) the number of required layers and ii) which messages should be propagated. In other words, our contribution is unique in that it determines, during training, how much and when to send information. We hope to have clarified this point.
2. In Appendix D, we provide a critical view on how the oversquashing term is defined and understood by our community. In particular, we argue that if oversquashing refers to “informational bottlenecks”, then additive rewiring may hurt performances if the same amount of layers is kept; in contrast, oversquashing defined as sensitivity is always improved by additive rewiring. For this reason, we believe that AMP and rewiring-based approaches are different but equally valid approaches to achieve long-range propagation. We will clarify these aspects in the paper, and we hope the reviewer appreciates our considerations. Thanks to the suggestion of Reviewer ZB4Z, we now have a new Theorem inspired by Di Giovanni et al, 2023 that exposes this inconsistency in the literature!
3. In terms of impact of the depth as regards efficiency, we refer to the computational complexity analysis written in Reviewer XraS’s response. We will add these considerations to the paper. Overall, the complexity of using a deeper network increases due to i) the higher parametrization, as usual; ii) the need for a layer-wise readout in AMP, as also done in previous works [JK-Net, Xu et al., 2018].

In addition, we prepared an ablation study where we try different (fixed) depths for AMP, maintaining the other architectural changes intact. In particular, we allow AMP to learn a normalized importance over the fixed number of layers. Here, we want to see the impact of dynamically learning the depth on effectiveness. The depths we tried depend on the range shown in Figure 4 (so, up to 45 different depths for a single model and dataset), while the other hyper-parameters were fixed to the best one we found for the task. The table is shown below for the peptides datasets:

It appears that learning the importance of fixed-depth network layers does not allow to obtain better performances than the fully adaptive AMP, and it potentially requires many more configurations to be tried, which is the main disadvantage of considering the depth as a hyper-parameter rather than a learnable parameter (as we do).

Questions For Authors: The specific tasks are chosen because there is reason to believe long-range information plays an important role. In other words, all our tasks require the ability to “reason” globally about the graph and have specifically tested for these reasons. In addition, we used the molecular tasks of the LRGB paper because i) it is possible to perform a fair evaluation on these tasks, thanks to the robust re-evaluation of Tönshoff et al.. IPR-MPNN evaluation on additional benchmarks is necessary to empirically validate their increased expressive power compared to 1-WL MPNNs, whereas we believe to already have sufficient empirical evidence that AMP can be beneficial to improve the performance of MPNNs on long-range tasks; AMP always improves the results compared to any of its base versions on all datasets.

An important clarification about AMP’s competitiveness: our goal is not to demonstrate that AMP surpasses the state of the art, as AMP is not a model but a framework that can be wrapped around most MPNNs. In our paper, we used simple MPNNs for their ease of use, showing the performance gains that AMP grants without changing a single line of code of these convolutional layers. In our view, the correct metric to assess AMP’s effectiveness is the gap in performance compared to the base versions AMP is wrapped around, showing that it helps to better capture long-range dependencies.

Finally, thank you for the reference, we will include it in the revised manuscript. We were unable to find such models for the considered datasets, but we should definitely mention architectures specific to molecular tasks in our work, as it is relevant due to the electrostatics interactions necessary to model energy and forces.

---

Conclusion: We will report these clarifications and additional statements in the paper. We hope the reviewer appreciates our response and hopefully will increase the score.

We thank the reviewer for the comments. Below, we clarify some of the points raised.

Claims And Evidence: We apologize for the incorrect statement in the abstract about “surpassing” the state of the art. That statement was true in the past, before we revised Table 2. Indeed, we had fixed our claims elsewhere in the paper (lines 80-83, 242-244, 377-379) by talking about “competitive” performances.

However, the goal of the analysis is not to show absolute improvements compared to recent works. Rather, it is to show, as the reviewer correctly identified, that wrapping the AMP framework around basic MPNNs grants a performance improvement in capturing long-range dependencies, due to the particular characteristics of AMP that endow MPNNs to mitigate oversmoothing and oversquashing as measured in Figure 3, in addition to the adaptive learning of the depth to mitigate underreaching. It is our belief that our work should be judged in terms of these contributions, and not solely in terms of absolute numbers, since we, for the ease of prototyping, chose to wrap AMP around classical MPNNs, rather than more recent and convoluted ones.

In terms of mitigating under-reaching, it is possible that the network finds a degenerate solution during training that corresponds to an unsatisfactory depth, but this is similar to how MLPs can end up in poor local minima when trying to learn functions. The design of AMP enables learning the depth and we see that it works empirically.

Finally, in our results, it seems the method learns more layers than those tried in previous works using grid search.

Experimental Designs Or Analyses: Thanks for pointing this out, we partly referred to the hyper-parameter ranges and data splits of Gravina et al. and Tonshoff et al. without explicitly mentioning them in the paper. To improve self-containment and reproducibility, we will add this information to Appendix E. The search methodology was a standard grid search with early stopping on validation performance.

Relation To Broader Scientific Literature: We would like to clarify that our approach does not pertain to probabilistic graph rewiring approaches; it is probabilistic, but it does not perform a topological rewiring of the graph. The graph topology remains the same, and messages are partially filtered to mitigate the oversmoothing and oversquashing issues. Concretely, we never add new edges to the input graph: please also refer to our answer to question 1 below.

Essential References Not Discussed: Thank you for providing these additional references, we agree with the reviewer’s analysis and we will include and discuss them in the revised manuscript. Briefly, in terms of main differences with AMP, [3] relies on a separate sequence model to develop a node selection mechanism whereas we work on the message passing itself; [4] develops a new graph Laplacian that is better suited for long-range propagation; [5] converts a graph into a latent representation that can be passed to downstream classifiers.

Questions For Authors:

1. As we also clarified earlier, our approach differs significantly from rewiring approaches, including probabilistic ones. In AMP, the probabilistic formulation is used to dynamically learn the depth (following Nazaret and Blei), something which rewiring approaches cannot do, whereas the “filtering of messages” is not to be intended as a variation to the original graph topology, rather as a “node-based” filtering of outgoing messages. The major difference is that AMP alters the computational graph but not the graph topology, whereas typical graph rewiring methods alter the graph topology. In Appendix D, we have discussed this through the lens of oversquashing.
2. Modeling a one-dimensional distribution leads to no particular overhead during training. We point the reviewer to our answer to Reviewer XraS for a discussion on computational complexity. In short, the main complexity is caused by the addition of one readout per layer, something which was also done in previous works such as JK-Net (Xu et al, 2018). In terms of training, optimizing the ELBO reduces to performing backpropagation w.r.t. a standard prediction loss, plus one (optional) weight regularization term and another (optional) depth regularization term. The overall asymptotic complexity is therefore not different from other GNNs that appeared in the literature in previous years.

---

Conclusion: We hope to have clarified the goal of our contribution and to have addressed the other questions. We would appreciate it if the reviewer could consider increasing the score in light of these considerations and revisions. We remain available should further information be needed.

We sincerely thank the reviewer for the detailed review and the valuable suggestions. We will do our best to clarify doubts and address the points of the reviewer.

Claims And Evidence:

On Figure 3: We will improve the coloring scheme as suggested, thank you. Also, the reason why we do not show GCN results for more layers is because we report the best configuration on the validation set, but we can mention in the paper that it has been already proven how GCNs oversmooth in the long run. Finally, the Dirichlet Energy (DE) becomes negative simply because the y-axis is in log scale to aid readability, as DE values and sensitivities vary a lot.

Theoretical Analysis: We cannot thank the reviewer enough for the suggestion! We indeed extended Theorem 3.2 of Di Giovanni ‘23 to our message filtering scheme, and (in short) we arrived at the conclusion that, if $c_F$ is the Lipschitz constant of our filtering function, $k_F \le 1$ is the maximal element-wise value attained by the Filtering function $F$, and $k_h$ is the maximal element-wise value of a node embedding, we arrive at a similar result for (using the notation of Theorem 3.2)

$\boldsymbol{S}_{r,a} := c_r\boldsymbol{I} + c_a(c_F*k_h + k_F)\boldsymbol{A}$

Where $\boldsymbol{S}_{r,a}$ controls the upper bound on sensitivity due to the topology of the graph. We are very happy with the Reviewer’s suggestion because this result formally supports what we argued in Appendix D: if we consider for simplicity a constant filtering function, namely $c_F=0$ and we filter enough, meaning $k_F < 1$, then filtering will decrease the sensitivity’s upper bound. However, this is actually helping to reduce the “informational oversquashing” defined in Alon et al., contradicting the widely used statements that “improving sensitivity mitigates oversquashing”. As a result, we argue that the community should probably distinguish the “informational oversquashing” of Alon et al. from the “topological oversquashing” of Topping et al.. We will add these results to our paper. Thank you so much for the invaluable suggestion! This will help our community reason more clearly about these issues.

Methods And Evaluation Criteria: Following the reviewer’s suggestion, we perform additional experiments with GCN and AMP_GCN on homophilic and heterophilic datasets. Risk assessment is performed using 10-fold Cross Validation, with an internal hold-out hyper-parameter tuning for both models. For each of the 10 folds, the best configuration on that fold is re-trained 10 times and the test values averaged. The final test score is the average of the test scores across the 10 folds. This procedure is inspired from [*] and results are shown below.

It appears that AMP allows GCN to improve results especially on heterophilic datasets. Therefore results seem to align with the reviewer’s intuition. We will add other common datasets and baselines in the revised version of the paper. Thank you for suggesting these extra experiments.

[*] Errica F., Podda M., Bacciu D., Micheli M., A fair comparison of graph neural networks on graph classification. ICLR 2020

Essential References Not Discussed: Thanks for pointing out that some references in the Table are not cited, we will amend that.

Remaining Sections: We hope to have addressed the weaknesses highlighted by the reviewer, namely the theoretical analysis and the extra experiments on node classification tasks on homo/heterophilic datasets. We will provide a complete proof of the extended version of Theorem 3.2 for AMP in the revised paper.

Questions For Authors: As a matter of fact, the entire architecture is trained using conventional backprop on a prediction loss + some regularization terms. There is no MLP involved in the distribution of the layers’ importance because we simply need to learn its parameters, but note that the variational formulation merely serves to show that AMP arises from a well defined graphical model, grounding our design choices in a principled formulation. The practical implementation is very simple, thanks to the brilliant work of Nazaret and Blei. Please let us know if we can further clarify this point.

---

Conclusion: Thank you again for the constructive comments, especially those related to the theoretical analysis. We hope the Reviewer appreciates our efforts and that this will lead to a score increase.

We thank the reviewer for recognizing the merits of our contribution and for providing constructive criticism. Below we comment on some of the points raised by the reviewer.

Methods And Evaluation Criteria: Following Reviewer ZB4Z’s suggestion, we will include node classification tasks related to homophily and heterophily and increase the number of real-world datasets considered. Please see our response to Reviewer ZB4Z for a first batch of results.

Theoretical Claims: AMP focuses on long-range interactions problems, where far-away nodes may need to communicate effectively (for instance, without oversquashing arising). Theorem 3.1 shows that, for a proper parametrization, it is possible to realize such long-range communication between two specific nodes without oversquashing behavior, which would not be possible on classical (synchronous) message-passing neural networks. This is meant to support the design choices made for AMP, although we do not want to claim that AMP is always able to implement such behavior in practice (we do not observe it empirically). At the same time, such a behavior is reminiscent of the asynchronous updates of [FW]. We will revise the paper to reflect these considerations, and we hope that this clarifies the reviewer’s doubts.

Other Strengths And Weaknesses: Thank you for the kind words about our work. We provide an answer for each potential improvement below.

• It is true that we did not mention additional costs in detail, thank you for pointing this out. The cost of filtering messages is $O(N)$, with $N$ nodes (lines 184-185). Therefore, the message passing operation is not altered significantly, since it has a cost of $O(N+E)$, with $E$ edges. However, the additional burden introduced by AMP, compared to classical MPNNs, is the layer-wise readout (lines 344-345) that we implemented as an MLP. Classical MPNNs employ a single readout with cost $O(N)$ or $O(1)$ depending on the task nature, whereas we use one per layer, so we have $O(NL)$ and $O(L)$ respectively. Note that other popular MPNN architectures, such as JK-Net (Xu et al., 2018), employ a similar scheme by means of concatenating the node embeddings across all layers. In terms of training costs, standard backpropagation with at most two light-weight additional regularizers is employed. We will make this clear in the revised version of the paper.
• Please refer to our additional experiments in Reviewer ZB4Z’s response, as mentioned before.

---

Conclusion: Thank you once more for the constructive feedback! We hope our clarifications and additional experiments may constitute ground for a score improvement. We remain available for further discussions.