Probabilistically Rewired Message-Passing Neural Networks
Authors: Chendi Qian
Andrei Manolache
Kareem Ahmed
Zhe Zeng
Guy Van den Broeck
Mathias Niepert
Christopher Morris
Venue: ICLR 2024 poster (2024)
Keywords: GNN
MPNN
graph
rewiring
probabilistic
weisfeiler
leman
lehman
k-subset
sampling
expressivity
Primary area: learning on graphs and other geometries & topologies

Abstract: Message-passing graph neural networks (MPNNs) emerged as powerful tools for processing graph-structured input. However, they operate on a fixed input graph structure, ignoring potential noise and missing information. Furthermore, their local aggregation mechanism can lead to problems such as over-squashing and limited expressive power in capturing relevant graph structures. Existing solutions to these challenges have primarily relied on heuristic methods, often disregarding the underlying data distribution. Hence, devising principled approaches for learning to infer graph structures relevant to the given prediction task remains an open challenge. In this work, leveraging recent progress in exact and differentiable k-subset sampling, we devise probabilistically rewired MPNNs (PR-MPNNs), which learn to add relevant edges while omitting less beneficial ones. For the first time, our theoretical analysis explores how PR-MPNNs enhance expressive power, and we identify precise conditions under which they outperform purely randomized approaches. Empirically, we demonstrate that our approach effectively mitigates issues like over-squashing and under-reaching. In addition, on established real-world datasets, our method exhibits competitive or superior predictive performance compared to traditional MPNN models and recent graph transformer architectures.

