Probabilistically Rewired Message-Passing Neural Networks
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.
Chendi Qian, Andrei Manolache, Kareem Ahmed, Zhe Zeng, Guy Van den Broeck, Mathias Niepert, Christopher Morris
# Probabilistically Rewired Message-Passing Neural Networks

###### 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.

## 1 Introduction

Graph-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 social-network 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, 2017; Gilmer et al., 2017; Scarselli et al., 2008; Velickovic 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 et al., 2019; Xu et al., 2019; Morris et al., 2021). In addition, MPNNs struggle to capture global or long-range information, possibly leading
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 to enhance message passing and connectivity, while Karhadkar et al. (2022) resort to spectral techniques, and Banerjee et al. (2022) proposed a greedy random edge flip approach to 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-neighbor-based message passing to maintain long-range dependencies (Abboud et al., 2022; Abu-El-Haija et al., 2019; Gasteiger et al., 2019; Xue et 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 a model.

In contrast to the above lines of work, graph transformers (Chen et al., 2022; Dwivedi et al., 2022; He et al., 2023; Muller et al., 2023; Ramgasek et al., 2022) and similar global attention mechanisms (Liu et al., 2021; Wu et al., 2021) marked a shift from local to global message passing, aggregating over all nodes. While not understood in a principled way, empirical studies indicate that graph transformers possibly alleviate over-squashing; see, e.g., Muller et al. (2023). However, due to their global aggregation mode, computing an attention matrix with \(n^{2}\) entries for an \(n\)-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 over-squashing rely on heuristic rewiring methods that may not adapt well to a prediction task or employ computationally intensive global attention mechanisms. Furthermore, the impact of probabilistic rewiring on a model's expressive power remains unclear.

Figure 1: Overview of the probabilistically rewired MPNN framework. PR-MPNNs use an _upstream model_ to learn priors \(\mathbf{\theta}\) for candidate edges, parameterizing a probability mass function conditioned on exactly-\(k\) constraints. Subsequently, we sample multiple \(k\)-edge adjacency matrices (here: \(k=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 \(\ell\) with respect to the parameters \(\mathbf{\theta}\) are approximated through the derivative of the exactly-\(k\) marginals in the direction of the gradients of the point-wise loss \(\ell\) with respect to the sampled adjacency matrix. We use recent work to make the computation of these marginals exact and differentiable, reducing both bias and variance.
**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 a probability distribution constrained by so-called \(k\)-subset constraints. Subsequently, we sample multiple \(k\)-edge adjacency matrices from this distribution and process them using a _downstream model_, typically an 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 (Xie and Ermon, 2019; Niepert et al., 2021; Ahmed et al., 2023); see Figure 1 for an overview. Our theoretical analysis explores how PR-MPNNs overcome MPNNs' inherent limitations in expressive power and identifies 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 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._

### Related Work

MPNNs are inherently biased towards encoding local structures, limiting their expressive power (Morris et al., 2019; Xu et al., 2019; Morris et al., 2021). Specifically, they are at most as powerful as distinguishing non-isomorphic graphs or nodes with different structural roles as the _1-dimensional Weisfeiler-Leman algorithm_(Weisfeiler and Leman, 1968), a simple heuristic for the graph isomorphism problem; see Section 2. Additionally, they cannot capture global or long-range information, often linked to phenomena such as under-reaching (Barcelo et al., 2020) or over-squashing (Alon and 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 neighbors. For example, Bruel-Gabrielsson et al. (2022) rewires the graphs with \(k\)-hop neighbors and virtual nodes and also augments them with positional encodings. MixHop (Abu-El-Haija et al., 2019), 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-aware neighbors in a single layer. Particularly, Gutteridge et 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), balanced Forman curvature (Topping et al., 2021). In addition, Deac et al. (2022); Shirzad et al. (2023) utilize expand graphs to enhance message passing and connectivity, while Karhadkar et al. (2022) resort to spectral techniques, and Banerjee et al. (2022) propose a greedy random edge flip approach to overcome over-squashing. 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. Overall, current strategies to mitigate over-squashing either rely on heuristic rewiring methods or even purely randomized approaches that may not adapt well to a given prediction task. Furthermore, the impact of probabilistic rewiring on a model's expressive power remains unclear.

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; Muller et al., 2023; Rampasek et al., 2022; Chen et al., 2022) and similar global attention mechanisms (Liu et al., 2021; Wu et al., 2021) marked a shift from local to global message passing. aggregating over all nodes. While not understood in a principled way, empirical studies indicate that graph transformers possibly alleviate over-squashing; see, e.g., Muller 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 \(\mathbb{N}\coloneqq\{1,2,3,\dots\}\). For \(n\geq 1\), let \([n]\coloneqq\{1,\dots,n\}\subset\mathbb{N}\). We use \(\llbracket\dots\rrbracket\) to denote multisets, i.e., the generalization of sets allowing for multiple instances for each of its elements. A _graph_\(G\) is a pair \((V(G),E(G))\) with _finite_ sets of _vertices_ or _nodes_\(V(G)\) and _edges_\(E(G)\subseteq\{\{u,v\}\subseteq V(G)\mid u\neq v\}\). If not otherwise stated, we set \(n\coloneqq|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,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 \(N(v)\) and \(\ell(v)\) denotes its discrete vertex label, and so on; see Appendix B for details.

\(1\)**-dimensional Weisfeiler-Leman algorithm** The \(1\)-WL or color refinement is a well-studied heuristic for the graph isomorphism problem, originally proposed by Weisfeiler & Leman (1968). Formally, let \(G=(V(G),E(G),\ell)\) be a labeled graph. In each iteration, \(t>0\), the \(1\)-WL computes a node coloring \(C^{1}_{t}\colon V(G)\to\mathbb{N}\), depending on the coloring of the neighbors. That is, in iteration \(t>0\), we set

\[C^{1}_{t}(v)\coloneqq\mathsf{RELABEL}\Big{(}\!\big{(}C^{1}_{t-1}(v), \llbracket\!\,\{C^{1}_{t-1}(u)\mid u\in N(v)\}\!\,\big{)}\!\,\Big{)},\]

for all nodes \(v\in V(G)\), where \(\mathsf{RELABEL}\) injectively maps the above pair to a unique natural number, which has not been used in previous iterations. In iteration \(0\), the coloring \(C^{1}_{0}\coloneqq\ell\). To test if two graphs \(G\) and \(H\) are non-isomorphic, we run the above algorithm in "parallel" on both graphs. If the two graphs have a different number of nodes colored \(c\) in \(\mathbb{N}\) at some iteration, the \(1\)-WL _distinguishes_ the graphs as non-isomorphic. Moreover, if the number of colors between two iterations, \(t\) and \((t+1)\), does not change, i.e., the cardinalities of the images of \(C^{1}_{t}\) and \(C^{1}_{i+t}\) are equal, or, equivalently,

\[C^{1}_{t}(v)=C^{1}_{t}(w)\iff C^{1}_{t+1}(v)=C^{1}_{t+1}(w),\]

for all nodes \(v\) and \(w\) in \(V(G)\), the algorithm terminates. For such \(t\), we define the _stable coloring_\(C^{1}_{\infty}(v)=C^{1}_{t}(v)\), for \(v\) in \(V(G)\). The stable coloring is reached after at most \(\max\{|V(G)|,|V(H)|\}\) iterations (Grobe, 2017). It is easy to see 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\colon V(G)\to\mathbb{R}^{d}\), for \(d>0\), is \(1\)-WL-_equivalent_ if \(f\equiv C^{1}_{\infty}\); see Appendix B for details.

**Message-passing graph neural networks** Intuitively, MPNNs learn a vectorial representation, i.e., a \(d\)-dimensional real-valued vector, representing each vertex in a graph by aggregating information from neighboring vertices. Let \(\mathbf{G}=(G,\mathbf{L})\) be an attributed graph, following, Gilmer et al. (2017) and Scarselli et al. (2008a), in each layer, \(t>0\), we compute vertex features

\[\mathbf{h}^{(t)}_{v}\coloneqq\mathsf{UPD}^{(t)}\Big{(}\mathbf{h}^{(t-1)}_{v},\mathsf{AGG}^{(t)}\big{(}\!\big{\{}\mathbf{h}^{(t-1)}_{u}\mid u\in N(v) \}\!\,\big{)}\!\,\Big{)}\Big{)}\in\mathbb{R}^{d},\]

where \(\mathsf{UPD}^{(t)}\) and \(\mathsf{AGG}^{(t)}\) may be differentiable parameterized functions, e.g., neural networks, and \(\mathbf{h}^{(t)}_{v}=\mathbf{L}_{v}\). In the case of graph-level tasks, e.g., graph classification, one uses

\[\mathbf{h}_{G}\coloneqq\mathsf{READOUT}\big{(}\!\big{\{}\mathbf{h}^{(T)}_{v} \mid v\in V(G)\!\,\big{\}}\!\big{)}\in\mathbb{R}^{d},\]

to compute a single vectorial representation based on learned vertex features after iteration \(T\). Again, \(\mathsf{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 a variant of stochastic gradient descent, e.g., Kingma & Ba (2015), together with the parameters of a neural network used for classification or regression.

## 3 Probalistically rewired MPNNs

Here, we outline probabilistically rewired MPNNs (PR-MPNNs) based on recent advancements in discrete gradient estimation and tractable sampling techniques (Ahmed et al., 2023). Let \(\mathfrak{A}_{n}\) denote the set of adjacency matrices of \(n\)-order graphs. Further, let \((G,\mathbf{X})\) be a \(n\)-order attributed graph with an adjacency matrix \(\mathbf{A}(G)\in\mathfrak{A}_{n}\), and node attribute matrix \(\mathbf{X}\in\mathbb{R}^{n\times d}\), for \(d>0\). A PR-MPNN maintains a (parameterized) _upstream model_\(h_{\mathbf{u}}\colon\mathfrak{A}_{n}\times\mathbb{R}^{n\times d}\to\Theta\), typically a neural network,
parameterized by \(\mathbf{u}\), mapping an adjacency matrix and corresponding node attributes to unnormalized edge priors \(\mathbf{\theta}\in\mathbf{\Theta}\subseteq\mathbb{R}^{n\times n}\).

In the following, we use the _priors_\(\mathbf{\theta}\) as parameters of a (conditional) probability mass function,

\[p_{\mathbf{\theta}}(\mathbf{A}(H))\coloneqq\prod_{i,j=1}^{n}p_{\theta_{ij}}( \mathbf{A}(H)_{ij}),\]

assigning a probability to each adjacency matrix in \(\mathfrak{A}_{n}\), where \(p_{\theta_{ij}}(\mathbf{A}(H)_{ij}=1)=\mathrm{sigmoid}(\theta_{ij})\) and \(p_{\theta_{ij}}(\mathbf{A}(H)_{ij}=0)=1-\mathrm{sigmoid}(\theta_{ij})\). Since the parameters \(\mathbf{\theta}\) depend on the input graph \(G\), we can view the above probability as a conditional probability mass function conditioned on the graph \(G\).

Unlike previous probabilistic rewiring approaches, e.g., Franceschi et al. (2019), we introduce dependencies between the graph's edges by conditioning the probability mass function \(p_{\theta_{ij}}(\mathbf{A}(H))\) on a _\(k\)-subset constraint_. That is, the probability of sampling any given \(k\)-edge adjacency matrix \(\mathbf{A}(H)\), becomes

\[p_{(\mathbf{\theta},k)}(\mathbf{A}(H))\coloneqq\left\{\begin{array}{ll}p_{\mathbf{ \theta}}(\mathbf{A}(H))/Z&\text{if }\|\mathbf{A}(H)\|_{1}=k,\\ 0&\text{otherwise},\end{array}\right.\text{ with }\quad Z\coloneqq\sum_{ \mathbf{B}\in\mathfrak{A}_{n}\colon\|\mathbf{B}\|_{1}=k}p_{\mathbf{\theta}}( \mathbf{B}). \tag{1}\]

The original graph \(G\) is now rewired into a new adjacency matrix \(\bar{\mathbf{A}}\) by combining \(N\) samples \(\mathbf{A}^{(i)}\sim p_{(\mathbf{\theta},k)}(\mathbf{A}(G))\) for \(i\in[N]\) together with the original adjacency matrix \(\mathbf{A}(G)\) using a differentiable aggregation function \(g\colon\mathfrak{A}_{n}^{(N+1)}\to\mathfrak{A}_{n}\), i.e., \(\bar{\mathbf{A}}\coloneqq g(\mathbf{A}(G),\mathbf{A}^{(1)},\ldots,\mathbf{A}^{( N)})\in\mathfrak{A}_{n}\). Subsequently, we use the resulting adjacency matrix as input to a _downstream model_\(f_{\mathbf{d}}\), parameterized by \(\mathbf{d}\), typically an MPNN, for the final predictions task.

We have so far assumed that the upstream MPNN computes one set of priors \(h_{\mathbf{u}}\colon\mathfrak{A}_{n}\times\mathbb{R}^{n\times d}\to\mathbb{R} ^{n\times n}\) which we use to generate a new adjacency matrix \(\bar{\mathbf{A}}\) through sampling and then aggregating the adjacency matrices \(\mathbf{A}^{(1)},\ldots,\mathbf{A}^{(N)}\). In Section 5, we show empirically that having multiple sets of priors from which we sample is beneficial. Multiple sets of priors mean that we learn an upstream model \(h_{\mathbf{u}}\colon\mathfrak{A}_{n}\times\mathbb{R}^{n\times d}\to\mathbb{R} ^{n\times n\times 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 \(\mathbf{\omega}=(\mathbf{u},\mathbf{d})\) of the PR-MPNN architecture, we minimize the expected loss

\[L(\mathbf{A}(G),\mathbf{X},y;\mathbf{\omega})\coloneqq\mathbb{E}_{\mathbf{A}^{(i )}\sim p_{(\mathbf{\theta},k)}(\mathbf{A}(G))}\Big{[}\ell\Big{(}f_{\mathbf{d}}\Big{(} g\Big{(}\mathbf{A}(G),\mathbf{A}^{(1)},\ldots,\mathbf{A}^{(N)}\Big{)}, \mathbf{X}\Big{)},y\Big{)}\Big{]},\]

with \(y\in\mathcal{Y}\), the targets, \(\ell\) a point-wise loss such as the cross-entropy or MSE, and \(\mathbf{\theta}=h_{\mathbf{u}}(\mathbf{A}(G),\mathbf{X})\). To minimize the above expectation using gradient descent and backpropagation, we need to efficiently draw Monte-Carlo samples from \(p_{(\mathbf{\theta},k)}(\mathbf{A}(G))\) and estimate \(\nabla_{\mathbf{\theta}}L\) the gradients of an expectation with respect to the parameters \(\mathbf{\theta}\) of the distribution \(p_{(\mathbf{\theta},k)}\).

**Sampling** To sample an adjacency matrix \(\mathbf{A}^{(i)}\) from \(p_{(\mathbf{\theta},k)}(\mathbf{A}(G))\) conditioned on \(k\)-edge constraints, and to allow PR-MPNNs to be trained end-to-end, we use Simple (Ahmed et al., 2023), a recently proposed gradient estimator. Concretely, we can use Simple to sample _exactly_ from the \(k\)-edge adjacency matrix distribution \(p_{(\mathbf{\theta},k)}(\mathbf{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) with respect to the prior weights \(\mathbf{\theta}\) using

\[\nabla_{\mathbf{\theta}}L\approx\partial_{\mathbf{\theta}}\mu(\mathbf{\theta})\nabla_{ \mathbf{A}}\ell\ \ \text{with}\ \ \mu(\mathbf{\theta})\coloneqq\{p_{(\mathbf{\theta},k)}(\mathbf{A}(G)_{ij})\}_{i,j=1}^{n} \in\mathbb{R}^{n\times n},\]

with an exact and efficient computation of the marginals \(\mu(\mathbf{\theta})\) that is differentiable on the backward pass, achieving lower bias and variance. We show empirically that Simple (Ahmed et al., 2023) outperforms other sampling and gradient approximation methods such as Gumbel SoftSub-ST (Xie & Ermon, 2019) and I-MLE (Niepert et al., 2021), improving accuracy without incurring a computational overhead.

**Computational complexity** The vectorized complexity of the exact sampling and marginal inference step is \(\mathcal{O}(\log k\log l)\), where \(k\) is from our \(k\)-subset constraint, and \(l\) is the maximum number of edges that we can sample. Assuming a constant number of layers, PR-MPNN's worst-case training complexity is \(\mathcal{O}(l)\) for both the upstream and downstream models. Let \(n\) be the number of nodes in the initial graph, and \(l=\max(\{l_{\text{add}},l_{\text{m}}\})\), with \(l_{\text{add}}\) and \(l_{\text{m}}\) being the maximum number of added and deleted edges. If
we consider all of the possible edges for \(l_{\text{add}}\), the worst-case complexity becomes \(\mathcal{O}(n^{2})\). Therefore, to reduce the complexity in practice, we select a subset of the possible edges using simple heuristics, such as considering the top \(l_{\text{add}}\) 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 \(\mathcal{O}(l)\) for the upstream model and \(\mathcal{O}(L)\) for the downstream model, with \(L\) being the number of edges in the rewired graph.

## 4 Expressive Power of Probabilistically Rewired MPNNs

We now, for the first time, explore the extent to 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 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\colon\mathfrak{A}_{n}\to[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 function over graphs. Moreover, let \(f\colon\mathfrak{A}_{n}\to\mathbb{R}^{d}\), for \(d>0\), be a permutation-invariant, parameterized function mapping a sampled graph's adjacency matrix to a vector in \(\mathbb{R}^{d}\). The function \(f\) could be the composition of an aggregation function \(g\) that removes the sampled edges from the input graph \(G\) and of a downstream MPNN. Now, the conditional probability mass function \(p\)_separates_ two graphs \(G\) and \(H\) with probability \(\rho\)_with respect to_\(f\) if

\[\mathbb{E}_{\widetilde{G}\sim p(\cdot|G),\widetilde{H}\sim p(\cdot|H)}\big{[} f(\mathbf{A}(\widetilde{G}))\neq f(\mathbf{A}(\widetilde{H}))\big{]}=\rho,\]

that is, if in expectation over the conditional probability distribution, the vectors \(f(\mathbf{A}(\widetilde{G}))\) and \(f(\mathbf{A}(\widetilde{H}))\) are distinct with probability \(\rho\).

In what follows, we analyze the case of \(p\) being the exactly-\(k\) probability distribution defined in Equation 1 and \(f\) being the aggregation function removing edges and a 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 \(k\) edges and use them to remove edges from the original graph. We aim to understand the separation properties of the probability mass function \(p_{(k,\mathbf{\theta})}\) in this setting and for various types of graph structures. Most obviously, we do not want to separate isomorphic graphs and, therefore, remain isomorphism invariant, a desirable property of MPNNs.

**Theorem 4.1**.: _For sufficiently large \(n\), for every \(\varepsilon\in(0,1)\) and \(k>0\), we have for almost all pairs, in the sense of Babai et al. (1980), of isomorphic \(n\)-order graphs \(G\) and \(H\) and all permutation-invariant, \(1\)-WL-equivalent functions \(f\colon\mathfrak{A}_{n}\to\mathbb{R}^{d}\), \(d>0\), there exists a probability mass function \(p_{(\mathbf{\theta},k)}\) that separates the graph \(G\) and \(H\) with probability at most \(\varepsilon\) with respect to \(f\)._

Theorem 4.1 relies on the fact that most graphs have a discrete \(1\)-WL coloring. For graphs where the \(1\)-WL stable coloring consists of a discrete and non-discrete part, the following result shows that there exist distributions \(p_{(\mathbf{\theta},k)}\) not separating the graphs based on the partial isomorphism corresponding to the discrete coloring.

**Proposition 4.2**.: _Let \(\varepsilon\in(0,1)\), \(k>0\), and let \(G\) and \(H\) be graphs with identical \(1\)-WL stable colorings. Let \(V_{G}\) and \(V_{H}\) be the subset of nodes of \(G\) and \(H\) that are in color classes of cardinality \(1\). Then, for all choices of \(1\)-WL-equivalent functions \(f\), there exists a conditional probability distribution \(p_{(\mathbf{\theta},k)}\) that separates the graphs \(G[V_{G}]\) and \(H[V_{H}]\) with probability at most \(\varepsilon\) with respect to \(f\)._

Existing methods such as DropGNN (Papp et al., 2021) or DropEdge (Rong et al., 2020) 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-invariance 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 non-isomorphic graphs than probability mass functions that remove edges or nodes uniformly at random.

**Theorem 4.3**.: _For every \(\varepsilon\in(0,1)\) and every \(k>0\), 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 \(f\),_

1. _there exists a probability mass function_ \(p_{(k,\boldsymbol{\theta})}\) _that separates_ \(G\) _and_ \(H\) _with probability at least_ \((1-\varepsilon)\) _with respect to_ \(f\)_;_
2. _removing edges uniformly at random separates_ \(G\) _and_ \(H\) _with probability at most_ \(\varepsilon\) _with respect to_ \(f\)_._

Finally, we can also show a negative result: there exist classes of graphs for which PR-MPNNs cannot do better than random sampling.

**Proposition 4.4**.: _For every \(k>0\), there exist non-isomorphic graphs \(G\) and \(H\) with identical \(1\)-WL colorings such that every probability mass function \(p_{(\boldsymbol{\theta},k)}\) separates the two graphs with the same probability as the distribution that samples edges uniformly at random._

\begin{table}
\begin{tabular}{l l l l l l} \hline \hline  & & \multicolumn{2}{c}{Zinc} & \multicolumn{2}{c}{OGBG-molity} & \multicolumn{1}{c}{Alchemy} \\  & & - Edge \(\downarrow\) & + Edge \(\downarrow\) & + Edge \(\uparrow\) & + Edge \(\downarrow\) \\ \hline \multirow{4}{*}{\begin{tabular}{l} \(\boldsymbol{\theta}\) \\ \(\boldsymbol{\theta}\) \\ \(\boldsymbol{\theta}\) \\ \end{tabular} } & K-ST SAT & 0.162\(\pm\)0.007 & 0.115\(\pm\)0.005 & 0.625\(\pm\)0.039 & N/A \\  & K-SG SAT & 0.162\(\pm\)0.013 & **0.095\(\pm\)**0.002 & 0.613\(\pm\)0.010 & N/A \\  & Base & 0.258\(\pm\)0.006 & 0.207\(\pm\)0.006 & 0.775\(\pm\)0.011 & 11.12\(\pm\)0.690 \\  & Base w. PE & 0.162\(\pm\)0.001 & 0.101\(\pm\)0.004 & 0.764\(\pm\)0.018 & 7.197\(\pm\)0.094 \\  & PR-MPNN\({}_{\text{\sc{max}}}\) (ours) & 0.153\(\pm\)0.003 & 0.103\(\pm\)0.008 & 0.760\(\pm\)0.025 & 6.858\(\pm\)0.059 \\  & PR-MPNN\({}_{\text{\sc{max}}}\) (ours) & **0.151\(\pm\)**0.001 & 0.104\(\pm\)0.008 & **0.774\(\pm\)**0.015 & 6.692\(\pm\)0.061 \\  & PR-MPNN\({}_{\text{\sc{min}}}\) (ours) & **0.139\(\pm\)**0.001 & **0.085\(\pm\)**0.002 & **0.795\(\pm\)**0.009 & 6.447\(\pm\)**0.057 \\ \hline \multirow{4}{*}{
\begin{tabular}{l} \(\boldsymbol{\theta}\) \\ \(\boldsymbol{\theta}\) \\ \(\boldsymbol{\theta}\) \\ \end{tabular} } & GPS & N/A & **0.070\(\pm\)**0.004 & **0.788\(\pm\)**0.010 & N/A \\  & K-ST SAT & 0.164\(\pm\)0.007 & 0.102\(\pm\)0.005 & 0.613\(\pm\)0.025 & N/A \\  & K-SG SAT & **0.131\(\pm\)**0.002 & **0.094\(\pm\)**0.008 & 0.591\(\pm\)0.005 & N/A \\ \hline \hline \end{tabular}
\end{table}
Table 1: 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., 2021), and Simple (Ahmed et al., 2023) and compare them with the base downstream model, and two Graph Transformers. 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-best, and red for third. We note with + Edge the instances where edge features are provided, and with - Edge when they are not.

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.
## 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.

* Can probabilistic graph rewiring mitigate the problems of over-squashing and under-reaching in synthetic datasets?
* 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 4 can be learned with PR-MPNNs and that multiple prior help?
* Does the increase in predictive performance due to probabilistic rewiring apply to (a) graph-level molecular prediction tasks and (b) node-level prediction tasks involving heterophilic data?

The repository of our code can be accessed at [https://github.com/chendiqian/PR-MPNN](https://github.com/chendiqian/PR-MPNN).

**Datasets** To answer **Q1**, we utilized the Trees-NeighborsMatch dataset (Alon and Yahav, 2021). Additionally, we created the Trees-LeafCount dataset to investigate whether our method could mitigate under-reaching issues; see Appendix D for details. To tackle **Q2**, we performed experiments with the Exp(Abboud et al., 2020) and Csl datasets (Murphy 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); Papp et al. (2021) and set it against a standard DropGNN model (Papp et al., 2021) for comparison while also ablating the performance difference concerning the number of priors and samples per prior. To answer **Q3** (a), we used the established molecular graph-level regression datasets (Chen et al., 2019), Zinc(Jin et al., 2017; Dwivedi et al., 2020), ogbg-molhiy(Hu et al., 2020a), QM9 (Hamilton et al., 2017), LRGB (Dwivedi et al., 2022b) and five datasets from the TDataset repository (Morris et al., 2020). To answer **Q3** (b), we used the Cornell, Wisconsin, Texas node-level classification datasets (Pei et al., 2020).

**Baseline and model configurations** For our upstream model \(h_{\mathbf{u}}\), we use an MPNN, specifically the GIN layer (Xu et al., 2019). For an edge \((v,w)\in E(G)\), we compute \(\mathbf{\theta}_{vw}=\phi([\mathbf{h}_{v}^{T}||\mathbf{h}_{w}^{T}])\in\mathbb{R}\), where \([\cdot||\cdot]\) is the concatenation operator and \(\phi\) is an MLP. After obtaining the prior \(\mathbf{\theta}\), we rewire our graphs by sampling two adjacency matrices for deleting edges and adding new edges, i.e., \(g(\mathbf{A}(G),\mathbf{A}^{(1)},\mathbf{A}^{(2)}):=(\mathbf{A}(G)-\mathbf{A}^ {(1)})+\mathbf{A}^{(2)}\) where \(\mathbf{A}^{(1)}\) and \(\mathbf{A}^{(2)}\) are two sampled adjacency matrices with a possibly different number of edges, respectively. Finally, the rewired adjacency matrix (or multiple adjacency matrices) is used in a _downstream model_\(f_{\mathbf{d}}\colon\mathfrak{A}_{n}\times\mathbb{R}^{n\times d}\to\mathcal{Y}\), typically an MPNN, with parameters \(d\) and \(\mathcal{Y}\) the prediction target set. For the instance where we have multiple priors, as described in Section 3, we can either aggregate the sampled adjacency matrices \(\mathbf{A}^{(1)},\ldots,\mathbf{A}^{(N)}\) into a single adjacency matrix \(\mathbf{\tilde{A}}\) that we send to a downstream model as described in Figure 1, or construct a downstream ensemble with multiple aggregated matrices \(\mathbf{\tilde{A}}_{1},\ldots,\mathbf{\tilde{A}}_{M}\).

All of our downstream models \(f_{\mathbf{d}}\) 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 \(\mathbf{h}_{v}^{T}\) 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 (Rampasek et al., 2022) and SAT (Chen et al., 2022), two state-of-the-art graph transformers. For the TUDataset, we compare with the reported scores from Giusti et al. (2023b) and use the same evaluation strategy as in Xu et al. (2019); Giusti et al. (2023b), i.e., running 10-fold cross-validation 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 8 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 et al., 2017; Jang et al., 2017; Xie and Ermon, 2019), I-MLE (Niepert et al., 2021), and Simple(Ahmed et al., 2023). The results in terms of predictive power are detailed in Table 1, and the computation time comparisons can be found in Table 9 in the appendix. Further experimental results on QM9 and LRGB are included in Appendix F in the appendix.
of the datasets, with the exception being NCI1, 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 11 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.

## 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, on 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 a significant step towards systematically developing more adaptable MPNNs, rendering them less susceptible to potential noise and missing data, thereby enhancing their applicability and robustness.

#### Acknowledgments

CQ and CM are partially funded by a DFG Emmy Noether grant (468502433) and RWTH Junior Principal Investigator Fellowship under Germany's Excellence Strategy. AM and MN acknowledge funding by Deutsche Forschungsgemeinschaft (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 HR00112220005, the DARPA Assured Neuro Symbolic Learning and Reasoning (ANSR) Program, and a gift from RelationalAI. GVdB discloses a financial interest in RelationalAI.