Published as a conference paper at ICLR 2024
PROBABILISTICALLY REWIRED MESSAGE-PASSING NEURAL NETWORKS
Chendi Qian"
Andrei Manolache'
Computer Science Department RWTH Aachen University, Germany chendi, qí an www.h-acchee de
Computer Science Department
University of Stuttgart, Germany
Bitdefender, Romania andrei marreaimmollcceek.... wag.aollacccee..ummssttttttrrr www.unnutttaar.....mm
Kareem Ahmed, Zhe Zeng & Guy Van den Broeck Computer Science Department University of California, Los Angeles, USA
Mathias Niepert!
Christopher Morris
Computer Science Department University of Stuttgart, Germany
Computer Science Department RWTH Aachen University, Germany
ABSTRACT
Message-passing graph neural networks (MPNNs) emerged as powerful tools for processing graph-structured input. However, they operate fixed input graph structure, ignoring potential noise and missing information. Furthermore, their local aggregation mechanism can lead problems such over-squashing and limited expressive power in capturing relevant graph structures. Existing solutions to these challenges have primarily relied heuristic methods, often disregarding the underlying data distribution. Hence, devising principled approaches for learning to infer graph structures relevant to the given prediction task remains an open challenge. In this work, leveraging recent progress in exact and differentiable k-subset sampling, we devise probabilisticall rewired MPNNS (PR-M,NNs), which learn add relevant edges while omitting less beneficial ones. For the first time, our theoretical analysis explores how PR-MPNNs enhance expressive power, and we identify precise conditions under which they outperform purely randomized approaches. Empirically, we demonstrate that our approach effectively mitigates issues like over-squashing and under-reaching. In addition, on established real-world datasets, method exhibits competitive superior predictive performance compared to traditional MPNN models and recent graph transformer architectures.
INTRODUCTION
Graph-strue structured data is prevalent across various application domains, including fields like chemo- and bioinformatics (Barabasi & Oltvai 2004 Jumper et al., 2021 Reiser et al., 2022), combinatorial
optimization (Cappart et al., 2023), and netwo-n analysis (Easley et al., 2012). highlighting the
need for machine learning techniques designed explicitly for graphs. In recent years, message-passing
graph neural networks (MPNNs) (Kipf Welling 20177 Gilmer et al. 2017 Scarselli et al.) 2008b)
Veličković et al. 2018) have become the dominant approach in this area, showing promising performance
in tasks such as predicting molecular properties (Klicpera et al. 2020 Jumper et al. 2021 or enhancing combinatorial solvers (Cappart et al. 2023).
However, MPNNs have a limitation due to their local aggregation mechanism. They focus on encoding
local structures, severely limiting their expressive power (Morris al., 2019 2021 Xu et al. 2019). In addition, MPNNs struggle to capture global or long-range information, possibly leading to phenomena
'These authors contributed equally.
"Co-senior authorship.

---

Published as a conference paper at ICLR 2024
Forward pass
(G)
A (A(A)
À 9 (AGG), A ....,.(.))))
Upstreum MPNN
Exactly-& Sampling ling
A
Downstream MPNN
< (A.
A(G) x
Aggregation
A x
0
Ag 5.6 0 4.1
Backward pass
VaL x (o) with p(() (PIAR) (AG)))
Figure I: Overview of the probabilistically rewired MPNN framework. PR-MPNNS use an upstream model to learn priors 0 for candidate edges, parameterizing a probability mass function conditioned on exactly-k constraints. Subsequently, we sample multiple k-edge adjacency matrices (here: A = 1) from this distribution, aggregate these matrices (here: subtraction), and use the resulting adjacency matrix as input to a downstream model, typically an MPNN, for the final predictions task. On the backward pass, the gradients of the loss l regarding the parameters are approximated through the derivative of the exactly-k marginals in the direction of the gradients of the point-wise loss e regarding the sampled adjacency matrix. We use recent work to make the computation of these marginals exact and differentiable, reducing both bias and variance.
like under-reaching (Barceló et al. 2020) or over-squashing (Alon & Yahav 2021). Over-squashing, as explained by |Alon & Yahay 2021). refers to excessive information compression from distant nodes due to source node's s extensive receptive field, occurring when too many layers are stacked.
Topping et al. (2021): ); Bober et al. (2022) investigated over-squashing from the perspective of Ricci and Forman curvature. Refining Topping et al. (2021), Di Giovanni et al. (2023) analyzed how the architectures' width and graph structure contribute to the over-squashing problem, showing that over- squashing happens among nodes with high commute time, stressing the importance of graph rewiring techniques, i.e., adding edges between distant nodes to make the exchange of information more accessible. In addition, [Deac et al. (2022); (Shirzad et al. (2023) utilized expander graphs enhance message passing and connectivity, while Karhadkar et al. (2022) resort to spectral techniques, and [anerjee et al. (2022) proposed greedy random edge flip approach overcome over-squashing. Recent work (Gutteridge et al., 2023) aims to alleviate over-squashing by again resorting to graph rewiring. In addition, many studies have suggested different versions of multi-hop hoalii-hop-niighbor based message passing to maintain long-range dependencies (Abboud et al. 2022, Abu-El-Haija et al. 2019 Gasteiger et al.. 2019 Xue al. 2023), which can also be interpreted as a heuristic rewiring scheme. The above works indicate that graph rewiring is an effective strategy to mitigate over-squashing However, most existing graph rewiring approaches rely on heuristic methods to add edges, potentially not adapting well to the specific data distribution or introducing edges randomly. Furthermore, there is limited understanding to what extent probabilistic rewiring, i.e., adding or removing edges based on the prediction task, impacts the expressive power of model. In contrast to the above lines of work, graph transformers (Chen et al. 2022 Dwivedi
et al. 2022b He et al. 2023, 3,,,3 Müller et al. 2023 Rampašek et al. 2022) and similar global attention mechanisms (Liu al. 202 Wu et al.) 2021) marked shift from local to global message passing, aggregating over all nodes. While not understood in principled way, empirical studies indicate that graph transformers possibly alleviate over-squashing see, e.g., Müller et al. (2023). However, due to their global aggregation mode, computing an attention matrix with n² entries for an ro-order graph makes them applicable only to small or mid-sized graphs. Further, to capture non-trivial graph structure, they must resort to hand-engineered positional or structural encodings.
Overall, current strategies to mitigate werrsquashing rely on heuristic rewiring methods that may not adapt well to prediction task or employ computationally intensive global attention mechanisms. Furthermore, the impact of probabilistic rewiring on a model's expressive power remains unclear.
Present work By leveraging recent progress in differentiable k-subset sampling (Ahmed et al. 2023), we derive probabilistically rewired MPNNs (PR-MPNNs). Concretely, we utilize an upstream model to learn prior weights for candidate edges. We then utilize the weights to parameterize probability distribution constrained by so-called k-subset constraints. Subsequently, we sample multiple k-edge adjacency matrices from this distribution and process them using downstream model, typically MPNN, for the final predictions task. To make this pipeline trainable via gradient descent, we adapt recently proposed discrete gradient estimation and tractable sampling techniques (Ahmed et al. 2023 Niepert al. 2021
2

---

Published as a conference paper at ICLR 2024
Xie & Ermon, 2019); Figure I for overview of our architecture. Our theoretical analysis explores how PR-MPNNs Ns overcome MPNNs' inherent limitations in expressive power and identifies precise conditions under which they outperform purely randomized approaches. Empirically, demonstrate that our approach effectively mitigates issues like over-squashing and under-reaching. In addition, on established real-world datasets, method exhibits competitive or superior predictive performance compared to traditional MPNN models and graph transformer architectures.
Overall, PR-MPNNs pave the way for the principled design of more flexible MPNNs, making them less vulnerable to potential noise and missing information.
1.1 RELATED WORK
MPNNs are inherently biased towards encoding local structures, limiting their expressive power (Morris et al., 2019, 2021 Xu et al. 2019). Specifically, they are at most as powerful as distinguishing non-isomorphic graphs or nodes with different structural roles as the 1-dimensional Weisfeiler-LLman algorithm (Weisfeiler & Leman 1968). simple heuristic for the graph isomorphism problem; see Sec- tion 2 Additionally, they cannot capture global or long-range information, often linked to phenomena such as under-reaching (Barceló al. 2020) over-squashing (Alon & Yahav 2021 with the latter being heavily investigated in recent works.
Graph rewiring Several recent works aim to circumvent over-squashing via graph rewiring. Perhaps the most straightforward way of graph rewiring is incorporating multi-hop neighhors. For example, Babiel-Ga et al. (2022) rewires the graphs with k-hopp neighbors and virtual nodes and also augments them with positional encodings. MixHop (Abu-El-Haija et al.] (019). SIGN (Frasca et al. 2020), DIGL (Gasteiger et al. 2019), and SP-MPNN (Abboud et al. 2022) can also be considered as graph rewiring as they can reach further-away neighbors in a single e layer. Particularly, Gutteridge al.) (2023) rewires the graph similarly to Abboud et al. (2022) but with a novel delay mechanism, showcasing promising empirical results. Several rewiring methods depend on particular metrics, e.g.,
Ricci or Forman curvature (Bober et al. 2022) and balanced Forman curvature [Topping et al. 2021). In
addition, Deac et al. (2022), [hhrzad et al. (2023) utilize expander graphs to enhance message passing and connectivity, while Karhadkar et al. (2022) resort to spectral techniques, and Banerjee et al. (002)
propose greedy random edge flip approach overcome over-squashin. Refining Topping et al. 2021
Di Giovanni et al. (2023) analyzed how the architectures' width and graph structure contribute to the
over-squashing problem, showing that over-squashing happens among nodes with high commute time, stressing the importance of rewiring techniques. Contrary to our proposed method, these strategies to mitigate over-squashing either rely on heuristic rewiring methods use purely randomized approaches that may not adapt well to a given prediction task. Furthermore, the impact of existing rewiring methods on a model's expressive power remains unclear and we close this gap with our work.
There also exists a large set of works from the field of graph structure learning proposing heuristical graph rewiring approaches; see Appendix A for details.
Graph transformers Different from the above, graph transformers (Dwivedi et al. 2022b He et al. 2023 Müller al. 2023 Rampásek al. [022 Chen et al. 2022) and similar global attention mechanisms (Liu et al. [022 Wu et al. 20211) marked a shift from local to global message passing, aggregating over all nodes. While not understood in principled way, empirical studies indicate that graph transformers
possibly alleviate over-squashing: see, e.g.. Müller et al. (2023). However, all transformers suffer from their quadratic space and memory requirements due to computing an attention matrix.
2 BACKGROUND
In the following, we provide the necessary background.
Notations Let N = 3.......... For n 2 1, let [m] == (1...,, n. C N. We use f... } to denote multisets, i.e., the generalization of sets allowing for multiple instances for each of its elements. A graph
G' pair (G), E(G)) with finite sets of vertices nodes V(G) and edges E(G) C {(u, b) C V(G) u # v}. If not otherwise stated, we set n = [V] and the graph is of order n. We also call the graph G an n-order graph. For ease of notation. we denote the edge {u, v in E(G) by (u, v) or (v, u). Throughout the paper, we use standard notations, e.g., we denote the neighborhood of a vertex v
by (w) and E(v) denotes its discrete vertex label, and so on; see Appendix B for details.
3

---

Published as a conference paper at ICLR 2024
1-dimensional Weisfeiler-Leman algorithm The 1-WL color refinement is a well-studies heuristic for the graph isomorphism problem, originally proposed by Weeifeiler & Leman (1968). Formally, let G = (V) E(G), E) be labeled graph. In each iteration, 0, the 1-WL computes a node coloring C (G) + N, depending on the coloring of the neighbors. That is, in iteration > 0, we set
(v) RELABEL
COOBEE(CC--L) (CCLI(0), ( u (())))
for all nodes v V(G), where RELABEL injectively maps the above pair unique natural number, which has not been used in previous iterations. In iteration 0, the coloring Cj := l. To test if two graphs G and are non-isomorphic, we run the above algorithm in "parallel" on both graphs. If the two graphs have different number of nodes colored E N at some iteration, the 1-WL distinguishes the graphs as non-isomorphic. Moreover, the number of colors between two iterations, t and (t + 1), does not change, i.e., the cardinalities of the images of C, and Ct+++ are equal, or, equivalently,
(o) (w) = (++1()) = C(+1(w),
for all nodes and w in V(G), the algorithm terminates. For such i, we define the stable coloring Ch (w) C, (r). for e in V(G). The stable coloring reached after at most max max[[V(C (V))), (V(V) iterations (Grohe 2017). It is easy to that the algorithm cannot distinguish all non-isomorphic graphs (Cai et al. 1992). Nonetheless, it is a powerful heuristic that can successfully test isomorphism for a broad class of graphs (Babai & Kucera 1979). A function f: V(G) + R for 0, is L-WL-equivalent if C see Appendix for details.
Message-passingg graph neural networks Intuitively, MPNNS learn vectorial representation. i.c., a d-dimensional real-valued vector, representing each vertex a graph by aggregating information from neighboring vertices. Let G (G, L) be attributed graph, following, Gilmer et al. (2017) and Scarselli et al. (2008)), in each layer, > O, we compute vertex features
(+) (b((-), AGG ((h!-1) | N(o))) R
where (?D and (G may be differentiable parameterized functions, e.g., neural networks, and h (f) L,. In the case of graph-level tasks, e.g., graph classification, one uses
he = READOUT ((h(7) E V(G))) E Rd.
to compute single vectorial representation based on learned vertex features after iteration T. Again, READOUT may be a differentiable parameterized function, e.g., a neural network. To adapt the parameters of the above three functions, they are optimized end-to-end, usually through variant of stochastic gradient descent, e.g., Kingma & Ba (015), together with the parameters of a neural network used for classification regression.
PROBALISTICALLY REWIRED MPNNS
Here, we outline probabilistically rewired MPNNs (PR-MPNNs) based on recent advancements in discrete gradient estimation and tractable sampling techniques (Ahmed set of adjacency matrices of n-order graphs. Further, let (G, X) be a n-order attributed graph with an adjacency matrix (G) E Ai and node attribute matrix E Raxd for d > 0. PR-MPNN maintains a (parameterized) upstream model h, 21, x Rnxd + O, typically a neural network, parameterized by u. mapping an adjacency matrix and corresponding node attributes to unnormalized edge priors 0 E O C Ra*n
2023). Let A, denote the
In the following, we use the priors as parameters of a (conditional) probability mass function,
12 II pe., (A(H)
po (A(H))
1.j=1
assigning probability each adjacency matrix in Mt, where pe,, (A(H),; 1) = sigmoid and Pa., (A(H), = 0) = sigmoid((;)). Since the parameters depend on the input graph G, we can view the above probability conditional probability mass function conditioned on the graph G.
4

---

Published as a conference paper at ICLR 2024
Unlike previous probabilistic rewiring approaches, e.g., Franceschi et al. (2019) introduce dependen-
cies between the graph's 's edges by conditioning the probability mass function pe,, (A(H) on k-subset constraint. That is, the probability of sampling any given k-edge adjacency matrix A(H), becomes
pe (a(A(1))/
if [IA la = k. with Z =
P) (A())
PB BE |BIl=
E
(I)
0
otherwise,
The original graph G is now rewired into a new adjacency matrix A by combining N samples (0) (ro.)(((()) for E [M]] together with the original adjacency matrix (()) using differentiable
) E 21, Subsequently. we use the resulting adjacency matrix as input to a downstream model fa, parameterized by d. typically an MPNN, for the final predictions task.
aggregation function g: 2, (N+1) + A, i.c., A = ((A A AMM)
We have so far assumed that the upstream MPNN computes one set of priors h 2, R R which we use to generate a new adjacency matrix A through sampling and then aggregating the adjacency matrices A ..... A A(), In Section 5 we show empirically that having multiple sets of priors from which we sample is beneficial. Multiple sets priors mean that we learn an upstream model hu 21, x Rn -> " M where M is the number of priors. We can then sample and aggregate the adjacency matrices from these multiple sets of priors.
Learning to sample To learn the parameters of the up- and downstream model w =uu, d) of the PR-MPNN architecture, we minimize the expected loss
L(A(G), X, w): E AAA (1(()(
A ; ..),))))
with E y, the targets, a point-wise loss such as the cross-entropy or MSE, and = hr (AG), To minimize the above expectation using gradient descent and backpropagation, need to efficiently draw Monte-Carlo samples from pre xa (A (G)) and estimate VoL the gradients of an expectation regarding the parameters O of the distribution P(0.k)
Sampling To sample an adjacency matrix A (1) from P().x) (A (G)) conditioned on k-edge constraints, and to allow PR-MPNNs to be trained end-to-end, we use SIMPLE (Ahmed et al. (2033), recently proposed gradient estimator. Concretely, we can use SIMPLE to sample exactly from the k-edge adjacency matrix distribution P(o,k) (A (G)) on the forward pass. On the backward pass, we compute the approximate gradients of the loss (which is an expectation over a discrete probability mass function) regarding the prior weights f using
Vol. m (0) with (0))= = (p(e,)) (A (A(G), E R
with an exact and efficient computation of the marginals (0) that is differentiable the backward pass, achieving lower bias and variance. We show empirically that SIMPLE (Ahmed et al. (023) outperforms other sampling and gradient approximation methods such GUMBEL. SOFTSUR-S" ST (Xie & Ermon 2019) and I-MLE (Niepert et al. 202I), improving accuracy without incurring computational overhead.
Computational complexity The vectorized complexity of the exact sampling and marginal inference step is (log & log 1). where k is from our k-subset constraint, and / is the maximum number of edges that we can sample. Assuming a constant number of layers, PR-MPNN's worst-case training complexity is (I) for both the upstream and downstream models. Let n be the number of nodes in the initial graph, and = max( (/mta, )), with las and I. being the maximum number of added and deleted edges. we consider all of the possible edges for Latd. the worst-case complexity becomes O(n²). Therefore, to reduce the complexity in practice, we select a subset of the possible edges using simple heuristics, such as considering the top ladd edges of the most distant nodes. During inference, since we do not need gradients for edges not sampled in the forward pass, the complexity is (()) for the upstream model and (L) for the downstream model, with L being the number of edges in the rewired graph.
4 EXPRESSIVE POWER OF PROBABILISTICALIS REWIRED MPNNS
We now, for the first time, explore the extent which probabilistic MPNNs overcome the inherent limitations of MPNNs in expressive power caused by the equivalence to 1-WL in distinguishing non- isomorphic graphs (Xu et al. 2018 Morris et al. 2019). Moreover, we identify formal conditions under
5

---

Published as a conference paper at ICLR 2024
which PR-MPNNs outperform popular randomized approaches such as those dropping nodes and edges uniformly at random. We first make precise what we mean by probabilistically separating graphs by introducing a probabilistic and generally applicable notion of graph separation.
Let us assume a conditional probability mass function p: 21, + [0, 1] conditioned on a given n-order graph, defined over the set of adjacency matrices of n-order graphs. In the context of PR-MPNNS, p is the probability mass function defined in Section 3 but it could also be any other conditional probability mass parameterized
function over graphs. Moreover, let f: 21, + Rd, for d >> O, be permutation-invarice function mapping sampled graph's 's adjacency matrix to a vector in R The function f could be the composition of an aggregation function و that removes the sampled edges from the input graph G and of a downstream MPNN. Now, the conditional probability mass function separates two graphs G and H with probability p with respect to f if
(B & f(A(H)) -
that is, in expectation over the conditional probability distribution, the vectors [(A(G) and ((A(H) are distinct with probability p.
In what follows, we analyze the case of p being the exactly-k probability distribution defined in Equation ! and f being the aggregation function removing edges and downstream MPNN. However, our framework readily generalizes to the case of node removal, and we provide these theoretical results in the appendix. Following Section 3 we sample adjacency matrices with exactly i edges and them remove edges from the original graph. We aim to understand the separation properties of the probability mass function P(k.0) in this setting and for various types of graph structures. Most obviously, we do not want separate isomorphic graphs and, therefore, remain isomorphism invariant, a desirable property of MPNNs.
Theorem 4.1. For sufficiently large n. for every E (0, 1) and k > O, we have that for almost all pairs, in the sense of Babai et al. (1980), of isomorphic n-order graphs G and H and all permutation-inven invariant, 1- I-WL-equivalent functions f: 2n + Rd, d > 0, there exists a probability mass function P(o.)) that separates the graph G and H with probability at most e with respect to f.
Theorem 4.1 relies on the fact that most graphs have discrete 1-WL coloring. For graphs where the 1- 1-WL stable coloring consists of a discrete and non-discrete part, the following result shows that there exist distributions P(e.k) not separating the graphs based on the partial isomorphism corresponding to the discrete coloring.
Proposition 4.2. Let E (0, 1), k > O, and let G and H be graphs with identical 1-WL stable colorings. Let VG and VH be the subset of nodes of G and H that are in color classes of cardinality 1. Then, for all choices of 1- I-WL-equivalent functions f. there exists a conditional probability distribution p(o,)) that separates the graphs G[Vel and [VH)) with probability at most e with respect to f.
Existing methods such as DropGNN (Papp et al.) 2021) or DropEdge (Rong et al.) (020) are more likely to separate two (partially) isomorphic graphs by removing different nodes or edges between discrete color classes, i.e., on their (partially) isomorphic subgraphs. For instance, in the appendix, we prove that pairs of graphs with m edges exist where the probability of non-separation under uniform edge sampling is at most 1/ m. This is undesirable as it breaks the MPNNs' permutation-ine -ivaariance in these parts.
Now that we have established that distributions with priors from upstream MPNNs are more likely to preserve (partial) isomorphism between graphs, we turn to analyze their behavior in separating the non-discrete parts of the coloring. The following theorem shows that PR-MPNNs are more likely to separate nonmisooee graphs than probability mass functions that remove edges nodes uniformly at random.
Theorem 4.3. For every s E (0, 1) and every k >> O, there exists a pair of non-isomorphic graphs G and H with identical and non-discrete 1-WL stable colorings such that for every 1-WL-equivalent function J.
(1) there exists a probability mass function P(k.e) that separates G and H with probability at least
(1 - E) with respect to f:
(2) removing edges uniformly random separates G and H with probability at most with respect to ʃ-
Finally, we can also show negative result, namely that there exist classes of graphs for which PR-MPNNs cannot do better than random sampling.

---

Published as a conference paper at ICLR 2024
Table I: Comparison between PR-MPNN and baselines on three molecular property prediction datasets. We report results for PR-MPNN with different gradient estimators for k-subset sampling: GUMBEL SOFTSUB-ST (Maddison et al., 2017 Jang et al., 2017, Xie & Ermon 2019) I-MLE (Niepert et al. 2. and SIMPLE (Ahmed al. ) and compare them with the base downstream model, and two graph transformer architectures. The variant using SIMPLE consistently outperforms the base models and is competitive or better than the two graph transformers. We use green for the best model, blue for the second-hest, and red for third. We note with + EDGE the instances where edge features are provided and with EDGE when they are not.
ZINC
OGBG-MOLHIV
ALCHEMY
EDGE 4
EDGE
EDGE f
EDGE
GIN BACKBONE
K-ST SAT K-SG SAT
0.166:20.007
0.115++0.005
0.625=0.039
N/A
0.162:10.013
0.095.10.00
0.613-0.010
N/A
BASE
0.258.10.000
0.207
0.775=0.011
11.1240.00
BASE W.
0.162:40.001
0.101 +000044
0.764=0.019
7.197400094
PR-MPNNGw Gwn (OURS)
0.153.10.003
0.103
0.760-0.025
6.858 090
PR- PR-MPNN (OURS)
0.15110.00
0. 0.10410.008
0.774
1.69210.061
PR- PR-MPNN (OURS)
0.139=0.001
-0.08500000
0.795=0.009
6.447=0.057
PNA
GPS
N/A
0.070.1..004
0.788_00.00
N/A
K-ST SAT
0.164
0. 0.100:20.005
0, 0.625=0.039
N/A
K-SG SAT
0.131 1.0000
0.09410.008
0.613=00010
N/A
4-Cycles 1 prior
4-Cycles 5 5 priors
4-Cycles - 10 priors
1.00
Accuracy
0.95
Avcural
0.90
PR-MPNN
PR-MPNN
PR-MPNN
DNN
DropGNNN
DroPCINN
0.85
10
20
30
40
50
10
20
30
40
50
10
20
30
40
50
#samples
#samples
#samples
Figure 2: Comparison between PR-MPNN and DropGNN on the 4-CYCLES dataset. PR-MPNN rewiring is almost always better than randomly dropping nodes, and is always better with 10 priors.
Proposition 4.4. For every le > there exist non-isomorphic graphs G and with identical 1-WL colorings such that every probability mass function P(e.)) separates the two graphs with the same probability as the distribution that samples edges uniformly at random.
5 EXPERIMENTAL EVALUATION
Here, we explore to what extent our probabilistic graph rewiring leads to improved predictive performance on synthetic and real-world datasets. Concretely, we answer the following questions.
Q1 Can probabilistic graph rewiring mitigate the problems of over-squashing and under-reaching in synthetic datasets?
Q2 Is the expressive power of standard MPNNs enhanced through probabilistic graph rewiring? That is, can we verify empirically that the separating probability mass function of Section can be learned with PR- PR-MPNNS and that multiple priors help?
Q3 Does the increase predictive performance due to probabilistic rewiring apply to (a) graph-level
molecular prediction tasks and (b) node-level prediction tasks involving heterophilic data?
An anonymized repository of our code can accessed at htttps://anonymont 4open science/r/PE: Stancee/r/PR PPPMM/PNRRRP
Datasets To answer Q1, we utilized the SEEIESS-NEIGIBORSMM.... Additionally, we created the TREES-LEAFCOUNS dataset investigate whether our method could mitigate under-reaching issues; see Appendix E for details. To tackle Q2, we performed experiments with the ExP (Abboud et al. 2020) and CSL datasets (Murhhy et al. 2019) to assess how much probabilistic graph rewiring can enhance the models' expressivity. In addition, we utilized the 4-CYCLES dataset from Loukas (2020):
dataset (Alon & Yahav 202
Papp et al. (2021) and set it against standard DropGNN model (Papp et al.
7

---

Published as a conference paper at ICLR 2024
Accuncy
0.5
Transformat PR-MPNN
GUN
Problem rodies
Figure 3: Example graph from the TREES- LEAFCOUNT test dataset with radius 4 (left). PR- MPNN rewires the graph, allowing the downstream MPNN to obtain the label information from the
Figure 4: Test accuracy of our rewiring method (Alon & Yahav 2021) dataset, compared to the reported accuracies from Müller et al. (2023).
the
leaves in one passage-passing step (right).
202I) for comparison while also ablating the performance difference concerning the number of priors and samples per prior. To answer Q3 (a), used the established molecular graph-level regression
datasets ALCHEMY (Chen et al. 2019), ZINC (Jin et al. 2017 Dwivedi et al. 2020). OGBG-MOLHIV (Hu
al. 2020a). QM9 (Hamilton et al. 2017). LRGB (Dwivedi et al. (022b) and five datasets from the
TUDATASET repository (Morris et al. 2020). To answer Q3 (b), we used the CORNELL, WISCONSIN,
TEXAS node-level classification datasets (Pei al. 2020).
Baseline and model contigurations For our upstream model fug we use an MPNN, specifically the GIN
layer (Xu et al. 2019). For edge (u, w) EE)), we compute onn ([h [bb])) E R, where [-II-] is the concatenation operator and Φ is an MLP. After obtaining the prior 0, we rewire our graphs by
sampling two adjacency matrices for deleting edges and adding new edges, i.e., g( A (G), A (1), A (2)
(A(G) - A (1) + A (2) where (1) and A (2) are two sampled adjacency matrices with a possibly different number of edges, respectively. Finally, the rewired adjacency matrix (or multiple adeaeeddaaaa adjacency
matrices) is used in downstream model fa: 24, x Rnxd - y, typically an MPNN, with parameters and y' the prediction target set. For the instance where we have multiple priors, as described in Section 3 A (N) into single adjacency matrix A that we send to a downstream model as described in Figure D or construct a downstream ensemble with multiple aggregated matrices A Ai...., AM. In practice. we always use downstream ensemble before the final projection layer when we rewire with more than one adjacency matrix, and we do rewiring by both adding and deleting edges, please consult Table in the Appendix for more details.
can either aggregate the sampled adjacency matrices A (1)
All of our downstream models fa and base models are MPNNs with GIN layers. When we have access to edge features, we use the GINE variant (Hu et al. 2020b) for edge feature processing. For graph-level tasks, we use mean pooling, while for node-level tasks, we take the node embedding h for a node v. The final embeddings are then processed and projected to the target space by an MLP.
For ZINC, ALCHEMY, and OGBG-MOLHIV, we compare our rewiring approaches with the base downstream model, both with and without positional embeddings. Further, we compare to GPS (Rampášek et al.) 2022) and SAT (Chen et al. 2022), two state-ofithe-art graph transformers. For the TUDATASET, we compare with the reported scores from Giusti et al. (2023b) and use the same evaluation strategy in Xu al. 2 Giusti al. (2023b), i.e., rumning 10- cross- volss-valiotiion and reporting the maximum average validation accuracy. For different tasks, we search for the best hyperparameters for sampling and our upstream and downstream models. See Table in the appendix for the complete description. For ZINC, ALCHEMY, and OGBG-MOLHIV, we evaluate multiple gradient estimators in terms of predictive power and computation time. Specifically, we compare GUMBEL SOFTSUB-ST (Maddison ctt al. al.
2017 Jang et al. 2017, Xie & Ermon 2019). I-MLE (Nicpert et al. (021). and SIMPLE (Ahmed)
2023 The results in terms of predictive power are detailed in Table and the computation time
comparisons can be found in Table in the appendix. Further experimental results on QM9 and LRGB are included in Appendix Gn the appendix.
Experimental results and discussion Concerning Q1, our rewiring method achieves perfect test accuracy up to a problem radius of 6 on both TREES-NEIGHBORSMM MATCH and - demonstrating that it can successfully alleviate over-squashing and under-reaching, see Figure 4 For TREES-LEAM LEASSULT....TT... our model can create connections directly from the leaves to the root, achieving perfect accuracy with a downstream model containing a single MPNN layer. We provide a qualitative result in Figure [3 and a detailed discussion in Appendix Concerning Q2, on the 4-CYCLES dataset, our probabilistic rewiring method matches or outperforms DropGNN. This advantage is most pronounced with and 10 priors, where we achieve 100% task accuracy using 20 samples, as detailed Figure 2 On the ExP dataset, showcase the expressive power of probabilistic rewiring by achieving perfect
8

---

Published as a conference paper at ICLR 2024
Table 2: Comparison between the base GIN model, PR-MPNN, and other more expressive models on the ExP dataset.
Table Comparison between the base GIN model and probabilistic rewiring model on CSL dataset, w/o positional encodings.
ACCURACY 1
MODEL
ACCURACY 1
GIN
0.511=0.021
GIN
0.10010 000
GIN ID-GNN
1.000=0.000
GIN POSENC
1,000+0000
OSAN
1 1.000..0000
PR MPNN (GURS)
0 .088
MPNN (OURS)
1.000=0.000
PR-MPNN POSENC (OURS)
1.000+0.000
Table 4: Comparison between PR-MPNN and other approaches reported in Giusti et al. (2023b); ;
Karhadkar et al. (2022) ; Papp et al. (2021). Our model outperforms existing approaches while keeping a
lower variance in most of the cases, except for NCII, where the WI. Kernel is the best. We use green for
the best model, blue for the second-best and red for third.
MODEL
MUTAG
PTC ICR
PROTEINS
NCI1
NCI109
GK (Ac 3) SHERVASHIDZE ET AL.
2009
81.411.7
55.7...5
71.410.3
12.5
62.410.3
PK (NEUMANN ET AL. 2016
76.022.7
59.5
73.7
82.520.
N/A
WL KERNEL (SHERVASHIDZE ET AL.
2011
90.4
59.924.3
75.0+3.1
86.0=1.8
N/A
DGCNN (ZIIANG ET AL..
2018
85.8+188
58.6#2.6
75.5+0.9
74.4=0.5
N/A
IGN (MARON ET AL.. 2019B
83.9+13.0
58.5+669
76.6+5.5
74.3+2.
72.8+15
GIN (XUET AL. 2019)
89.4+5.6
64.627.0
76.2+2.8
82.7=1.7
N/A
PPGNS (MARON ET AT..
2019A
90.618.7
66.226.6
77.224.7
83.221.1
82.2+1.4
NATURAL GN (DE HAAN
AL.
2020
89.4+1.6
66.8+1.7
71.7+1.
82.4#1.
83.0#19
GSN [BOURITSAS
2022
92.2+7.6
68.2+7.2
76.6+3.0
83.5#2.0
83.5
CIN BODNAR ET AL.
2021
92.7 02.7.16.1
68.2.15.6
77.014.3
83.611.4
84.0
CAN GIUSTTET AL.
2023A
94. 94.1
72.8
78.2
84.5
83.61112
CIN+
ET AL.
2023B
94.4
73.2
80.5
85.3#1.2
84.5
FoSR (KARHADKAR AL.
2022
86.2+1.6
58.5
75.1+0.8
72.9
71.1
DROPONN (PAPP ET AL.
90.447.0
66.3#8.6
76.3+6.1
81.6#1.8
80.8+2.6
PR-MPNN (10-FOLD CV)
98.4
74.3
80.743.9
85.620.8
84.6
accuracy, see Table 2 Besides, our rewiring approach can distinguish the regular graphs from the CSL dataset without any positional encodings, whereas the I-WL-equivalent GIN obtains only random accuracy. Concerning Q3 (a), the results in Table [ show that our rewiring methods consistently outperform the base models on ZINC, ALCHEMY, and OGBG-MOLHIV and are competitive or better than the state-of-the probabilistic rewiring method outperforms existing approaches and obtains lower variance on most of the datasets, with the exception being NCII, where our method ranks second, after the WL kernel. Hence, our results indicate that probabilistic graph rewiring can improve performance for molecular prediction tasks. Concerning Q3 (b), we obtain performance gains over the base model and other existing MPNNs, see Table II in the appendix, indicating that data-driven rewiring has the potential of alleviating the effects of over-smoothing by removing undesirable edges and making new ones between nodes with similar features. The graph transformer methods outperform the rewiring approach and the base models, except on the TEXAS dataset, where our method gets the best result. We speculate that GIN's aggregation mechanism for the downstream models is a limiting factor on heterophilic data. We leave the analysis of combining probabilistic graph rewiring with downstream models that address over-smoothing for future investigations.
GPS and SAT graph transformer methods. On TUDATASET, see Table our
6 CONCLUSION
Here, we utilized recent advances in differentiable k-subset sampling to devise probabilistically rewired message-passing neural networks, which learn to add relevant edges while omitting less beneficial ones, resulting in the PR-MPNN framework. For the first time, our theoretical analysis explored how PR-MPNNs enhance expressive power, and we identified precise conditions under which they outperform purely randomized approaches. On synthetic datasets, we demonstrated that our approach effectively alleviates the issues of over-squashing and under-reaching while overcoming MPNNs' limits in expressive power. In addition, established real-world datasets, we showed that our method is competitive or superior to conventional MPNN models and graph transformer architectures regarding predictive performance and computational efficiency. Ultimately. PR-MPNNs represent significant step towards systematically developing more adaptable MPNNs, rendering them less susceptible to potential noise and missing data, thereby enhancing their applicability and robustness.
9

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
CQ and CM are partially funded by a DFG Emmy Noether grant (468502433) and RWTH Junior Principal Investigator Fellowship under Germany's 's Excellence Strategy. AM and MN acknowledge funding by Deutsche Farschungsgemeinscaaftt (DFG, German Research Foundation) under Germany's Excellence Strategy EXC 2075 390740016, the support by the Stuttgart Center for Simulation Science (SimTech), and the International Max Planck Research School for Intelligent Systems (IMPRS-IS). This work was funded in part by the DARPA Perceptually-enabled Task Guidance (PTG) Program under contract number HR001 12220005, the DARPA Assured Neuro Symbolic Learning and Reasoning (ANSR) Program, and gift from RelationalAI. GVdB discloses financial interest in RelationalAl.
REFERENCES
Ralph Abboud, Ismail Ilkan Ceylan, Martin Grohe, and Thomas Lukasiewiez. The surprising power of graph neural networks with random node initialization. arXiv preprint arXiv:2010.01 179, 2020. 7 20
Ralph Abhoud, Radoslav Dimitrov, and Ismail IIkan Ceylan, Shortest path networks for graph property prediction. In Learning on Graphs Conference, pp. PMLR, 2022.
22
Sami Abu-El-Haija, Bryan Perozzi, Amol Kapoor, Nazanin Alipourfard, Kristina Lerman, Hrayr Harutyunyan, Greg Ver Steeg. and Aram Galstyan. Mixhop: Higher-order graph convolutional architectures via sparsified neighhorhood mixing. In International Conference on Machine Learning, pp. 21-29. PMLR, 2019.
Kareem Ahmed, Zhe Zeng, Mathias Niepert, and Guy Van den Broeck. Simple: A gradient estimator for k-subset sampling. International Conference on Learning Representations, 2023.
20
Uri Alon and Eran Yahav. On the bottleneck of graph neural networks and its practical implications. In International Conference on Learning Representations, 2021. 2
20
22
Laszlo Babai and Ludik Kucera. Canonical labelling of graphs in linear average time. In 20th Annual Symposium Foundations of Computer Science (sfes 1979), pp. 39-46. IEEE, 1979.
Lászłó Babai, Paul Erdos, and Stanley M Selkow. Random graph isomorphism. SlaM Journal on
computing, 9(3):628-633 1980. 6 16 17
Pradeep Kr Banerjee, Kedar Karhadkar, Yu Guang Wang, Uri Alon, and Guido Montúfar. Oversquashing in gnns through the lens of information contraction and graph expansion. In 2022 58th Annual Allerton Conference on Communication, Control, and Computing (Allerton), pp. 1-8. IEEE, 2022. 2
Albert-Laszlo Barabasi and Zoltan N Oltvai. Network biology: Understanding the cell's functional
organization. Nature reviews genetics, 5(2):101-113, 2004. [
Pablo Barceló, Egor V. Kostylev, Mikael Monet, Jorge Pérez, Juan Reutter, and Juan Pablo Silva.
The logical expressiveness of graph neural networks. In International Conference on Learning Representations, 2020. 2 21
Jakub Bober, Anthea Monod, Emil Saucan, and Kevin N Webster. Rewiring networks for graph neural network training using discrete geometry. arXiv preprint arXiv:2207.08026 2022. 3
Cristian Bodnar, Fabrizio Frasca, Yuguang Wang, Nina Otter, Guido F Montufar, Pietro Lio, and Michael Bronstein. Weisfeiler and lehman go topological: Message passing simplicial networks. In
International Conference on Machine Learning. pp. 1026-1033 PMLR, 2021. 9 22
Giorgos Bouritsas, Fabrizio Frasca, Stefanos Zafeiriou, and Michael M Bronstein. Improving graph
neural network expressivity via subgraph isomorphism counting. IEEE Transactions on Pattern
Analysis and Machine Intelligence, ((())6657-663 2022. 9, 22
Mare Brockschmidt. Gnn-film:
Graph neural networks with feature-wise linear modulation. In
International Conference on Machine Learning, pp. 1144-1152, PMLR, 2020. 22
Rickard Brüel-Gatee Gabrielsbbi,,,,nn, Mikhail Yurochkin, and Justin Solomon. Rewiring with positional encodings for graph neural networks. arXiv preprint arXiv: 12674, 2022.
10

---

Published as a conference paper at ICLR 2024
Jin Yi Cai, Martin Fürer, and Neil Immerman. An optimal lower bound on the number of variables for graph identification. Combinatorica, 12(4))389-410, 1992.
Quentin Cappart. Didier Chételat, Elias B. Khalil, Andrea Lodi, Christopher Morris, and Petar Veličkovié. Combinatorial optimization and reasoning with graph neural networks. Journal of Machine Learning Research, 24(130): 2023.
Daniele Castellana and Federico Errica. Investigating the interplay between features and structures in graph learning. 2023. 15
Dexiong Chen, Leslie O'Bray, and Karsten Borgwardt, Structure-aware transformer for graph representa- tion learning. In International Conference Machine Learning, pp. 3469-3489. PMLR, 2022. 2
Guangyong Chen, Pengfei Chen, Chang Yu Hsich, Chee-Kong Lee, Benben Liao, Renjie Liao, Weiwen Liu, Jiezhong Qiu, Qiming Sun, Jie Tang, et al. Alchemy: A quantum chemistry dataset for benchmarking ai models. arXiv preprint arXiv: 120/./006.9.9 094670947, 2019.
Yu Chen, Lingfei Wu, and Mohammed Zaki. Iterative deep graph learning for graph neural networks: Better and robust node embeddings. Advances Neural Information Processing Systems, 33: 19314-19326, 2020.
15
Mark Craven, Dan DiPasquo, Dayne Freitag, Andrew McCallum, Tom Mitchell, Kamal Nigam, and Seán Slattery. Learning extract symbolic knowledge from the world wide web. AAAI/IAAI, 3(3.6)2, 1998.
20
George Cybenko. Approximation by superpositions of a sigmoidal function. Math. Control. Signals Syst., 5(4):455, 1992. 16
Pim de Haan, Taco S. Cohen, and Max Welling. Natural graph networks. Advances in Neural Information Processing Systems, 113636-36466 2020.
22
Andreea Deac, Mare Lackenby, and Petar Veličković. Expander graph propagation. In Learning
Graphs Conference, pp. 38 PMLR, 2022. 3
Francesco Di Giovanni, Lorenzo Giusti, Federico Barbero, Giulia Luise, Pietro Lio, and Michael M Bronstein. On over-squashing in message passing neural networks: The impact of width, depth, and topology. In International Conference on Machine Learning, pp. 7865-7885. PMLR, 2023.
Justin Domke, Implicit differentiation by perturbation. In Advances in Neural Information Processing Systems 23, pp. 523-531. 2010. 19
Vijay Prakash Dwivedi, Chaitanya K Joshi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks. arXiv preprint 20ti::20030009822 2020. g
Vijay Prakash Dwivedi, Anh Tuan Luu, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Graph neural networks with learnable structural and positional representations. In International Conference on Learning Representutions, 2022a. 21
Vijay Prakash Dwivedi, Ladislav Rampášek, Michael Galkin, Ali Parviz, Guy Wolf, Anh Tuan Luu, and Dominique Beaini. Long range graph benchmark. Advances in Neural Information Processing Systems, 35:22326-22340, 2022b.
20 22
David Easley, Jon Kleinberg, et al. Networks, crowds, and markets. Cambridge Books, 2012.
Bahare Fatemi, Layla El Asri, and Seyed Mehran Kazemi. Slaps: Self-supervision improves structure learning for graph neural networks. Advances in Neural Information Processing Systems, 34: 22667-22681, 2021. [5]
Bahare Fatemi, Sami Abu-El-Haine Anton Tsitsulin, Mehran Kazemi, Dustin Zelle, Neslihan Bulut, Jonathan Halcrow, and Bryan Perozzi. Ugsl: A unified framework for benchmarking graph structure learning. arXiv preprint arXiv:2308. 10737, 2023. 15

---

Published as a conference paper at ICLR 2024
Luca Franceschi, Mathias Niepert, Massimiliano Pontil, and Xiao He. Learning discrete structures for graph neural networks. In International Conference on Machine Learning, pp. 1972-1982 1982. PMLR, 2019.
15
Fabrizio Frasca, Emanuele Rossi, Davide Eynard, Ben Chamberlain, Michael Bronstein, and Federico Monti. Sign: Scalable inception graph neural networks. arXiv preprint arXiv:2004. //98, 2020.
Johannes Gasteiger, Stefan Weißenberger, and Stephan Günnemann, Diffusion improves graph learning.
Advances in Neural Information Processing Systems, 32, 2019. 22 25
Justin Gilmer, Samuel S Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning, pp. 1263-1272. PMLR, 2017. 2
Lorenzo Giusti, Claudio Battiloro, Lucia Testa, Paolo Di Lorenzo, Stefania Sardellitti, and Sergio Barbarossa. Cell attention networks. In 2023 International Joint Conference on Neural Networks
(IJCN), pp. 1-8. IEEE, 2023a.
Lorenzo Giusti, Teodora Reu, Francesco Ceccarelli, Cristian Bodnar, and Pietro Liò. Cin++: Enhancing topological message passing. arXiv preprint arXiv:2306.03661, 2023b.
22
Martin Grohe. Descriptive complexity, canonisation, and definable graph structure theory, volume 47. Cambridge University Press, 2017.
Benjamin Gutteridge, Xiaowen Dong, Michael M Bronstein, and Francesco Di Giovanni. Drew: dynamically rewired message passing with delay. In International Conference on Machine Learning, pp. 12252-12267. PMLR, 2023.
25
William Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Advances in Neural Information Processing Systems, pp. 1024- 1034, 2017. 8
Xiaoxin He, Bryan Hooi, Thomas Laurent, Adam Perold, Yann LeCun, and Xavier Bresson. generalization of vit/mlp-mixer to graphs. In International Conference on Machine Learning, pp. 12724-12745. PMLR, 2023. 2
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. Advances in Neural Information Processing Systems, 33:22118- 18-22133, 2020a.
20
Weihua Hu, Bowen Liu, Joseph Gomes, Marinka Zitnik, Percy Liang, Vijay Pande, and Jure Leskovee. Strategies for pre-training graph neural networks. In International Conference Learning Representa- tions, 2020b. g
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameteriization with gumbel-softmax. In International Conference on Learning Representations, 2017.
Wei Jin, Yao Ma, Xiaorui Liu, Xianfeng Tang, Suhang Wang, and Jiliang Tang. Graph structure learning for robust graph neural networks. arXiv preprint 1000:20005.0003, 2020. 15
Wengong Jin, Connor Coley, Regina Barzilay, and Tommi Jaakkola. Predicting organic reaction outcomes with weisfeiler-lehman network. Advances in Neural Information Processing Systems, 30, 2017. 8
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figumov, Olaf Ronneberger,
Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583-588) 2021.
Kedar Karhadkar, Pradeep Kr Banerjee, and Guido Montúfar. Fosr: First-order spectral rewiring for addressing oversquashing in gnns. arXiv preprint arXiv:2210.11790 2022. 3
Anees Kazi, Luca Cosmo, Seyed-Ahee Ahmad Ahmadi, Nassir Navab, and Michael M Bronstein. Differentiable graph module (dgm) for graph convolutional networks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(2): 15(2):1606-1617, 2022. 15
12

---

Published as a conference paper at ICLR 2024
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015. 4
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations, 2017.
Johannes Kliepera, Janek GroB, and Stephan Günnemann. Directional message passing for molecular graphs. In International Conference on Learning Representations, 2020.
Moshe Leshno, Vladimir Ya. Lin, Allan Pinkus, and Shimon Schocken. Multilayer feedforward networks with a nonpolynomial activation function can approximate any function. Neural Networks, 6(6): 861-867, 1993. T6
Meng Liu, Zhengyang Wang, and Shuiwang Ji. Non-local graph neural networks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(12):00200-00276, 2021.
Nian Liu, Xiao Wang, Lingfei Wu, Yu Chen, Xiaojie Guo, and Chuan Shi. Compact graph structure learning via mutual information compression. In Proceedings of the ACM Web Conference 2022, pp. 1601-1610, 2022a.
15
Yixin Liu, Yu Zheng, Daokun Zhang, Hongxu Chen, Hao Peng, and Shirui Pan. Towards unsupervised deep graph structure learning. In Proceedings of the ACM Web Conference 2022, pp. 1392-1403, 2022b. [5
Andreas Loukas. What graph neural networks cannot learn: depth vs width. In International Conference on Learning Representations, 2020.
Chris J. Maddison, Andriy Mnih, and Yee Whye Tch. The concrete distribution: A continuous relaxation of discrete random variables. International Conference on Learning Representations, 2017. 7
Haggai Maron, Heli Ben-Hamu, Hadar Serviansky, and Yaron Lipman. Provably powerful graph networks. In Advances in Neural Information Processing Systems, pp. 2153-2164, 2019a.
Haggai Maron, Heli Ben- Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. In International Conference on Learning Representations, 2019b. 9 22
Christopher Morris, Martin Ritzert, Matthias Fey, William I. Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and leman go neural: Higher-order graph neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 4602-4609, 2019.
16
Christopher Morris, Nils M Kriege, Franka Bause, Kristian Kersting, Petra Mutzel, and Marion Neumann. Tudataset: A collection of benchmark datasets for learning with graphs. arXiv preprint vrXiv:2007.08663 2020. g 20
Christopher Morris, Yaron Lipman, Haggai Maron, Bastian Rieck, Nils M Kriege, Martin Grohe, Matthias Fey, and Karsten Borgwardt. Weisfeiler and leman go machine learning: The story so far. arXiv preprint arXiv:21 12.09992, 2021.
Luis Müller, Mikhail Galkin, Christopher Morris, and Ladislav Rampášek. Attending to graph transformers. arXiv preprint arXiv:2302.0418/, 2023.
Ryan Murphy, Balasubramaniam Srinivasan, Vinayak Rao, and Bruno Ribeiro. Relational pooling for graph representations. In International Conference on Machine Learning, pp. 4663-4673, PMLR, 2019.
20
Marion Neumann, Roman Garnett, Christian Bauckhage, and Kristian Kersting. Propagation kernels: efficient graph kernels from propagated information. Machine learning, 102:209- 245, 2016. 99
Mathias Niepert, Pasquale Minervini, and Luca Franceschi. Implicit mle: backpropagating through discrete exponential family distributions. Advances in Neural Information Processing Systems, 34: 14567-14579, 2021.
20
13

---

Published as a conference paper at ICLR 2024
Pál András Papp. Karolis Martinkus, Lukas Faber, and Roger Wattenhofer. Dropgnn: Random dropouts increase the expressiveness of graph neural networks. Advances in Neural Information Processing
Systems, 34:21997-220099 2021. 6 7
Hongbin Pei, Bingzhe Wei, Kevin Chen-Chuan Chang, Yu Lei, and Bo Yang. Geom-gen: Geometric
graph convolutional networks. In International Conference on Learning Representations, 2020. 20
22
25
Ladislav Rampášek, Michael Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Dominique Beaini. Recipe for general, powerful, scalable graph transformer. Advances in Neural Information Processing Systems, 35: 14501-14515, 2022.
Patrick Reiser, Marlen Neubert, André Eberhard, Luca Torresi, Chen Zhou, Chen Shao, Houssam Metni, Clint van Hoesel, Henrik Schopmans, Timo Sommer, et al. Graph neural networks for materials science and chemistry. Communications Materials, 3(1):93, 2022. [
Yu Rong, Wenbing Huang, Tingyang Xu, and Junzhou Huang. Dropedge: Towards deep graph convolu- tional networks on node classification. In International Conference on Learning Representations, 2020.
Dan Roth. On the hardness of approximate reasoning. Artif. Intell. 8211-)2273323330 1996. doi: 10.1016/ M0043702(94)0009921 /ttp://t///////...o///0000000000
00092-1
20
Avishkar Saha, Oscar Mendez, Chris Russell, and Richard Bowden. Learning adaptive neighborhoods for graph neural networks. arXiv preprint arXiv:2307.0906. 2023. [5
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE Transactions on Neural Networks, 20(1): 2008a.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. Computational capabilities of graph neural networks, IEEE Transactions on Neural Networks, 20(1): 81-102, 2008b.
Nino Shervashidze, SVN Vishwanathan, Tobias Petri, Kurt Mchlhorn, and Karsten Borgwardt. Efficient graphlet kernels for large graph comparison. In Artificial Intelligence and Statistics, pp. 488-495.
PMLR, 2009. g
Nino Shervashidze, Pascal Schweitzer, Erik Jan Van Leeuwen, Kurt Mehlhorn, and Karsten M Borgwardt. Weisfeiler-lehman graph kemnlls. Journal of Machine Learning Research, 12(9), 2011. 9
Hamed Shirzad, Ameya Velingker, Balaji Venkatachalam, Danica Sutherland, and Ali Kemal Sinop. Exphormer: Sparse transformers for graphs. arXiv preprint (:Xi::2303306147, 2023.
Jake Topping. Francesco Di Giovanni, Benjamin Paul Chamberlain, Xiaowen Dong, and Michael M
Bronstein. Understanding wverssquahinn and bottlenecks on graphs via curvature, arXiv preprini
1rXiv:2/11.14 14522, 2021. 3 22 25
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference on Learning Representations, 2018. [
Boris Weisfeiler and Andrei Leman. The reduction of a graph to canonical form and the algebra which appears therein. nti, Series, 2(9):12-16, 1968.
Zhanghao Wu, Paras Jain, Matthew Wright, Azalia Mirhoseini, Joseph Gonzalez, and Ion Stoica. Representing long-range context for graph neural networks with global attention. Advances in Neural Information Processing Systems, 14:13266-13279, 2021.
Sang Michael Xie and Stefano Ermon. Reparameterizable subset sampling via continuous relaxations. International Joint Conference Artificial Intelligence, 2019.
Keyulu Xu, Chengtao Li, Yonglong Tian, Tomohiro Sonobe, Ken-ichi Kawarabayashi, and Stefanie Jegelka. Representation learning on graphs with jumping knowledge networks. In International Conference Machine Learning, pp. 5453-5462. PMLR, 2018.
14

---

Published as a conference paper at ICLR 2024
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In International Conference on Learning Representations, 2019.
3 8
22
Rui Xue, Haoyu Han, MohamadAli Torkamani, Jian Pei, and Xiaorui Liu. Lazygnn: Large-scale graph neural networks via lazy propagation. arXiv preprint arXiv:2302.01550 2023.
Donghan Yu, Ruohong Zhang, Zhengbao Jiang, Yuexin Wu, and Yiming Yang. Graph-revised convolutional network. In Machine Learning and Knowledge Discovery in Databases: European Conference, pp. 378-393, 2021. 15
Muhan Zhang, Zhicheng Cui, Marion Neumann, and Yixin Chen. An end-to-end deep learning architecture for graph classification. In Proceedings of the AAAI Conference on Artificial Intelligence,
volume 32, 2018. 9 22
Tong Zhao, Yozen Liu, Leonardo Neves, Oliver Woodford, Meng Jiang, and Neil Shah. Data augmentation for graph neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 11015-1 1023, 2021. 15
Zhiyao Zhou, Sheng Zhou, Bochao Mao, Xuanyi Zhou, Jiawei Chen, Qiaoyu Tan, Daochen Zha, Can Wang, Yan Feng, and Chun Chen. Opengsl: A comprehensive benchmark for graph structure learning.
arXiv preprint urXiv:2306. 10280, 2023. 15
Dongcheng Zou, Hao Peng, Xiang Huang, Renyu Yang. Jianxin Li, Jia Wu, Chunyang Liu, and Philip Yu. Se-gsl: A general and effective graph structure learning framework through structural entropy optimization. arXiv preprint arXiv:2303.09778, 2023. 15
A ADDITIONAL RELATED WORK
In the following, we discuss additional related work.
Graph structure learning The field of graph structure learning (GSL.) is a topic related to graph rewiring. Motivated by robustness and more general purposes, several GSL works have been proposed. [Jin et al. (2020) optimizes graph structure from scratch with some loss function as bias. More generally, an edge scorer function is learned, and modifications are made to the original graph structure (Chen et al. 2020 Yu et al. 2021 Zhao et al. (02T). To introduce discreteness and sparsity, Kazi et al. (022 Franceschi al. (2019) Zhao et al. (2021) leverage Gumbel and Bernoulli discrete sampling, respectively. (ahaa et al.
(2023) incorporates end-to-end differentiable discrete sampling through the smoothed Heaviidde function.
Moreover, GSL also benefits from self-supervised or unsupervised learning approaches; see, e.g., Zou ct al.
2023 Fatemi et al. (2021); Liu et al. (2022a b). For comprehensive survey of GSL. see Fatemi
ct al. 2203) Zhou et al. 2023). In the context of node classification. there has been recent progress in
understanding the interplay between graph structure and features (Castellana & Errica| 2023).
The main differences to the proposed PR-MPNN framework are as follows: (a) for sparsification, existing GSL approaches typically use k-NN algorithm, a simple randomized version of k-NN, or model edges with independent Bernoulli random variables. In contrast, PR-MPNN uses a proper probability mass function derived from exactly-k constraints. Hence, we introduce complex dependencies between the edge random variables and trade-off exploration and exploitation during training, and (b) GSL approaches do not use exact sampling of the exactly-k distribution and recent sophisticated gradient estimation techniques. However, the theoretical insights we provide in this paper also largely translate to GSL approaches with the difference that sampling is replaced with arg max operation and, therefore, should be of independent interest to the GSL community.
B EXTENDED NOTATION
A graph is a pair (G),, E(G)) with finite sets of vertices or nodes V() and edges E(G) S Hu,,) C V(G) u + v). If not otherwise stated, we set == [V(G)), and the graph is of order n. We also call the graph G an n-order graph. For ease of notation, we denote the edge {u, 8} in
E(G) by (u,v) or u). A (vertex- llabeled graph G is triple (V(G), E) with (vertex- label function E: VGG + N. Then {(v) is a label of v, for o in (G).. An attributed graph G is a triple (V(G), F((G), a) with graph (V) F(G)) and wertex--attrrbuun function a: V(G) Rxxd for
15

---

Published as a conference paper at ICLR 2024
some d > O. That is, contrary to labeled graphs, allow for vertex annotations from an uncountable set. Then (vv) ) is an attribute or feature of v, for u in V(G).. Equivalently, we define an n-order attributed pair G = (G, L). where G (V(G), E(G)) and L in Raxd is node attribute matrix. Here, we identify (() with [n]. For matrix L in R and v in [n], we denote by Lg in Rixd the oth row of L such that L,. = a(v). Furthermore, we can encode an n-order graph G via an adjacency matrix AGG E {0, 1,xm, where Aij = 1 if, and only, if (i,j) E EG).. We also write Rd for Rixd
graph G = (V(G), E(G), a)
The neighborhood of in V(G) is denoted by N(v)== - (m (v, u) [(G)) and the degree of a vertex v is [N)"] Two graphs G and H are isomorphic and we write G = if there exists a bijection yo: (G) + (HH) preserving the adjacency relation, i.c., (u, v) is in E if and only if (p(u), (Q))) is in E Then & is an isomorphism between G and H. In the case of labeled graphs, we additionally require that I(v) = I()(()) for v V(G), and similarly for attributed graphs. Further, we call the equivalence classes induced by = isomorphism types.
node coloring is a function e: V(G) Rd, d > O, and we say that c(u) is the color of v V(G). A node coloring induces edge coloring Ec: F(G) - N. where (u, v) + {c(u), c(u)) for (u, v) F.GG). node coloring (edge coloring) e refines a node coloring (edge coloring) d, written E if c(()) <(w) implies d(o) d(w) for every w, V(G) (v, w. E E(G)). Two colorings are equivalent if с E and E c, in which case we write c = A color class Q S V(G) node coloring is maximal set of nodes with c(t) c(w) for every o, E Q. A node coloring is called discrete all color classes have cardinality 1.
C MISSING PROOFS
In the following, we outline missing proofs from the main paper.
Theorem C.1 (Theorem 4.1 in the main paper). For sufficiently large n. every E (0, 1) and > we have that for almost all pairs, in the sense of Babai et al. (1980), of isomorphic n-order graphs G and and all cermutation-ivvarinn conditional probability mass function k)) that separates the graph G and H with probability at most regarding f.
WWL-equivalen functions f: 24, - RR,, O, there exists
Before proving the above result, we first need three auxiliary results. The first one is the well-known universal approximation theorem for multi-layer perceptrons.
Theorem C.2 Cyhenko 1992;; Leshno et al. (193)). Let a: T -> TR. be continuous and not polynomial. Then for every continuous function J K RR,", where K R. is a compact set, and every & O there is a depth-two multi-layer perceptron N with activation function funntiooon g(1) on layer 1 and no activation function on layer 2 (i.e., g(2) is the identity function) computing a function fN such that
sup x
[[f(x) -x((x) <.
Building on the first, the second result shows that an MPNN can approximate real-valued node colorings of given finite graph arbitrarily close.
Lemma C.3. Let G be an n-order graph and let e: ((G) R4, d O, be 1-WL-equivalent node coloring. Then, for all E >> O, there exists a ipermutation-equivean MPNN f: V(G) + IRd, such that
max (CV
[ <E
Proof sketch. First, by (Morris et al. 2019 Theorem 2), there exists an 1-WL-equivalent MPNN m: (G) + Rd such that
C m.
Since the graph's number of vertices, by assumption, is finite, the cardinality of the image K = m also finite. Hence, we can find a continuous function g: K - Rd such that (g m = c(v) for v E V(G). Since K is finite and hence compact and g is continuous, by Theorem C.2 we can approximate it arbitrarily close with a two-layer multi-layer perceptron, implying the existence of the MPNN f.
16

---

Published as a conference paper at ICLR 2024
5+1
(8)
(b)
(c)
(d)
Figure Example graphs used in the theoretical analysis.
The third result lifts the previous result to edge colorings.
Lemma C.A. Let G be an n-order graph and let c: E(G) IR d> O, be 1-WL-equivalent edge coloring. Then, for all >> O, there exists a (permatation-eeuuivrrianl) MPNN f: E(G) IRd, such that
[lf(e) - c(e)ll <. eEB(G)
max
Proof sketch. The proof analogous the proof of Lemma
We note here that can extend the above results to any finite subset of n-order graphs. We are now ready to prove Theorem C.I
Proof sketch. Following Babai et al. 1980). for sufficiently large order n, the 1-WL will compute a discrete coloring for almost any n-order graph. Concretely, they showed that algorithm equivalent to the 1-WL computes a discrete coloring of graphs sampled from the Erdōs- Rényi random graph model Gnn, 1/2) with the probability of failure bounded by O Since the G(n, 1/2) model assigns uniform distribution over all graphs, the 1-WL succeeds on "almost all" graphs.
By the above, and due to Lemmas C.3 and C.4 every node, and thereby any edge, can be assigned a distinct arbitrary prior weight with upstream MPNN. Consequently, there exists an upstream MPNN that returns high prior O,y for exactly k: edges (0, for exactly nodes) such that sampling from the exactly-k distribution returns these & edges (nodes) with probability at least VI- Specifically, we know that the upstream MPNN can return arbitrary priors 0zj. and want to show that, given some E (0, 1), there exists at least one such that for a set S of & edges (nodes) G, we have Po. (S) 2
Let be the number of edges in the graph G' That is, from (o())) > 8Z. Without losing generality, let wi and W22 be two prior weights with w1 > W22 such that 0: wi for the edges (nodes) in S and 0; w2 otherwise. Then po(S) 2 87 becomes use the upper bound S w + ((T) 1 w k-1 and obtain & (1 6) 1 ((0)))((())))) 1)--1 Therefore, prior f exists, and we can obtain it by using the derived set 5 v1 11-e. The sampled & edges (nodes) then identical for both graphs with probability at least vI-E" = E and, therefore, the edges (nodes) that are removed are isomorphic edges (nodes) with probability at least 1 When remove pairs of edges (nodes) from two isomorphic graphs that are mapped to each other via an isomorphism, the graphs remain isomorphic 1-WL coloring. Since the two graphs have the same 1-WL coloring.
obtain probability definition,
of 5 6( (+) (m-)))
inequality. Now,
and, therefore, must have the an MPNN downstream model f cannot separate them.
Proposition C.5 (Proposition 4.2 in the main paper). Lete (0,1), O, and let G and H be graphs with identical 1 1-WL stable colorings. Let VG and VH be the subset of nodes of G and H that are in color classes of cardinality Then, for all choices of 1--L-equivvalen functions there exists conditional probability mass function P(O.k) k) that separates the graphs G[Vc] and [VH) with probability at most regarding f.
Proof sketch. Since the graphs G[Vc] and H[Va] have discrete coloring under the 1-WL and the graphs and have identical 1-WL colorings, follows that there exists an isomorphism :p: G[Va] [Vn).
17

---

Published as a conference paper at ICLR 2024
nodes
nodes
nodes
nodes
(a)
(b)
(d)
Figure 6: The graphs used in the proof of Theorem C.6 for node sampling.
pairs
1 pairs
pairs
pairs
(a)
Figure 7: The graphs used in the proof of Theorem C.6 for edge sampling.
Analogous to the proof of Theorem 4.1 we can now show that there exists a set of prior weights that ensures that the exactly-k sample selects the same subset of edges (nodes) from, respectively, the same subset of edges from Gi [V]] and G2[V6] (the same subset of nodes from Va and Vir) with probability at least /I - e. Note that the cardinality of the sampled subsets could also empty since the priors could be putting higher weight nodes (edges) with non-discrete color classes.
Regarding uniform edge removal, consider the graphs in Figure (b) (b) and (c). With distribution based on an MPNN upstream model, the probability separating the graphs by removing edges in the isomorphic subgraphs induced by the nodes with colors to 5 t can be made arbitrarily small for any t. However, the graphs would still be separated through samples in the parts whose coloring is non-discrete. In contrast, sampling uniformly at random separates the graphs in these isomorphic subgraphs with probability converging towards 1 with increasing! l.
Theorem C.6. For every E (0, 1) and every k > O, there exists pair of non-isomorphic graphs G and H. with identical and non-discrete 1-WL stable colorings such that for every I-WL-equivalent fanction
(1) there exists probability mass function P(k.0) that separates G and H with probability at least
(1 E) with respect to fi
(2) removing edges uniformly random separates G and H with probability at most with respect
to f-
Proof. We distinguish the cases: (1) sampling nodes be removed and (2) sampling edges be removed from the original graphs.
For case (1), where we sample nodes to removed, consider the graphs Figure For k we take the graphs (a) and (b). Both of these graphs have the same 1-WL coloring, indicated by the color numbers of the nodes. To separate the graphs, we need to sample and remove one of the nodes with color 1 or 2. Removing node with color 3 would lead again an identical color partition for the two graphs. Removing nodes with color or 2 achievable by placing high prior weight O,, on nodes of one of the corresponding color classes; see Lemma C.3 Without loss of generality, we choose all nodes in color class 2 and set the prior weight 0,, such that a node in this color class is sampled with probability vI Since random sampler would uniformly sample any of the nodes, we simply have increase the number t of nodes with color such that the probability randomly sampling a node of the color classes 1 or 2 is smaller than equal Ve.
18

---

Published as a conference paper at ICLR 2024
For k >> 1, construct the graphs depicted in Figure (c) and (d) for k. 2. These are constructed by first taking a (A + 1) -cycle and connecting to each node of the cycle the nodes with color class 1 in the two ways shown in Figure (c) and (d). Finally,
connect nodes of color class 3 to each of the nodes in the cycle. These graphs can be separated by sampling k nodes from either the color class 1 or 2. For instance, removing k nodes from color class 2 always creates l disconnected subgraphs of size 2 in the first parameterized graph but not the second. By Lemma C.3] we know that we can find an upstream model that leads to prior weights 0, such that sampling nodes from a color class has probability at least vI-c; E; see the proof of Theorem C.1 As argued before, by increasing the number of nodes with color class 3, we can make the probability that uniform sampler picks a node with color classes 1 or be
(c)
less than or equal to Ve.
For case (2), where we sample edges to be removed from the original graph, consider the graphs in Figure 7 For k = 1, we take the graphs (a) and (b). Both of these graphs have the same 1-WL coloring, indicated by the color numbers of the nodes. To separate the graphs, we need to sample from each graph or (2)
1 and C² (v) = 2. Removing an edge between two nodes with color class 3 in both graphs would lead again to an identical color partition of the two graphs. Removing
an edge (u, v) such that either C2 (u) Co (t)
edge between the color classes (1, 1) and (1, 2) is possible
by Lemma C.4] and choosing a prior weight large enough such that the probability of sampling an edge between these color classes is at least v1 e; the proof of Theorem C.1 Since random sampler would sample an edge uniformly at random, we simply have to increase the number of nodes with color class 3 such that the probability of sampling an edge between the color classes (1, 2) or (1, 2) smaller than VE.
For k >>> we construct the graphs depicted in Figure 7 c) and (d) for k = 2. We first take (k -yycle and connect to each node of the said cycle the nodes with color classes 1 in the two different ways shown in Figure (cc and (d). Finally, we again add pairs of connected nodes with color class 3. The two graphs can be separated by sampling k edges between the color classes (1. 1), (1, 2), and (2, 2). For instance, sampling edges between nodes in color class leads to a disconnected subgraph of size 3 in the first graph but not the second. By Lemma C.4 we know that we can learn an upstream MPNN that results in prior edge weights Ou. for all edges (u, v) where both u and v are in color class 2, such that sampling k of these edges has probability least /1 C, see the proof of Theorem C.1 Again, by increasing the
number of nodes with color class 3. we can make the probability that uniform sampler picks an edge
between the color classes (1, 1), (1, 2) or (2, 2) to be less than or equal to VE.
Finally, we can also show negative result, i.e., graphs exist such that PR-MPNNs cannot do better than random sampling.
Proposition C.7 (Proposition 4.4 in the main paper). For every k > O, there exist non-is con-isomorphic graphs H and H with identical 1- WL colorings such that every probability mass function p(o.)) separates the two graphs with the same probability as the distribution that samples nodes (edges) uniformly at random.
Proof. Any pair of graphs where the 1-WL coloring consists of single color class suffices to show the result. For instance, consider the graphs in Figure(5jd where all nodes have the same color. In fact, any pair of non-isomorphic d-regular graphs for d >> 0 works here An MPNN upstream model cannot separate the prior weights of the nodes and, therefore, behaves as uniform sampler.
D SIMPLE: SUBSET IMPLICIT LIKELIHOOD
In this section, we introduce SIMPLE, which is a main component of our work. The goal of SIMPLE is to build a gradient estimator for VoL(X, y; w. It inspired by hypothetical sampling-free architecture, where the downstream neural network fa is a function of the marginals, l == (()) = {po(z) Σ, 2 = ,) instead of a discrete sample z, resulting in a loss Lm s.l.
VeLm (X, y: w) 0 TV (famu,
When the marginals (0)) can be efficiently computed and differentiated, such hypothetical pipeline can be trained end-to-end. Furthermore, Domke (010)) observed that, for an arbitrary loss function e. defined on the marginals, the Jacobian of the marginals w.r.t. the logits is symmetric. Consequently. computing the gradient of the loss w.r.t. the logits, V Lm (X, y w), reduces to computing the directional derivative, or the Jacobian-vector product, of the marginals w.r.t. the logits in the direction of the gradient
19

---

Published as a conference paper at ICLR 2024
of the loss. This offers alluring opportunity, i.e., the conditional marginals characterize the probability of each 2; in the sample, and could be thought of as a differentiable proxy for the samples. Specifically, by reparameterizing x as a function of the conditional marginal f under approximation 0, se I as proposed by /Niepert et al. (2021), and using the straight-through estimator for the gradient of the sample w.r.t. the marginals on the backward pass, SIMPLE approximate
VoL(X, w) x ( ,XX
where the directional derivative of the marginals can be taken along any downstream gradient, rendering the whole pipeline end-to-end learnable despite the presence of sampling.
Now, estimating the gradient of the loss w.r.. the parameters can be thought of as decomposing into two sub-problems: (P1) Computing the derivatives of conditional marginals ōap((0) ()), which requires the computation of the conditional marginals, and (P2) Computing the gradient of the loss w.r.t. the samples V&L(X, y/ w) using sample-wise loss, which requires drawing exact samples. These two problems are complicated by conditioning the h-subset constraint, which introduces intricate dependencies to the distribution and is infeasible to solve naively, e.g., by enumeration. Next, we will show the solutions that SIMPLE provide each problem, the heart of which is the insight that we need not care about the variables' order, only their sum, introducing symmetries that simplify the problem.
Derivatives of conditional marginals In many probabilistic models, the marginal inference is a #P-hard problem (Roth, (996), and this is not the case for the k-subset distribution. 'Teorem 1 in Ahmed et al. (2023) shows that the conditional marginals correspond to the partial derivatives of the log probability of the k-subset constraint. To see this, note that the derivative of a multi-linear function regarding a single variable retains all the terms referencing that variable and drops all other terms; this corresponds exactly to the unnormalized conditional marginals. By taking the derivative of the log probability, this introduces the ks-ubbset probability in the denominator, leading to conditional marginals. Intuitively, the rate of change of the k-subset probability w.r.t. variable only depends that variable through its length-k subsets. They further show in Proposition I in Ahmed et al. (2023) that the log probability of the exactly-k constraint pe E5 2j k) is tractable is tractable as well as amenable to auto-differentiation, solving problem (P1) exactly and efficiently.
Gradients of loss w.r.f. samples What remains is estimating the loss value, requiring faithful sampling from the k-subset distribution. To perform exact sampling from the k-subset distribution, SIMPLE starts by sampling the variables in reverse order, that is, it samples in through 21. The intuition is that, having 2 1) with a Hamming weight of k j. it samples Z4 with a probability of choosing k - j of n 1 variables and the rith variable given that we choose k - + 1 of n variables, providing an exact and efficient solution to problem (P2).
sampled (Zn, Rn-1,
By combining the use of conditional marginal derivatives in the backward pass and the exact sampling in the forward pass. SIMPLE can achieve both low bias and low variance in its gradient estimation. We refer the readers SIMPLE (Ahmed al. 2023) for the proofs and full details of the approach.
E DATASETS
Here, we give additional information regarding the datasets. The statistics of the datasets in our paper can be found in 5 Among them, ZINc, ALCIIEMY, MUTAG, PTC_MR, NCII, NC1109, PROTEINS, IMDB-B, and IMDB-M are from TUDatasets Morris et al. 2020). Whereas PEPTIDES-FUNT and PEPTIDES-STRUCT STRUCT are featured in Dwivedi et al. 2022b). Besides, CORNELL, TEXAS and WISCONSIN are WebKB datasets (Craven al. 1998) also used in Pei al. (2020). The OGB datasets are credited to Hu et al. (2020a). Moreover, we also incorporate synthetic datasets from the literature. EXP dataset consists of partially isomorphic graphs as described in Abboud et al. (2020), while the graphs in the CSL dataset are synthetic regular graphs proposed in Murphy et al. (2019). The construction of BEGGLS-NEIGHBBRRMM...
dataset is introduced in Alon & Yahav (2021).
dataset, we propose own TREES-LEAFCEE LEAFCOUNT dataset. We fix a problem radius R > O and retrieve the binary representation of all numbers fitting into 2 bits. This construction allows to create 2d unique binary trees by labeling the leaves with "0" and "1" corresponding to the binary equivalents of the numbers. A label is then assigned to the root node, reflecting the count of leaves tagged with "1". From the resulting graphs, we sample to ensure equal
Similar to the TREEES-NEIGIBORSM..N
20

---

Published as a conference paper at ICLR 2024
Table 5: Dataset statistics and properties for graph-level prediction tasks, Continuous vertex labels following Gilmer et al. 2017), the last three components encode 3D coordinates.
PROPSaTIES
GRAZHS
OF TARGETS
LABELS
ALCHEMS
222.559
12
MAD
15.1
10.4
QMP
439
13
MAE
15.0
18.6
((13/D))
7(4)
Zine
2494556 456
MAE
23.1
24.9
EXP
200
ACC
64.5
DGBG MOLIDV
41
ROCAUC
25.5
23.5
CORSELL
ACC
163.0
200.0
TEXAS
315.0
The BR
16 16,000
ACD
31
Tares Nemoeeooooooonoeoeoooooo
14100
ACT
11
61
AP
150.9
1937
MUTAG
ACE
17.9
19.8
FTC.JM MR
346
ACC
14.5
14.
NCII
4130
ACC
23.9
22.3
29.7
PROTEINS
ACC
72.6
IMDB-N M
1500
ACC
13.0
65.9
IMD8-B B
1000
ACC
19.7
96.5
class distribution. The task requires a model to predict the root label, thereby requiring strategy capable of conveying information from the leaves to the root.
We aim to have controlled environment to observe our upstream model hu can sample meaningful edges for the new graph configuration. Conventionally, minimum of R wessage-passing layers is required to accomplish both tasks (Barceló et al. 2020; Alon Yahav] 202T). ). However, a single-layer upstream MPNN could trivially resolve both datasets, provided the rewired graphs embed direct pathways from the root node to the leaf nodes containing the label information. To circumvent any potential bias within the sampling procedure, upstream model hu. along with a single-layer GIN architecture serving as the downstream model fa. For each problem radius, sample exactly l: 2n edges. Indeed, method consistently succeeded in correctly rewiring the graphs in all tested scenarios, extending up to a problem radius of R = 6, and achieved perfect test accuracy both datasets. Figure 3 presents qualitative result from the TREES- LEAFCOUNT dataset, further illustrating the capabilities of our approach.
utilize the self-attentio mechanism described in Section 3 as
F HYPERPARAMETER AND TRAINING DETAILS
Experimental Protocol Table 8 lists our hyperparameters choices. For all our experiments, we use early stopping with an initial learning rate 0.001 that we decay by half on a plateau.
We compute each experiment's mean and standard deviation with different random seeds over a minimum of three runs. We take the best results from the literature for the other models, except for SAT on the OGBG-MOLHIV, where we use the same hyperparameters as the authors use on ZINC We evaluate test predictive performance based on validation performance. In the case of the WEBKB datasets, we employ 10-fold cross-validation with the provided data splits. For PEPTIDES, OGBG-MOLIIV, ALCIEMY, and ZINC, our models use positional and structural embeddings concatenated to the initial node features. Specifically, add both RWSE and LAPPE (Dwivedi al. 2022a). We use the same downstream model as the base model for the rewiring models.
Our code can be accessed at
con/ http:/u...oommeeeniiiiaaaaaaaaae MPNN/
G ADDITIONAL EXPERIMENTAL RESULTS
Here, report the computation times of different variants of our probabilistic graph rewiring schemes and results on synthetic datasets.
Training times We report the average training time per epoch in Table The RANDOM entry refers to using random adjacency matrices as rewired graphs.
Extended TUDatasets In addition to Table
4 in the main paper, report the results of IMDB-B and
IMDB-M datasets in Table We also propose proper train/validation/test splitting and show the results in Table
21

---

Published as a conference paper at ICLR 2024
Table 6: Extended results between our probabilistic rewiring method and the other approaches reported in
Giusti et al. (2023b). Besides the 10-fold cross-validation, as in Giusti et al. (2023b); Xu et al. (2019), we provide train/validationneest split. addition, we also provide results on the IMDB-B and IMDB-M datasets. We use green for the best model, blue for the second-best and red for third.
MODEL
MUTAG
PTC.MR
PROTEINS
NCII
NCI109
IMDB-B
IMDB-M
PR-MPNN (10-FOLD CV)
98.4
74.3
80.7
85.6105
84,6112
75.2
52.9-32
PR-MPNN (TRAIN/VAL/TEST)
91.0+3.7
58.9+5.0
79.1
81.5+16
81.8+15
71.6
45.8
Table 7:
Extended comparison on the IMDB-B and IMDB-M M datasets from the TUDATASET collection. We use green for the best model, blue for the second-best and red for third.
MODEL
IMDB-B
IMDB-M
DGCNN ZHANG ET AL.
2018
70.0
47.8
IGN MARON ET AL. 2019B)
71.324.5
48.6+3.9
GIN EX AL.
2019
75.1
52.3+2.8
PPGNS (MARON ET AL..
2019A
73.0
50.413.6
NATURAL GN (DE HAAN ET AL..
2020
74.8
51.211.
GSN (BOURITSAS ET AL.
2022
77.8=3.3
54.3+3..
CIN (BODNAR ET AL., 2021
75.613.2
52.51:.0
PR-MPNN (10-FOLD CV) PR-MPNN (TRAIN/VAL/TEST)
75.2+3.2
52.9+3.2
71.621.2
45.8+0..8
Q.99 We compare our PR-MPNN with multiple current methods QM9 dataset, see Table 10 The baselines are R-GNN in Alon & Yahav (202T), FILM (Brockschmidt, 2020). SPN (Abboud et al.] 2022) and the recent D Rew paper (Gutteridge et al. 2023 Following the settings of Abboud et al. (2022) and Gutteridge et al. (023). we train the network on each task separately. We use the normalized regression labels for training and report the de-nommalized numbers. Similar to Abboud al. (2022) and Gutteridge et al. (023), we also exclude the 3D coordinates of the datasets. It is worth noting that our PR-MPNN reaches the overall lowest mean absolute error on HOMO, LUMO, gap, and Omega tasks while gaining at most 14.13x better performance compared with base GIN model.
WebKB To show PR-MPNNS's 's capability on heterophilic graphs, we carry out experiments on the three WebKB datasets, namely CORNELL, TEXAS, and WISCONSIN, Table [T] We compare with diffusion-based GNN (Gasteiger al. 2019). Geom-GCN etl. 2020p, and the recent graph rewiring work SDRF (Topping et al. 22021. Besides the MPNN baselines above, we also compare them against graph transformers. PR-MPNNS consistently outperform the other MPNN methods and are even better than graph transformers on the TEXAS dataset.
LRGB We apply PR-MPNNs on the two Long Range Graph Benchmark tasks (Dwivedi et al. 2022b). PEPTIDES-FUNE and PEPTIDES- SEPTIDES-STRUCT, which are graph classification and regression tasks, respectively. The baseline methods are also reported in Gutteridge et al. (2023). Notably, on PEPTIDES- -STRUCT, PR-MPNNs reach the overall lowest mean absolute error.
H ROBUSTNESS ANALYSIS
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 parametrized "drop/add" distribution rather than a uniform one.
To observe the performance degradation when testing on noisy data, we conduct 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 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 Table
22

---

Published as a conference paper at ICLR 2024
59)
SAMILLIS
12.51
11.5
11.5
10)
30 40
18.5 10)
15 20
(1.5 10
().5 10,10 25
11.5
11.31
11.
IMPOPOUU
DUSTA
DISTANCE
DISTANCH
DISSANCE
BISTANC
BISTANCE
SIMILARIT
SIMILARIT
SEMILARIH
INETANCE
DISTANCE
ALL
256
235
235
230
230
250
235
4096
4006
4096
550
Table 8: Overview of used hyperparameters.
(L.5 10, 20 25 50, 100)
25
10,25)
120)
(256. S12)
(256, 512
(256. 5121
ALL
(1.5 10, 15 26, 25, 50, 100)
(20. 32. 64, 128, 2363
5 10044 R068)
11024 2048
11024, 2048)
5
00 5
00 5
LAY Satccue
(4,8 12)
(2.4-)
(2.3.4)
LAW LAY
(2.
951
64 95
98
96
64, 96
68 96)
64,
192
192
192
22
52
128
16
22
52,
52,
52,
(32
152
H1001 Mar.0Noerrean
96
164 256
164 256
331
164,26
32
52
16
64
Tisss Teess-Natoon ANOOMBORSHATCH
Tiass LEAPCOUNT
Permoss-nue
P1r11055-11
MUTAO
IMDI-M
IMDI-8
C
WICCONSNN
TOKAS
ZoNc
XC11
CSL
ExP
23

---

Published as a conference paper at ICLR 2024
Table 9: Train and validation time per epoch in seconds for GINE model, the SAT Graph Transformer, and PR-MPNN using different gradient estimators on OGBG-MOLHIV. The time is averaged over five epochs. PR-MPNN is approximately times slower than a GINE model with similar parameter count, while SAT is approximately times slower than the GINE, and times slower than the PR-MPNN models. Experiments performed on machine with single Nvidia RTX A5000 GPU and Intel 19-1K CPU.
MODEL
#PARAMS
TOTAL SAMPLED EDGES
TRAIN TIME/EP (8)
VAL. TIME/EP (S)
GINE
502k
3.19 #0.03
0.20 =0.01
K-ST SATGINE
506k
86.54 10.13
4.78 =0.01
K-SG SATGINE
481k
97.94 10.31
5.57 =0.01
K-ST SATPNA
534k
90.34 =0.29
4.85 =0.01
K-SG SATPNA
509k
118.75 =0.50
5.84 =0.04
PR-MPNN
582k
20
15.20 10.08
1.01 10.01
PR-MPE TR-MPNN Grab
582k
100
18.18 =0.08
1.08 =0.01
PR-MPNNN Imie
5824
15.01 =0.22
1.08 =0.06
PR PR-MPNN Imle
5824
100
15.13 #0.17
1.13 =0.66
PR-MPNN Sim
582k
20
15.98 10.13
1.07 20.01
PR-MPNN 5im
582k
100
17.23 10.15
1.15 20.01
Table 10: Performance of PR-MPNN QM9, in comparison with the base downstrem model (Base-GIN) and other competing methods. The relative improvement of PR-MPNN over the base downstream model is reported in the paranthesis. The metric used is MAE, lower scores better. We note the best performing method with green, the second- best with blue, and third with orange.
PROPERTY
R-GIN+PA
GNN-FILM
SPN
DRew-GIN
BASE-GIN
PR-MPNN
MU
2.54+0.08
2.38
2.32
1.93.0.06
2.64+0.01
1.99
(1.33x)
ALPHA
2.28
3.75
1.77.0.09
1.63
10.16
2.28
(3.36x)
HOMO
1.2640.02
1.22
1.26
1.1620.01
1.70
1.14 (1.49x)
LUMO
1.34
1.30
1.19
1.13
3.05
1.1240.01 (2.72x)
1.9644.04
1.96+0.08
1.8940.11
1.74=0.02
3.37
1.70#a.01 (L.8x)
R2
12.6110.37
15.59
10.66:2.00
9.3920.13
23.35
10.411:035 (2.24x)
ZPVE
5.03+t0.6
11.00
2.77
2.73=20.19
66.87+1.45
4.73=0.08 (14.13x)
UO
2.21:0.12
5.43+40.6
1.12-t0.13
1.01
21.48
2.23 13 (9.38x)
U
2.32:10.18
5.95
1.03
0.99
21.59
2.314:.00 06.00 (9.35x)
H
2.26:0.19
5.59=
1.05
1.0620.08
21.9621.24
2.66 #0.01 (8.26x)
G
2.04
5.17#1113
0.97400.8
1.06
19.53
2.24 (8.24x)
CV
1.8622003
3.4620.21
1.3640.06
1.24=0.02
7.34
1.44 (5.10x)
OMEGA
0.80+1.04
0.98
0.5740.04
0.55=00.1
0.60110.3
1.48 (1.25x)
24

---

Published as a conference paper at ICLR 2024
Table 11: Quantitative results on the heterophilic and transductive WEBKB datasets. Best overall; Second best; Third best. Rewiring outperforms the base models all of the datasets. Graph transformers have an advantage over both the base models and the ones employing rewiring.
HETEROPHILIC & TRANSDUCTIVE
CORNELL. 1
TEXAS
WISCONSIN
BASE
0.57440.006
0.674=0.010
0.697=0.013
BASE W. PE
0.540=00043
0.654=0.010
0.649
DIGL (AATTEGEE ET AT..
2019
0.000s
0.620=0.003
0.095
MPNNS
DIGL UNDIRECTED GASTEIGER ET AI..
2019)
0.595
0.635=0.004
0.022
GEOM-GCN (PEI
2020)
0.608
0.676+ N/A
0.641
SDRF (TOPPING AL.
202T
0.546=1.000
0.044
0.555=00000
SDRF + UNDIRECTED TOPPING ET AL.
2021
0.575
0.003
0.615 100008
PR-MPNN (OURS)
0.659:10.040
0.827 0.032
0.750
GPS (LAPPE)
0.662 0.008
0.010
0.747 100009
GPS (RWSE)
0.708.10.00
0.775
0.802.0.022
GTs
GPS (DEG)
0.718.10.024
0.013
0.798
GRAPHORMER (DEG)
0.683
0.767=0.017
0.770=0.019
GRAPHORMER (DEG + ATTN BIAS)
0.683=20.017
0.767=0.017
0.770=20.019
Table 12: Comparison between PR-MPNN and other methods as reported in Gutteridge et al. (2023).
Best overall; Second best; Third best. PR-MPNN obtains the best score the PEPTIDES-STRUS dataset from the LRGB collection, but ranks below Drew on the PEPTIDES-FUNE dataset.
PEPTIDES-FUNC AP
PEPTIDES-STRUCT MAE +
MODEL
GCN
0.5930-0.0023
0.34961000011
GINE
0.5498=0.0079
0.3547 00.0045
GATEDGCN
0.5864=0.0077
0.3420+40.011
GE
0.6069
0.3357 #0.0006
DIGL+MPNN
0.6469 10.0019
0.3173 10.0007
DIGL:MPNN+LAPPE
0.6830=0.000
0.26 1 6+0.00018
MIXHOP-GCN
0.6592
0.292 1 +0 0023
MIXHOP-GCN+LAPPE
0.6843 0.0049
0.2614:10.000
TRANSFORMER+LAPPE
0.6326-200123
0.2529
SAN+LAPPE
0.638410.011
0.2683
GRAPHGPS+LAPPE
0.6535 00.0444
0.2500:10.0005
DREW-GCN
0.699610.00m
0.278 1 110.0008
DREW-GCN+LAPEE
0.7150 001150=00000 20.0044
0.2536 100000
DREW-GIN
0.6940 00.0004
0.2799+0.00
CEEW-GIN+LAPPE
0.7126 2 0.0045
0.260611.0001
0.6733 0.0094
0.2699 10.0018
DREW-GATEDGON-LAPE
0.6977=0.00m
0.2539+0.0007
PR-MPNN
0.6825:2..0mm
0.2477:100005
25

---

Published as a conference paper at ICLR 2024
Table 13: Robustness results on the PROTEINS dataset when testing on various levels of noise, obtained by removing or adding random edges. The percentages indicate the change in the average test accuracy over 5 runs. PR-MPNNs consistently obtain the best results for both removing 25 and 50 edges.
MODEL
GINE
PR-MPNN
PR-MPNN A-0 50
NOISE
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%
26


Reviews:
[{"review_id": "27TVZ7Zit8", "reviewer": "Reviewer_MyQT", "summary": "The paper has proposed a probabilistic rewiring methanism for GNN. Its improvement in expressive power has been verified with theory and experiments.", "questions": "I wonder what is the performance gain over the overhead the method produces.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The method has been verified with both theory and experiments.", "weaknesses": "1. The computation overhead is also needed.\n2. How does the expressiveness guarantee translate to practical performances?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fmuZWo0zuB", "reviewer": "Reviewer_8vRc", "summary": "The paper proposes a probabilistic rewired message-passing network (PR-MPNN) to address the under-reaching and over-squashing problems of existing graph neural network (GNN) models. Specifically, PR-MPNN first uses a GNN to learn the priors over edges, and then it samples multiple adjacency matrices from the edge prior distributions using SIMPLE, a gradient estimator for k-subset sampling. The sampled adjacency matrices are combined with the original adjacency matrix to obtain the rewired graph, which is used in the downstream tasks. The paper also provides theoretical analysis to identify conditions under which the proposed method can outperform randomized approaches.", "questions": "Please see the questions in the Weaknesses section.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The proposed PR-MPNN model is simple yet effective. \n2. The paper provides theoretical analyses to prove that the proposed model is more effective in probabilistically separating graphs than randomized approaches such as dropping nodes and edges uniformly under certain conditions. \n3. Experimental results on both node classification and graph classification tasks indicate that PR-MPNN can achieve better or competitive performance compared with baselines.", "weaknesses": "1. The motivation for using $k$-subset constraint when sampling the adjacency matrix is not very clear. What are the advantages of using such constraints? \n\n2. The difference between the proposed PR-MPNN and previous works is not explicitly discussed in the related work section. \n\n3. 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?\n\n4. 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. \n\n5. PR-MPNN use 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. \n\n6. I also have some concerns regarding the experiments: \n\n(1) 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. \n\n(2) When answering Q2, why compare PR-MPNN with different baselines on different datasets? For example, the paper compares PR-MPNN with OSAN, GIN+POSENC, 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 achieve comparable performance with DropGNN in some cases. \n\n(3) 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?\n\n(4) There are no ablation studies to validate the effectiveness of each component of the proposed model. \n\n7. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4OHQhVn4ys", "reviewer": "Reviewer_j89x", "summary": "The work proposes a probabilistic rewiring technique that relies on differentiable k-subset sampling. The main motivation is that while current rewiring techniques rely on \"arbitrary\" heuristics (improving spectral gap, connecting nodes in the 2-hop), it would be ideal to *learn* such a rewiring. Further, the work proposes a way to study the expressive power -- in terms of WL -- for such a process. The work is evaluated on synthetic and real-world benchmarks.", "questions": "(Q1) 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? \n\n(Q2) Regarding (W1), would it be possible to show results for existing deterministic rewiring techniques?\n\n(Q3) Regarding (W2) and especially (Q1), it is not clear to me how the technique can reduce over-squashing if it is only able to remove edges. It would be important to clarify (Q1), and provide some theoretical evidence that it is indeed able to alleviate over-squashing. In general by removing edges, one is reducing the total effective resistance over the graph.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The work is well-written and overall I found it interesting. I agree that learning a rewiring is intuitively something that should be preferred over arbitrary heuristics. I also appreciated that the work discusses techniques such as the Differentiable Graph Module, which in principle may seem similar.", "weaknesses": "(W1) 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. \n\n(W2) The claim that the rewiring technique reduces over-squashing could be strengthened. At the moment this seems to be solely motivated by the empirical results in Figure 2 and Figure 4.\n\n[1] FoSR: First-order spectral rewiring for addressing oversquashing in GNNs. Kedar Karhadkar, Pradeep Kr. Banerjee, Guido Montúfar. ICLR 2023\n\n[2] Jake Topping, Francesco Di Giovanni, Benjamin Paul Chamberlain, Xiaowen Dong, and Michael M Bronstein. Understanding over-squashing and bottlenecks on graphs via curvature. ICLR 2022", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)