Probabilistic Graph Rewiring via Virtual Nodes

Authors: chendi  qian
andrei  manolache
christopher  morris
mathias  niepert
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 21:36:14

Abstract: Message-passing graph neural networks (MPNNs) have emerged as a powerful paradigm for graph-based machine learning. Despite their effectiveness, MPNNs face challenges such as under-reaching and over-squashing, where limited receptive fields and structural bottlenecks hinder information flow in the graph. While graph transformers hold promise in addressing these issues, their scalability is limited due to quadratic complexity regarding the number of nodes, rendering them impractical for larger graphs. Here, we propose implicitly rewired message-passing neural networks (IPR-MPNNs), a novel approach that integrates implicit probabilistic graph rewiring into MPNNs. By introducing a small number of virtual nodes, i.e., adding additional nodes to a given graph and connecting them to existing nodes, in a differentiable, end-to-end manner, IPR-MPNNs enable long-distance message propagation, circumventing quadratic complexity. Theoretically, we demonstrate that IPR-MPNNs surpass the expressiveness of traditional MPNNs.  Empirically, we validate our approach by showcasing its ability to mitigate under-reaching and over-squashing effects, achieving state-of-the-art performance across multiple graph datasets. Notably, IPR-MPNNs outperform graph transformers while maintaining significantly faster computational efficiency.

Main paper: 
 introduction: 
 Message-passing graph neural networks (MPNNs) [Gilmer et al., 2017 , Scarselli et al., 2008] recently emerged as the most prominent machine-learning architecture for graph-structured, applicable to a large set of domains where data is naturally represented as graphs, such as bioinformatics [Jumper et al., 2021 , Wong et al., 2023] , social network analysis [Easley et al., 2012] , and combinatorial optimization [Cappart et al., 2023 , Qian et al., 2024] .MPNNs have been studied extensively in theory and practice [Böker et al., 2023 , Gilmer et al., 2017 , Kipf and Welling, 2017 , Maron et al., 2019b , Morris et al., 2019 , 2021 , 2023 , Veličković et al., 2018 , Xu et al., 2019] . Recent works have shown that MPNNs suffer from over-squashing [Alon and Yahav, 2021] , where bottlenecks arise from stacking multiple layers leading to large receptive fields, and under-reaching [Barceló et al., 2020] , where distant nodes fail to communicate effectively because MPNNs' receptive fields are too narrow. These phenomena become prevalent when dealing with graphs with a large diameter, potentially hindering the performance of MPNNs on essential applications that depend on long-range interactions, such as protein folding [Gromiha and Selvaraj, 1999] . However, modeling long-range interactions in atomistic systems such as proteins remains a Figure 1 : Overview of how IPR-MPNNs implicitly rewire a graph through adding virtual nodes. IPR-MPNNs use an upstream MPNN to learn priors θ for connecting original nodes with virtual nodes via edges, parameterizing a probability mass function conditioned on exactly-k constraints. Subsequently, we sample exactly k edges from this distribution for each original node, connecting it to k virtual nodes. We input the resulting graph to a downstream model, typically an MPNN, for the final predictions task, propagating information from (1) original nodes to virtual nodes, (2) among virtual nodes, and (3) among original nodes. On the backward pass, the gradients of the loss ℓ regarding the parameters θ are approximated through the derivative of the exactly-k marginals. challenging problem often solved in an ad-hoc fashion using coarse-graining methods [Saunders and Voth, 2013, Husic et al., 2020] , effectively grouping the input nodes into cluster nodes.Recently, graph rewiring [Bober et al., 2022 , Deac et al., 2022 , Gutteridge et al., 2023 , Karhadkar et al., 2022 , Shirzad et al., 2023 , Topping et al., 2021] techniques emerged, adapting the graph structure to enhance connectivity and reduce node distance through methods ranging from edge additions to leveraging spectral properties and expander graphs. However, these approaches typically employ heuristic methods for selecting node pairs to rewire. Furthermore, graph transformers (GTs) [Chen et al., 2022a , Dwivedi et al., 2022b , He et al., 2023 , Müller et al., 2023 , Müller and Morris, 2024 , Rampášek et al., 2022] and adaptive techniques like those by Errica et al. [2023] improve handling of long-range relationships but face challenges of quadratic complexity and extensive parameter sets, limiting their scalability.Most similar to IPR-MPNNs is the work of Qian et al. [2023] , which, like IPR-MPNNs, leverage recent techniques in differentiable k-subset sampling [Ahmed et al., 2023] to learn to add or remove edges of a given graph. However, like GTs, their approach suffers from quadratic complexity due to their need to compute a score for each node pair.Present Work Our proposed IPR-MPNN architecture advances end-to-end probabilistic adaptive graph rewiring. Unlike the PR-MPNN framework of Qian et al. [2023] , which suffers from quadratic complexity since the edge distribution is modeled explicitly, IPR-MPNNs implicitly transmits information across different parts of a graph by learning to connect the existing graph with newly-added virtual nodes, effectively circumventing quadratic complexity; see Figure 1 for a high-level overview of IPR-MPNNs. Our contributions are as follows.1. We introduce IPR-MPNNs, adding virtual nodes to graphs, and learn to rewire them to the existing nodes end-to-end. IPR-MPNNs successfully overcome the quadratic complexity of graph transformers and previous graph rewiring techniques.2. Theoretically, we demonstrate that IPR-MPNNs exceed the expressive capacity of standard MPNNs, typically limited by the 1-dimensional Weisfeiler--Leman algorithm.3. Empirically, we show that IPR-MPNNs outperform standard MPNN and GT architectures on a large set of established benchmark datasets, all while maintaining significantly faster computational efficiency.4. We show that IPR-MPNNs are reducing the total effective resistance [Black et al., 2023] of multiple molecular datasets while also significantly improving the layer-wise sensitivity [Di Giovanni et al., 2023 , Xu et al., 2018] between distant nodes when compared to the base model.In summary, IPR-MPNNs represent a significant advancement towards scalable and adaptable MPNNs. They enhance expressiveness and adaptability to various data distributions while scaling to large graphs effectively.In the following, we discuss relevant related work.MPNNs Recently, MPNNs [Gilmer et al., 2017 , Scarselli et al., 2008] emerged as the most prominent graph machine learning architecture. Notable instances of this architecture include, e.g., Duvenaud et al. [2015] , Hamilton et al. [2017] , and Veličković et al. [2018] , which can be subsumed under the message-passing framework introduced in Gilmer et al. [2017] . In parallel, approaches based on spectral information were introduced in, e.g., Bruna et al. [2014] , Defferrard et al. [2016] , Gama et al. [2019] , Kipf and Welling [2017] , Levie et al. [2019] , and Monti et al.[2017]-all of which descend from early work in Baskin et al. [1997] , Goller and Küchler [1996] , Kireev [1995] , Merkwirth and Lengauer [2005] , Micheli and Sestito [2005] , Micheli [2009] , Scarselli et al. [2008] , and Sperduti and Starita [1997] .Limitations of MPNNs 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-Leman algorithm [Weisfeiler and Leman, 1968] , a well-studied heuristic for the graph isomorphism problem; see Section C. Additionally, they cannot capture global or long-range information, often linked to phenomena such as under-reaching [Barceló et al., 2020] or over-squashing [Alon and Yahav, 2021] , with the latter being heavily investigated in recent works.Graph Transformers Different from the above, graph transformers [Chen et al., 2022a ,b, Dwivedi et al., 2022b , He et al., 2023 , Hussain et al., 2022 , Kim et al., 2022 , Ma et al., 2023 , Mialon et al., 2021 , Müller et al., 2023 , Müller and Morris, 2024 , Rampášek et al., 2022 , Shirzad et al., 2023] 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 Müller et al. [2023] . However, all transformers suffer from their quadratic space and memory requirements due to computing an attention matrix.Rewiring Approaches for MPNNs Several recent works aim to circumvent over-squashing via graph rewiring. The most straightforward way of graph rewiring is incorporating multi-hop neighbors. For example, Brüel-Gabrielsson et al. [2022] rewires the graphs with k-hop neighbors and virtual nodes and 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-away 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] and balanced Forman curvature [Topping et al., 2021] . In addition, Deac et al. [2022] , Shirzad 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. [2022] propose a greedy random edge flip approach to overcome oversquashing. DiffWire [Arnaiz-Rodríguez et al., 2022] conducts fully differentiable and parameter-free graph rewiring by leveraging the Lovász bound and spectral gap. 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 rely on heuristic rewiring methods or purely randomized approaches that may not adapt well to a given prediction task. LASER [Barbero et al., 2023] performs graph rewiring while respecting the original graph structure. The recent work S2GCN [Geisler et al., 2024] combines spectral and spatial graph filters and implicitly introduces graph rewiring for message passing. Most similar to IPR-MPNNs is the work of Qian et al. [2023] , which, like IPR-MPNNs, leverage recent techniques in differentiable k-subset sampling [Ahmed et al., 2023] to learn to add or remove edges of a given graph. However, like GTs, their approach suffers from quadratic complexity due to their need to compute a score for each node pair. In addition to the differentiable k-subset sampling [Ahmed et al., 2023] method we use in this work, there are other gradient estimation approaches such as GUMBEL SOFTSUB-ST [Xie and Ermon, 2019] and I-MLE [Niepert et al., 2021 , Minervini et al., 2023] .There is also a large set of works from graph structure learning proposing heuristical graph rewiring approaches and hierarchical MPNNs; see Section A for details. 
background: 
 In the following, we introduce notation and formally define MPNNs.Notations Let N := {1, 2, 3, . . . }. For n ≥ 1, let [n] := {1, . . . , n} ⊂ N. We use { {. . . } }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 nodes or vertices V (G) and edgesE(G) ⊆ {{u, v} ⊆ V (G) | u ̸ = v}.If not otherwise stated, we set n := |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 ℓ(v) denotes its discrete vertex label, and so on; see Section 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 G = (G, X) be an n-order attributed graph with node feature matrix X ∈ R n×d , for d > 0, following, Gilmer et al. [2017] and Scarselli et al. [2008] , in each layer, t > 0, we compute vertex featuresh (t) v := UPD (t) h (t-1) v , AGG (t) { {h (t-1) u | u ∈ N (v)} } ∈ R d ,where UPD (t) and AGG (t) may be parameterized functions, e.g., neural networks, and hv := X v . In the case of graph-level tasks, e.g., graph classification, one usesh G := READOUT { {h (T ) v | v ∈ V (G)} } ∈ R d ,to compute a single vectorial representation based on learned vertex features after iteration T . Again, READOUT may be a 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 and Ba [2015] , together with the parameters of a neural network used for classification or regression. 
implicit probabilistically rewired mpnns: 
 Implicit probabilistically rewired message-passing neural networks (IPR-MPNNs) learn a probability distribution over edges connecting the original nodes of a graph to additional virtual nodes, providing an implicit way of enhancing the graph connectivity. To learn to rewire original nodes with these added virtual nodes, IPR-MPNNs use an upstream model, usually an MPNN, to generate scores or (unnormalized) priors θ := h u (A(G), X) ∈ R n×m , where G is an n-order graph with adjacency matrix A(G) ∈ {0, 1} n×n , node feature matrix X ∈ R n×d , for d > 0, and number of virtual nodes m with m ≪ n.IPR-MPNNs use the priors θ from the upstream model to sample new edges between the original nodes and the m virtual nodes from the posterior constrained to exactly k edges, thus obtaining an assignment matrix H ∈ {0, 1} n×m , i.e., an adjacency matrix between the input and virtual nodes. The assignment matrix H is then used in a downstream model f v , utilized for solving our downstream task, e.g., graph-level classification or regression. Leveraging recent advancements in gradient estimation for k-subset sampling [Ahmed et al., 2023] , the upstream and downstream models are jointly optimized, enabling the model to be trained end-to-end; see below.Hence, unlike PR-MPNNs [Qian et al., 2023] , which in the worst case explicitly model an edge distribution for all n 2 possible edge candidates for rewiring, IPR-MPNNs leverage virtual nodes for implicit rewiring and passing long-range information. Therefore, IPR-MPNNs benefit from better computation complexity while being more expressive than the 1-WL; see Section 4. Moreover, intuitively, it is also easier to model a distribution of edges connected to a few virtual nodes than to learn the explicit distribution of all possible edges in a graph. More specifically, while in PR-MPNNs, the priors θ can have a size of up to n 2 for an n-order graph, IPR-MPNNs use m • n parameters, therefore significantly enhancing both computational efficiency and model simplicity.In the following, we describe the IPR-MPNN architecture in detail.Sampling Edges Let A n represent the adjacency matrices of n-order graphs. Consider (G, X) as a graph of order n with adjacency matrix A(G) ∈ A n and node feature matrix X ∈ R n×d , for d > 0. IPR-MPNNs maintain a parameterized upstream model h u : A n × R n×d → Θ, usually implemented through an MPNN and parameterized by u. The upstream model transforms an adjacency matrix along with its node attributes into a set of unnormalized node priors θ ∈ Θ ⊆ R n×m , with m denoting the predefined number of virtual nodes. Formally,θ := h u (A(G), X) ∈ R n×m .The matrix of priors θ serves as the parameter matrix for the conditional probability mass function from which the assignment matrix H ∈ {0, 1} n×m is sampled. Crucially and contrary to prior work [Qian et al., 2023] , these edges connect the input and a small number m of virtual nodes. Hence, each row of matrix θ i: ∈ R m represents the unnormalized probability of assigning node i to each virtual node. Formally, we have,p θ (H i: ) := M j=1 p θij (H ij ), for i ∈ [n],where p θij (H ij = 1) = sigmoid(θ ij ) and p θij (H ij = 0) = 1 -sigmoid(θ ij ). Without loss of generality, we allow each node to be assigned to k virtual nodes, with k ∈ [m]. That is, each row of the sampled assignment matrix has exactly k non-zero entries, i.e.,p (θ,k) (H) := p θ (H) /Z if ∥H i: ∥ 1 = k, for all i ∈ [n], 0 otherwise, with Z := B∈{0,1} n×m : ∥Bi:∥1=k,∀ i∈[n]p θ (B).We can potentially sample independently q times H (i) ∼ p (θ,k) (H) and consequently obtain q multiple assignment matrices H := { {H (1) , H (2) , . . . , H (q) } }, which, together with corresponding number of copies of A(G) and X, will be utilized by the downstream model for the tasks.Message-passing Architecture of IPR-MPNNs Here, we outline the message-passing scheme after adding virtual nodes and edges. Consider an n-order graph G and a virtual node set C(G) of cardinality m, where each original nodev ∈ V (G) := [n] is assigned to k ∈ [m] virtual nodes.We assign original nodes v to a subset of the virtual node using the function a :V (G) → [C(G)] k , where v → {c ∈ C(G) | H vc = 1} and [C(G)] k is the set of k-element subsets of the virtual nodes.Conversely, each virtual node c ∈ C(G) links to several original nodes. Hence, we define an inverse assignment as the set of all original nodes assigned to virtual node c, i.e., a -1 (c) := {v ∈ V (G) | c ∈ a(v)}.Across the graph, the union of these inverse assignments equals the set of original nodes, i.e., c a -1 (c) = V (G).We represent the embedding of an original node v ∈ V (G) at any given layer t ≥ 0 as h(t)v , and similarly, the embedding for a virtual node c ∈ C(G) as g (t) c . To compute these embeddings, IPR-MPNNs compute initial embeddings for each virtual node. Subsequently, the architecture updates the virtual nodes via the adjacent original nodes' embeddings. Afterward, the virtual nodes exchange messages, and finally, the virtual nodes update adjacent original nodes' embeddings. Below, we outline the steps in order of execution involved in our message-passing algorithm in detail.Intializing Virtual Node Embeddings Before executing inter-hierarchical message passing, we need to initialize the embeddings of virtual nodes. To that, given the node assignments a(v), for v ∈ V (G), we can effectively divide the original graph into several subgraphs, where nodes sharing the same assignment label are grouped together to form an induced subgraph. Formally, for any virtual node c ∈ C(G), we have the induced subgraph G c with node subsetV c (G c ) := {v | v ∈ V (G) ∩ a -1 (c)} and edge set E c (G c ) := {{u, v} | {u, v} ∈ E(G), u ∈ a -1 (c) and v ∈ a -1 (c)}.The initial attributes of the virtual node c are defined by the node features of its corresponding subgraph G c , calculated using an MPNN, i.e.,EQUATIONAlternatively, we can generate random features for each virtual node as initial node features or assign unique identifiers to them.Updating Virtual Nodes In each step t > 0, we collect the embeddings from original nodes to virtual nodes according to their assignments and obtain intermediate virtual node embeddingsḡ(t) c := AGGn (t) { {h (t-1) v | v ∈ a -1 (c)} } , for c ∈ C(G),where AGGn (t) denotes some permutation-equivariant aggregation function designed for multisets.Updating Among Virtual Nodes We assume the virtual nodes form a complete, undirected, unweighted graph, and we perform message passing among the virtual nodes to update their embeddings.That is, at step t, we setg (t) c := UPDc (t) g (t-1) c , ḡ(t) c , AGGc (t) { {ḡ (t) j | j ∈ C(G), j ̸ = c} } ,where UPDc and AGGc are the update and neighborhood aggregation functions for virtual nodes.Updating Original Nodes Finally, we redistribute the embeddings from the virtual nodes back to the base nodes. This update process considers both the neighbors in the original graph and the virtual nodes to which the original nodes are assigned. The updating mechanism is detailed in the following equation,EQUATION)where UPD is the update function, AGG is the aggregation function, and DS is the distributing function that incorporates embeddings from virtual nodes back to the original nodes.Gradient Estimation In the context of our downstream MPNN model described above, we define the set of its learnable parameters as v and group these with the upstream model parameters u into a combined tuple w = (u, v). We express the downstream model as f v , resulting in the loss function 1) , . . . , H (q) } } , y .L A(G), X, H; w := E H (i) ∼p (θ,k) (H) ℓ f v A(G), X, { {H (While the gradients for the downstream model f v can be straightforwardly calculated via backpropagation, obtaining gradients for the upstream model parameters θ is more challenging, as the assignment matrices H are sampled from the priors θ, a process which is not differentiable.Similar to prior work [Qian et al., 2023] , we utilize SIMPLE [Ahmed et al., 2023] , which efficiently estimates gradients under k-subset constraints. This method involves exact sampling in the forward phase and uses the marginal of the priors µ(θ) ∈ R n×m during the backward phase to approximate gradients∇ θ L ≈ ∂ θ µ(θ)∇ H ℓ.The following analysis shows that our proposed method circumvents the problem of being quadratic in the number of input nodes.Complexity Assuming a constant number of hidden dimensions and layers of the MPNNs, recall that the runtime complexity of a plain MPNN is O(|E|), where |E| is the number of edges of a given graph. In the IPR-MPNN framework, we still have an MPNN backbone, augmented with Figure 2 : Comparing model sensitivity across different layers for the two most distant nodes from graphs from the ZINC dataset. On the left, we compare the sensitivity for models with a varying number of layers. We can observe that IPR-MPNNs maintain a high sensitivity even for the last layer, while the base models have the sensitivity decaying to 0. On the right, we compare models with a different number of virtual nodes, observing that the results are similar for all of the variants.inter + n • m + m 2 ).Since m ≪ n is a small constant, IPR-MPNNs show great potential due to their low complexity compared to the quadratic worst-case complexities of graph transformers [Müller et al., 2023] and other rewiring methods, e.g., Gutteridge et al. [2023] , Qian et al. [2023] .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
expressive power: 
 In this section, we analyze the extent to which IPR-MPNNs can separate non-isomorphic graphs on which the 1-WL isomorphism test fails [Xu et al., 2019 , Morris et al., 2019] , and whether IPR-MPNNs can maintain isomorphisms between pairs of graphs. We adopt the notion of probabilistic separation from Qian et al. [2023] .Our arguments rely on the ability of our upstream MPNN to assign arbitrary and distinct exactly-k distributions for each color class. By modifying the graph structure, we can make the rewired graphs 1-WL-distinguishable, enabling our downstream MPNN to separate them. However, since the expressiveness of the upstream model is also equivalent to 1-WL, there is a possibility of still separating isomorphic graphs.The following result demonstrates that we can preserve almost all partial subgraph isomorphisms.Theorem 4.1. Let k > 0, ε ∈ (0, 1), and G, H be two graphs with identical 1-WL stable colorings. Let M be the set of ordered virtual nodes, V G and V H be the subset of nodes in G and H that have a color class of cardinality 1, with (2) There exists a conditional probability mass function p (θ,k) that separates G[W G ] and H[W H ] with probability strictly greater than 0.|V G | = |V H | = d,We argue that preserving these partial subgraph isomorphisms is sufficient for most examples in practice. Indeed, our empirical findings show that we can successfully solve both the EXP and CSL datasets, whereas a 1-WL model obtains random performance; see Table A11 , Table A10 .The next corollary follows the above and recovers Theorem 4.1 from Qian et al. [2023] . The corollary tells us that, even if there are isomorphic graphs that we risk making separable, we will maintain the isomorphism between almost all isomorphic pairs. tTable 1 : We compare IPR-MPNN on QM9 with the base downstream GIN model [Xu et al., 2019] , two graph rewiring techniques [Gutteridge et al., 2023 , Qian et al., 2023] , a multi-hop MPNN [Abboud et al., 2022] , and the relational GIN [Schlichtkrull et al., 2018] . The best-performing method is colored in green, the second-best in blue, and third in orange. IPR-MPNN obtains the best result on all targets, except for MU, where it obtains the second-best result. Table 2 : Results on the PEPTIDES and PCQM-CONTACT datasets from the long-range graph benchmark [Dwivedi et al., 2022b] . For PCQM-CONTACT, we use both the initially proposed evaluation methodology (PCQM (1) ) and the filtering methodologies proposed in Tönshoff et al.[2023] (PCQM (2) for filtering and PCQM (3) for extended filtering). Green is the best model, blue is the second, and red the third. IPR-MPNNs obtain the best predictive performance on all datasets. Corollary 4.1.1. For sufficiently large n, for every ε ∈ (0, 1), a set m of ordered virtual nodes, and k > 0, we have that 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 :MODEL PEPTIDES-FUNC ↑ PEPTIDES-STRUCT ↓ PCQM (1) ↑ PCQM (2) ↑ PCQM (3) ↑ GINE [A n → R d , d > 0,there exists a probability mass function p (θ,k) that separates the graph G and H with probability at most ε regarding f .The previous theorems show that we are preserving isomorphisms better than purely randomized approaches while being more powerful than 1-WL since we can separate non-isomorphic graphs with a probability strictly greater than 0. We provide the proofs and examples in Section E. 
experimental setup and results: 
 To empirically validate the effectiveness of our IPR-MPNN framework, we conducted a series of experiments on both synthetic and real-world molecular datasets, answering the following researchquestions. An open repository of our code can be accessed at https://github.com/chendiqian/ IPR-MPNN.Q1 Do IPR-MPNNs alleviate over-squashing and under-reaching?Q2 Do IPR-MPNNs demonstrate enhanced expressivity compared to MPNNs?Q3 How do IPR-MPNNs compare in predictive performance on molecular datasets against other rewiring methods and graph transformers?Q4 Does the lower theoretical complexity of IPR-MPNNs translate to faster runtimes in practice? [Xu et al., 2019] , the GPS graph transformer [Rampášek et al., 2022] and the Drew model [Gutteridge et al., 2023] on the PEPTIDES-STRUCT dataset [Dwivedi et al., 2022b] . Our model has almost the same computation and memory efficiency as the base GINE model while being twice as fast and significantly more memory efficient when compared to GPS. To address Q1, we investigate whether our method alleviates over-squashing and under-reaching by experimenting on TREES-NEIGHBOURSMATCH [Alon and Yahav, 2021] and TREES-LEAFCOUNT [Qian et al., 2023] . On TREES-LEAFCOUNT with a tree depth of four, we obtain perfect performance on the test dataset with a one-layer downstream network, indicating we can alleviate under-reaching. Furthermore, on TREES-NEIGHBOURSMATCH, our method obtains perfect performance to a depth up to six, effectively alleviating over-squashing, as shown in Figure A4 . To quantitatively assess whether over-squashing is mitigated in real-world scenarios, we computed the average layer-wise sensitivity [Xu et al., 2018 , Di Giovanni et al., 2023 , Errica et al., 2023] between the most distant nodes in graphs from the ZINC dataset and compared these results with those from the baseline GINE model. Specifically, we compute the logarithm of the symmetric sensitivity between the most distant nodesu, v as log (| ∂h l v/∂h k u + ∂h l u/∂h k v |),where k to l represent the intermediate layers. We show that IPR-MPNNs maintain a high layer-wise sensitivity compared to the base model, as seen in Figure 2 , implying that they can successfully account for long-range relationships, even with multiple stacked layers. Lastly, we measured the average total effective resistance [Black et al., 2023] of five molecular datasets before and after rewiring, showing in Figure 3 that IPR-MPNNs are successfully improving connectivity by reducing the average total effective resistance of all evaluated datasets.For Q2, we conduct experiments on the EXP [Abboud et al., 2020] and CSL [Murphy et al., 2019] datasets to evaluate the expressiveness of IPR-MPNNs. The results, as detailed in Table A10 and Table A11 , demonstrate that our IPR-MPNN framework handles these datasets effectively and exhibits improved expressiveness over the base 1-WL-equivalent GIN model.For answering Q3, we utilize several real-world molecular datasets-QM9 [Hamilton et al., 2017] , ZINK 12K [Jin et al., 2017] , OGB-MOLHIV [Hu et al., 2020] , TUDATASET [Morris et al., 2020a] , and datasets from the LONG-RANGE GRAPH BENCHMARK [Dwivedi et al., 2022b] , namely PEPTIDES and PCQM-CONTACT. Our results demonstrate that IPR-MPNNs effectively account for longrange relationships, achieving state-of-the-art performance on the PEPTIDES and PCQM-CONTACT datasets, as detailed in Table 2 . Notably, on the PCQM-CONTACT link prediction tasks, IPR-MPNNs outperform all other candidates across three measurement metrics outlined in Tönshoff et al. [2023] . For QM9, we show in Table 1 that IPR-MPNNs greatly outperform similar methods, obtaining the best results on 12 of 13 target properties. On ZINC and OGB-MOLHIV, we outperform similar MPNNs and graph transformers, namely GPS Rampášek et al. [2022] and SAT [Chen et al., 2022a] , obtaining state-of-the-art results; see Table 4 . For the TUDATASET collection, we achieve the best results on four of the five molecular datasets; see Table A9 .Finally, to address Q4, we evaluate the computation time and memory usage of IPR-MPNNs in comparison with the GPS graph transformer [Rampášek et al., 2022] on PEPTIDES-STRUCT and extend our analysis to include PR-MPNNs [Qian et al., 2023] , SAT [Chen et al., 2022a] , and GPS on the ZINC dataset. The results in Tables 3 and A12 demonstrate that IPR-MPNNs adhere to their theoretical linear runtime complexity in practice. We observed a notable speedup in training and validation times per epoch while reducing the memory footprint by a large margin compared to the two mentioned transformers. This efficiency underscores the practical advantages of IPR-MPNNs in computational speed and resource utilization. We compute the log of total effective resistance [Black et al., 2023] of five molecular datasets before and after rewiring the graphs using virtual nodes. Our rewiring technique consistently lowers the total effective resistance, indicating a better information flow on all of the datasets.  
conclusion: 
 Here, we introduced implicit probabilistically rewired message-passing neural networks (IPR-MPNNs), a graph-rewiring approach leveraging recent progress in end-to-end differentiable sampling. IPR-MPNNs show drastically improved running times and memory usage efficiency over graph transformers and competing rewiring-based architectures due to IPR-MPNNs' ability to circumvent comparing every pair of nodes and significantly outperforming them on real-world datasets while effectively addressing over-squashing and overreaching. Hence, IPR-MPNNs represent a significant step towards designing scalable, adaptable MPNNs, making them more reliable and expressive. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? Answer: [Yes] Justification: We reviewed the NeurIPS Code of Ethics, and our work respects it. Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [NA] Justification: This paper conducts foundational research in the area of graph learning. While certainly our work could be used both for positive and negative societal impact, we do not foresee any immediate positive or negative societal impact. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: We release neither data nor large-scale models as part of this work. Further, our experiments are conducted on comparatively small, curated, task-specific datasets used for benchmarking graph learning models. Hence, our work does not pose immediate risks for misuse. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [NA] Justification: This work does not involve crowdsourcing or research with humans. Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: In the following, we discuss related work.Graph Structure Learning Graph structure learning (GSL) is closely related to graph rewiring, where the primary motivation is refining and optimizing the graph structure while jointly learning graph representations [Zhou et al., 2023a] . Several methods have been proposed in this field. Jin et al.[2020] develop a technique to optimize graph structures from scratch using a specific loss function as a bias, while the general approach is using edge scoring functions for refinment [Chen et al., 2020 , Yu et al., 2021 , Zhao et al., 2021] , but discrete sampling methods have also been applied. More specifically, DGM [Kazi et al., 2022] is predicting the latent graph structure by leveraging Gumbel discrete sampling, while Franceschi et al. [2019] is learning a Bernoulli distribution via Hypergradient Descent. Saha et al. [2023] learns adaptive neighborhoods for trajectory prediction and point cloud classification by sampling through the smoothed-Heaviside function, while Younesian et al. [2023] samples nodes that are used for downstream tasks using GFlowNets. Some GSL techniques also employ unsupervised learning [Zou et al., 2023 , Fatemi et al., 2021 , Liu et al., 2022a,b] . We encourage the reader to refer to Fatemi et al. [2023] , Zhou et al. [2023b] for detailed surveys regarding GSL.Our proposed IPR-MPNN is different from other GSL framework in two main ways: (1) instead of sparsifying the graph using randomized k-NN approaches or independent Bernoulli random variables, we learn a probability mass function with exactly-k constraints [Ahmed et al., 2023] . Moreover, we don't aim to discover the graph structure by considering a fully-connected latent graph from which we sample new edges [Qian et al., 2023] , instead we introduce sparse connections from base nodes to virtual nodes, with complexity N • k;(2) GSL methods do not investigate exact sampling of the exactly-k distribution; however, one of our aims is to demonstrate that these techniques can significantly alleviate information propagation issues caused by inadequate graph connectivity, such as over-squashing and under-reaching.Moreover, our IPR-MPNN also differs from previous graph rewiring method, specifically PR-MPNN [Qian et al., 2023] . First, we rewire the graph implicitly by connecting nodes to virtual nodes, respecting the original graph structure. Besides, we show a significant run-time advantage in that our worst-case complexity is sub-quadratic, while PR-MPNN optimally needs to consider n 2 node pairs for an n-order graph.Hierarchical MPNNs Our method draws connections to hierarchical MPNNs. The hierarchical model initially emerged in graph-level representation learning, as seen in approaches like AttPool [Huang et al., 2019] , DiffPool [Ying et al., 2018] , and ASAP [Ranjan et al., 2020] . Further developments, such as Graph U-Net [Gao and Ji, 2019] , H-GCN [Hu et al., 2019] , and GXN [Li et al., 2020] , introduced top-down and bottom-up methods within their architectures. However, they did not incorporate virtual node message passing. Other works create hierarchical MPNNs while incorporating inter-hierarchical message passing. For example, Fey et al. [2020] introduced HIMP-GNN on molecular learning, using a junction tree to create a higher hierarchy of the original graph and do inter-message passing between the hierarchies. Rampášek and Wolf [2021] proposed HGNet for long-range dependencies, generating hierarchies with edge pooling and training with relational GCN. Zhong et al. [2023] designed HC-GNN, more efficient than HGNet, for better node and higher level resolution community representations. These hierarchical MPNNs require well-chosen heuristics for hierarchy generation. Using an auxiliary supernode is particularly prominent in molecular tasks [Gilmer et al., 2017 , Pham et al., 2017 , Li et al., 2017] , which involves adding a global node to encapsulate graph-level representations. Further advancements in this area, as suggested in [Battaglia et al., 2018] , and developments like GWM Ishiguro et al. [2019] , have enhanced supernode MPNNs with a gating mechanism. Theoretically, Cai et al. [2023] proves MPNN with a virtual node can simulate self-attention, which is further investigated in Rosenbluth et al. [2024] . In addition, Southern et al. [2024] studied the effect of MPNNs using a virtual node on over smoothing and oversquashing.Simultaneously, the recent study by Li et al. [2023] has introduced the concept of a collection of supernodes, termed "neural atoms," which incorporate supernode message passing with an attention mechanism. Moreover, the idea of a coarsened hierarchical graph has become widely employed in scalable MPNN training and graph representation learning, as evidenced by works like Huang et al.[2021], Liang et al. [2021] , Namazi et al. [2022] . Unlike existing hierarchical MPNNs, our IPR-MPNN uniquely leverages differential k-subset sampling techniques for dynamic, probabilistic graph rewiring and incorporates hierarchical message passing in an end-to-end trainable framework. This approach enhances graph connectivity and expressiveness without relying on predefined heuristics or fixed structures.A graph G is a pair (V (G), E(G)) with finite sets of vertices or nodes V (G) and edgesE(G) ⊆ {{u, v} ⊆ V (G) | u ̸ = v}.If not otherwise stated, we set n := |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). A (vertex-)labeled graph G is a triple (V (G), E(G), ℓ) with a (vertex-)label function ℓ : V (G) → N. Then ℓ(v) is a label of v, for v in V (G). An attributed graph G is a triple (V (G), E(G), σ) with a graph (V (G), E(G)) and (vertex-)attribute function σ : V (G) → R 1×d , for some d > 0. That is, contrary to labeled graphs, we allow for vertex annotations from an uncountable set. Then σ(v) is an attribute or feature of v, for v in V (G). Equivalently, we define an n-order attributed graph G := (V (G), E(G), σ) as a pair G = (G, L), where G = (V (G), E(G)) and L in R n×d is a node feature matrix. Here, we identifyV (G) with [n]. For a matrix L in R n×d and v in [n], we denote by L v• in R 1×d the vth row of L such that L v• := σ(v). Furthermore, we can encode an n-order graph G via an adjacency matrix A(G) ∈ {0, 1} n×n , where A ij = 1 if, and only, if (i, j) ∈ E(G). We also write R d for R 1×d .The neighborhood of v in V (G) is denoted by N (v) := {u ∈ V (G) | (v, u) ∈ E(G)}(φ(u), φ(v)) is in E(H).Then φ is an isomorphism between G and H. In the case of labeled graphs, we additionally require that l(v) = l(φ(v)) for v in V (G), and similarly for attributed graphs.A node coloring is a function c : V (G) → R d , d > 0, and we say that c(v) is the color of v ∈ V (G).A node coloring induces an edge coloring e c : E(G) → N, where (u, v) → {c(u), c(v)} for (u, v) ∈ E(G). A node coloring (edge coloring) c refines a node coloring (edge coloring) d, written c ⊑ d if c(v) = c(w) implies d(v) = d(w) for every v, w ∈ V (G) (v, w ∈ E(G)).The 1-WL or color refinement is a fundamental, well-studied heuristic for the graph isomorphism problem, originally proposed by Weisfeiler and Leman [1968] . 1 The algorithm is an iterative method starting from labeling or coloring vertices in both graphs with degrees or other information, and updating the color of a node with its color as well as its neighbors' colors. During the iterations, two vertices with the same label get different labels if the number of identically labeled neighbors is unequal. Each iteration ends up with a vertex color partition, and the algorithm terminates when the partition is not refined by the algorithm, i.e., when a stable coloring or stable partition is obtained. We can finally conclude that the two graphs are not isomorphic if the color partitions are different, or the number of nodes of a specific color is different. Although in Cai et al. [1992] the limitation is shown that 1-WL algorithm cannot distinguish all non-isomorphic graphs, it is a powerful heuristic that can succeed on a broad class of graphs [Arvind et al., 2015 , Babai and Kucera, 1979 , Babai et al., 1980] .Formally, let G = (V (G), E(G), ℓ) be a labeled graph. In each iteration, t > 0, the 1-WL computes a vertex coloring C 1 t : V (G) → N, depending on the coloring of the ego node and of the neighbors. That is, in iteration t > 0, we setC 1 t (v) := RELABEL C 1 t-1 (v), { {C 1 t-1 (u) | u ∈ N (v)} } ,for all vertices v ∈ V (G), where 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 := ℓ is used. To test whether 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 vertices colored c ∈ 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) ⇐⇒ C 1 t+1 (v) = C 1 t+1 (w), for all vertices v and w in V (G ∪H), then the algorithm terminates. For such t, we define the stable coloringC 1 ∞ (v) = C 1 t (v), for v ∈ V (G ∪H).The stable coloring is reached after at most max{|V (G)|, |V (H)|} iterations [Grohe, 2017, Kiefer and McKay, 2020] . A function f : V (G) → R d , for d > 0, is 1-WL-equivalent if f ≡ C 1 ∞ .Here, we provide additional empirical results and experimental details.Details In all of our real-world experiments, we use two virtual nodes with a hidden dimension twice as large as the base nodes. We randomly initialize the features of the virtual nodes. For the upstream and downstream models, we do a hyperparameter search; see Table A5 . We use RWSE and LapPE positional encodings [Dwivedi et al., 2022a] for all of our experiments as additional node features, except for the synthetic TREES datasets [Alon and Yahav, 2021, Qian et al., 2023] , EXP [Abboud et al., 2020] and CSL [Murphy et al., 2019] , as well as node classification tasks on heterophilic datasets. For TREES-LEAFCOUNT, we use a single one-layer downstream GINE network, and for TREES-NEIGHBOURSMATCH, we use n + 1 layers, where n is the depth of the tree. For both TREES datasets, we have a hidden dimension of 32 for the original nodes. For most graph-level tasks, we apply a read-out function over the final pooled node embeddings, virtual node embeddings, or a combination of both. Distinctly, for PEPTIDES-FUNC, we apply read-out functions and use a supervised loss on all the intermediate embeddings, similarly to Errica et al. [2023] . We compute the logarithm of the symmetric sensitivity between the most distant two nodes u, v in Fig. 2 similar to Errica et al. [2023] , i.e. 1, and the total Effective Resistance of the datasets in Fig. 3 as in Black et al. [2023] . We optimize the network using Adam Kingma and Ba [2015] with a cosine annealing learning rate scheduler. We use the official dataset splits when available. Notably, for the TUDATASET, WEBKB datasets [Craven et al., 1998 ] and heterophilic datasets proposed in Platonov et al. [2023] , we perform a 10-Fold Cross-Validation and report the average validation performance, similarly to the other methods that we compare with. For some experiments, we search for hyperparameters using grid-search, for more details please see Table A5 . All experiments were performed on a mixture of A10, A100, A5000, and RTX 4090 NVIDIA GPUs. For each run, we used at most eight CPUs and 64 GB of RAM.We exhibit experimental results on the WEBKB datasets in Table A6 , and on the heterophilic datasets in Table A7 . Our method exhibits significant improvement on WEBKB datasets as well as ROMAN-EMPIRE datasets compared with other MPNN baselines.Ablations We conduct ablation experiments on the number of virtual nodes and repetition of samples, see Table A8 .We first discuss three scenarios where we (1) separate isomorphic graphs with the same non-discrete stable 1-WL colorings (Figure A5 ), ( 2) separate non-isomorphic graphs with the same non-discrete 1-WL colorings (Figure A6 ), and ( 3) preserve isomorphisms between isomorphic graphs with the same discrete 1-WL colorings (Figure A7 ). For each of the three examples, we assume two unique virtual nodes with k = 1, i.e., we sample a single edge between a base node and a virtual node.For example (1) , consider Figure A5 and assume that, for each color class, the upstream MPNN randomly selects a virtual node, i.e., we have a uniform prior. Since both base nodes have the same [Alon and Yahav, 2021] for a depth up to 6. color class, each base node connects to one of the virtual nodes uniformly at random. The scenario where two graphs remain isomorphic is when they connect to exactly the same virtual nodes or either connect as in case (a) or (b). Therefore, we have a high probability of separating these isomorphic graphs.1 1 3 2 1 1 3 2 1 1 3 2 1 1 3 2 (a) (b) (c) (d)For example (2) , we again produce uniform priors for each color class of the graphs in Figure A6 .Once again, there is a high probability of separation. One possible configuration that can separate between the two graphs is shown in Figure A6 where in the first graph (a), the nodes colored with 1 get assigned to virtual node 4, while the nodes colored with 2 are assigned to virtual node 3. In the second graph (b), all nodes get assigned to the same virtual node 3.For example (3) in Figure A7 , we consider producing priors that assign, with high probability, the same virtual node to all of the nodes that are in color classes of cardinality 1. This approach ensures that the discretely-colored graphs remain isomorphic with high probability.The intuition is that if we want to distinguish between graphs that have the same 1-WL stable color partitioning, the upstream model needs to produce "uninformative" prior weights for some color classes. However, preserving isomorphisms is most likely when the nodes in the same color class in the two graphs get assigned to the same virtual nodes. Since our upstream model is as powerful as 1-WL, we can control the prior distribution for the color classes but not for individual nodes, therefore we can only guarantee high assignment probabilities for nodes in color classes of cardinality 1.Next, we formally argue that we can preserve, with arbitrarily high probability, isomorphisms between graphs that have the same discrete stable 1-WL color partitions, as well as isomorphisms between subgraphs with the same discrete stable 1-WL color partitions. Lemma E.1. Let G and H be a pair of graphs with the same 1-WL graph coloring, i.e., they are1-WL non-distinguishable. Let k ∈ N, let G ′ k be color-induced subgraph where V (G ′ k ) := {v ∈ V (G) | c 1 ∞ (v) = k} ⊆ V (G), and V (H ′ ) similarly. Then the subgraphs induced by V (G ′ k ) and V (H ′ k ) are still not 1-WL-distinguishable.Proof. To prove the lemma, we use the concept of an unrolling tree for a node; see, e.g., Morris et al. [2020b] . That is, for a node, we recursively unroll its neighbors, resulting in a tree. It is easy to see that two nodes get the same colors under 1-WL if and only if such trees are isomorphic; see Morris et al. [2020b] for details. Figure A7 : An example where we preserve the isomorphism when all the nodes from the initial graph are discretely colored. For each base node, we assign a high probability of connecting to the same virtual node (node 8).Consider nodes v, u ∈ V (G ′ k ) ∪ V (H ′ k )that share the same stable 1-WL coloring k. Based on the above, v and u must have isomorphic unrolling trees. Now remove subgraphs of the two trees not rooted at the vertex with color k. Since both v and u have isomorphic unrolling trees, the resulting trees are isomorphic. Hence, running 1-WL on top of G ′ k and H ′ k will still not distinguish them.Lemma E.2. Let G and H be a pair of graphs with the same stable coloring under 1-WL. If we add a finite number of virtual nodes on both graphs C(G), C(H), and connect these virtual nodes based on 1-WL colors of the original graphs, i.e., two equally colored vertices get assigned the same virtual nodes. Then, the augmented graphs Ĝ and Ĥ have the same stable partition.Proof. The proof is by straightforward induction on the number of iterations using the fact that two nodes with the same color will be assigned to the same virtual node. That is, the neighborhood of two such nodes is extended by the same nodes.Besides, we leverage the following result by Qian et al. [2023] , Morris et al. [2019] .Lemma E.3 (Qian et al. [2023] , Morris et al. [2019] ). Let G be an n-order graph and let c : V (G) → R d , d > 0, be a 1-WL-equivalent node coloring. Then, for all ε > 0, there exists a (permutationequivariant) MPNN f :V (G) → R d , such that max v∈V (G) ∥f (v) -c(v)∥ < ε.Theorem E.4. Let k > 0, ε ∈ (0, 1), and G, H be two graphs with identical 1-WL stable colorings. Let M be the set of ordered virtual nodes, V G and V H be the subset of nodes in (2) There exists a conditional probability mass function p (θ,k) that separates G[W G ] and H[W H ] with probability strictly greater than 0.Proof. Let M = {v 1 , ..., v m } be the set of m ordered virtual nodes andd = |V G | = |V H |.To prove this theorem, we can leverage Lemma E.3 to assign distinct and arbitrary priors for every color class.For (1), we know that since G and H have identical 1-WL stable colorings and V G , V H have a color class of cardinality 1, then G[V G ] and H[V H ] must be discrete and isomorphic. Using Lemma E.3, we can obtain an upstream MPNN that assigns a sufficiently high prior θ i such that, when we sample from the exactly-k distribution, we can assign the corresponding k virtual nodes to a base node with probability at least 2d √ 1 -ε.To demonstrate the existence of such a set of priors θ, let δ ∈ (0, 1) and S ⊂ M be a subset of k virtual nodes. Let w 1 > w 2 be two prior weights such thatEQUATIONθ i = w 2 if v i ∈ M \ S. We have that p θ,k (S) ≥ δ k i=0 k i m -k k -i w i 1 w k-i 2 = δZ,with the upper bound [Qian et al., 2023 ]Z ≤ w k 1 + m k -1 w 2 w k-1 1 ,Thus, θ exists and can be obtained using this inequality.Next, we set δ = 2d √ 1 -ε. Consequently, the probability that the sampled virtual nodes are identical for both graphs is at least 2d √ 1 -ε 2d = 1 -ε. Finally, using Lemma E.2, we know that the two graphs also retain their color partitions and remain isomorphic with probability at least 1 -ε.For (2), it is easy to see that, for any prior weights θ that we assign to the color classes of cardinality greater than 1, there is at least one configuration separating the two graphs. ∈ (0, 1), v ∈ G[W G ]∪H[W H] and θ u be an uniform prior, i.e. θ 1 = θ 2 = . . . = θ m . Again, we use Lemma E.3 and obtain a distribution p θ,k , arbitrarily close to the uniform distribution.Then, the probability of making the two graphs distinguishable by obtaining the mentioned example is greater than1-ε ( m k )2n , which is strictly greater than 0. For a visual example, see Figure A6 .The next Corollary follows directly from Theorem E.4, and recovers Theorem 4.1 from Qian et al.[2023].Corollary E.4.1. For sufficiently large n, for every ε ∈ (0, 1), a set M of ordered virtual nodes, and k > 0, we have that 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 :A n → R d , d > 0,there exists a probability mass function p (θ,k) that separates the graph G and H with probability at most ε with respect to f .Proof. We know from Babai et al. [1980] that an 1-WL-equivalent algorithm will produce a discrete color partition for almost all pairs of isomorphic graphs G, H of sufficient size. We use Theorem E.4 and set W G = W H = ∅ and conclude that we maintain isomorphisms between almost all isomorphic graphs.A limitation of our approach is the assumption that the number of virtual nodes m is significantly smaller than the total number of nodes n. As the number of virtual nodes increases, the runtime is also expected to rise (see Table A12 for a detailed example). In the worst-case scenario, where m = n, our method exhibits quadratic complexity. However, in all real-world datasets we have tested, the required number of virtual nodes for achieving optimal performance is low. For more information, refer to Table A5 . Another question is whether IPR-MPNNs can perform well on node-level tasks.We have designed our rewiring method specifically to solve long-range graph-level tasks (such as the tasks on the molecular datasets from the LONG-RANGE GRAPH BENCHMARK [Dwivedi et al., 2022b] ). Nevertheless, IPR-MPNNs and adaptations might also work on node-level tasks, but we leave this question open for further work.Table A5 : Overview of used hyperparameters. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.Strictly speaking, the 1-WL and color refinement are two different algorithms. That is, the 1-WL considers neighbors and non-neighbors to update the coloring, resulting in a slightly higher expressive power when distinguishing vertices in a given graph; seeGrohe [2021] for details. Following customs in the machine learning literature, we consider both algorithms to be equivalent. 


Decision: This paper proposed implicitly rewired message-passing neural networks (IPR-MPNNs), which integrate implicit probabilistic graph rewiring into MPNNs to alleviate the under-reaching and over-squashing issues. Most of the reviewers agreed that the proposed approach is novel and this paper is technically solid. All the reviews tend to accept the paper during the reviewer-author discussion phase. Therefore, I recommend this paper to the NeurIPS 2024 conference.

Reviews: review: paper_summary: paper_summary: This paper introduces Implicitly Rewired Message-Passing Neural Networks (IPR-MPNNs) to address the limitations of traditional MPNNs, such as under-reaching and over-squashing. By integrating implicit probabilistic graph rewiring and adding a small number of virtual nodes, IPR-MPNNs enable long-distance message propagation without quadratic complexity. This approach enhances expressiveness and computational efficiency.
strength_weakness: strength_weakness: 1. **Risk of Over-Smoothing**:
   - While IPR-MPNNs alleviate squashing and under-reaching, the approach of connecting all nodes via virtual nodes increases the risk of over-smoothing. A discussion on over-smoothing analysis would strengthen the paper.

2. **Applicability to Node-Level Tasks**:
   - The method should be directly applicable to node-level tasks, where long-range interactions are particularly beneficial. Conducting experiments on node-level datasets requiring long-range interactions and comparing with existing works would enhance the paper.

3. **Unclear Method Description**:
   - The benefit of sampling \( H \) \( q \) times, its implications, and the associated costs and disadvantages need further clarification.
   - It is unclear how to ensure that the union of the inverse assignments covers the entire set of original nodes.
   - The initial set of \( G_c \) is not clearly defined.

4. **Runtime Measurement**:
   - Better reporting of runtime measurements with respect to the number of \( q \) would provide valuable insights.

5. **Notation Confusion**:
   - There is some notation confusion, such as between \( h^{(t)} \) and \( H^{(q)} \), which needs to be clarified.

6. **Optimal Parameters and Ablation Studies**:
   - Reporting the optimal \( m \) and \( q \) for each experiment and conducting ablation studies on these parameters would provide insights on how to choose them effectively.
questions: questions: 1. Does the method exacerbate over-smoothing?
2. Why is it necessary to sample $q$ times? Does this serve a similar function as multi-head attention?
3. How do you ensure that the union of the inverse assignments covers the entire set of original nodes?
4. What is the initial set of $G_c$ ?
5. What is the runtime when $q$ is large?
6. How does the method perform on node-level tasks that require long-range interactions?
7. Can more ablation studies be conducted to provide additional insights?
limitations: limitations: yes
paper_summary: paper_summary: The paper introduces a method (IPR-MPNN) which connects nodes to a small number of virtual nodes in a learnable way. The proposed approach is more expressive than MPNN whilst circumventing quadratic complexity. It can reduce oversquashing and performs well on various benchmarks.
strength_weakness: strength_weakness: - Whilst the method is shown well in the context of MPNNs and GTs, the probabilistic rewiring is less well motivated in the context of standard virtual nodes. Adding a baseline of MPNN + VN (with your number of VNs) would help and your results seem like they would be much better. Additionally, explaining why only connecting a fraction of nodes to the VN (instead of all) helps in terms of expressivity could also be mentioned/explored - this seems like it would be the case intuitively. You could also compare the effective resistance of your IPR-GNN graph to one where we just have a single global VN. As MPNN + VN is quite popular in the community, this would help us understand why we should favour your approach for certain tasks. 
- In line 68, you suggest that your approach helps with "scaling to large graphs effectively". To me, it is not clear if this is the case. Firstly, the results seem to be mainly on small molecular datasets. Additionally, whilst you only use a small number of VNs on these datasets (eg. 2), This can still be 10% of the number of nodes in the graph. It is not clear if for large graphs you also need to have 10% VNs compared to total number of nodes. If this is the case, would your approach still be practical?
questions: questions: - Do you have some explanation for why your approach improves over GTs? Can GTs not replicate this method?
- It would be useful to the community to mention the fraction of nodes which are connected to the VNs on some of these datasets. Connecting to a single node seems similar to node-marking and to all nodes would be equivalent to standard VN. It is interesting to see how this approach falls in this region for certain tasks.
limitations: limitations: some limitations have been addressed in the appendix. For example: "We have designed our rewiring method specifically to solve long-range graph-level tasks" - I assume the model is flexible here though in that it could also only use local interactions by connecting to. single node or node neighbours.
paper_summary: paper_summary: The paper proposes implicitly rewired message-passing neural networks (IPR-MPNNs), which integrate implicit probabilistic graph rewiring into MPNNs. This method involves introducing a small number of virtual nodes into the graph, allowing for long-distance message propagation without the quadratic complexity associated with graph transformers.
strength_weakness: strength_weakness: 1. It remains unclear to me how the unnormalized node priors $\theta$ are derived. On Line 170, the authors mentioned that $\theta \in \Theta$, but what is $\Theta$? According to Line 169, is it the output of the upstream GNN? Why can we trust that it accurately reflects the original node-virtual node connectivity? Is there any related optimization objective?

2. Gradient estimators for graph sampling are not a new topic. The authors may refer to [1, 2]. Providing related ablations might be more convincing.

3. I do not fully understand Figure 2. As stated by the authors, the values in the figure are about "the two most distant nodes," but why is this referred to as model sensitivity? How is this sensitivity defined?

4. Although the authors provided a complexity analysis, I believe it is necessary to provide numerical comparisons, such as the per epoch time and wall-clock time before and after using IPR-MPNN.


[1] Robust Graph Representation Learning via Neural Sparsification, NeurIPS22

[2] Differentiable Graph Module (DGM) for Graph Convolutional Networks, TPAMI
questions: questions: 1. Can you provide an algorithm table that details the training process of IPR-MPNN？
limitations: limitations: Yes
paper_summary: paper_summary: This work introduces an approach called implicitly rewired message-passing neural networks (IPR-MPNNs) to address limitations in message-passing graph neural networks (MPNNs): expressiveness and mainly over-squashing. Their method works by adding virtual nodes and learning to connect them to existing nodes in a differentiable, end-to-end manner. Authors claim that previous techniques are not scalable, and state that their approach is more scalable. They empirically show the performance of their proposed model.
strength_weakness: strength_weakness: ### Related work

* The state of the art is in general not properly structured and not enough detailed. Although they cite some of the important work, thay do not delve into the detail of the different papers about limitations on oversquashing (Di Giovanni 2023, Black 2023, they cite both of them, but authors dot seem to cite them when talking about limitations or even they dont mention the rewiring methods they propose SDRF and GTR), which is one of the main points in this work. For instance, when talking about graph rewiring they cite several works, no detailing what are the contributions of each and also mixing works of different solutions in the same sentence: MixHop is based on a adaptative message passing and DIGL is rewiring basd on difussion, but authors just say that both works are "considered as graph rewiring as they can reach further-away neighbors in a single layer". I mean, authors mix adaptative MP with spatial rewiring with spectral rewiring, expanders with no explanation about it. I do not think that the level of structure and detail is enough for a paper in this conference. See Di giovanni 2023, Banerjee 2022a(FOSR), 2022b(RLEF), Choi 2024(Pandas), for an example of better structured and detailed explanation of the different rewiring methods.

* In addition, they miss very relevant works in the literature about graph rewiring and the analysis of graph limitations that are related with this paper. Specifically, authors evaluate their contribution using commute time or effective resistance. They miss 2 very important rewiring works in this line DiffWire (Arnaiz-Rodriguez, 2022) and FOSR (Banerjee, 2022), where they introduce the oversquashing connectino with the Cheger constant and methods to reduce the effective resistance between nodes to overcome oversquashing. To be even more clear regarding Banerjee 2022: in the references appears RLEF and FOSR, but in the text there is only one reference to RLEF. They also do not mention GTR from the already cited paper Black 2023, which is also based on effective resistance, metric that authors then use to evaluate IPR-MPNN. I suggest the authors to look more carefully also Di Giovanni 2023 on all the impact of not only these but also more works in the graph rewiring literature, and to completely cite all the important first works in this area.

* On top of DiffWire and FOSR, authors also miss more important papers in grph rewiring such as Drew (Gutteridge 2023), LASER (Barbero 2024) or Affinity-Aware GNNs (Velingker). 

* In addition, they miss relevant papers in the literature about virtual nodes, which is one of the main points in this paper. For instance, the two first important works about virtual nodes (Battaglia 2017 and Brüel-Gabrielsson 2022) are just cited in the appendix and/or not even mentioning that they proposed the virtual nodes. Additionally, authors are not aware of works focused on the limitations of virtual nodes and how they are related with oversquashing: Southern 2024 and Geisler 2024. These papers completelly address how *theoretically* virtual nodes affect oversquashing and how they can be used to mitigate it (see methodology weaknesses). I suggest the authors to look more carefully at the literature about virtual nodes and oversquashing, and to completely and properly cite all the important first works in this area correctly stating the main contributions of the prior work.


### Methodology

* The contribution of this work is limited since it comes from the combination of virtual nodes and probabilistic message passing. Both ideas were already previously proposed. It would be helpful if authors are more transparent on what are the exact novelty of this method. Virtual nodes ar not novelty as proposed before (Battaglia 2017) and probabilitistic message passing as well. The changes to make this approach novel are not sufficient in my opinion.
  
* Regarding the contributions, only the expressiveness contribution is shown theoretically, and it is the less related with the problem they claim to solve (OVSQ). The rest of the contributions are just shown empirically, which does not advance on the knowledge of the limitations of the GNN field.
* Authors do not justify thoretically how their approach improves oversquashing, which should be the main approach of the paper. They just empirically show that the effective resitance and sensitivity are smalles, which is a trivial result of the Rayleigh’s monotonicity principle when densifying a graph. The paper would increase a lot the quality if authors really tryy to understand why their method improve oversauashing and if it really does. Souther 2024 and Geisler 2024 might be helpful for this. To reach the level of Neurips authors should have a more theoretical approach to the problem they are trying to solve. However, it is monstly empirical as they also aknowledge in the section 5 (they answer all the questions just empirically). This is very dangerous, since the empirical results might not be generalizable to other datasets or tasks. Also, in the GNN community is widely known the lack of understanding of how GNNs work if we only look at the results of datasets.

  
* One key contribution is that they do rewiring it in less than quadratic complexity (which is only ok for transphormers but is not ok for rewiring methods). However this is only empirically, since authors also mention l865 "In the worst-case scenario, where m = n, our method exhibits quadratic complexity. However, in all real-world datasets we have tested, the required number of virtual nodes for achieving optimal performance is low". 

* Also, mention do not delve into the limitation of the virtual and oversquashing (See Southern 2024 and Geisler 2024). For instance, altough the effective resistance is reduced using virtual nodes, the oversquashing is not fully solved since all the messages are squeezed in the virtual nodes, even worsening the problem of exponential aggregation of information in some bottleneck "nodes". As shown in Topping 2022 and Arnaiz-Rodriguez 2022, the blottleneckness of the nodes (measured either by curvature, spectral metrics or centralities), causes oversquashing.

* As their main experiments are in graph classification tasks,  and the main contribution is based on pairwise scores, they also should compare aagaints CT-Layer in Diffwire (Arnáiz-Rodrigues 2023), which is a rewiring based also in similarity scores that was shown to improve graph classification tasks. Also, authors claim that previous "strategies to mitigate over-squashing rely on heuristic rewiring methods or purely randomized approaches that may not adapt well to a given prediction task". Also, diffwire is in-processing rewiring and takes into account the task. Finally, also Geisler 2024 is task-oriented.

  
* They compare the running time against transformer but not against rewiring, which does not make sense. They propose an architecture to do graph rewiring and alliviate oversquashing, so they should compare with graph rewiring methods.


### Refs

P. W. Battaglia,  et al. Relational inductive biases, deep learning, and graph networks. 2017.

A. Arnaiz-Rodriguez et al. DiffWire: Inductive Graph Rewiring via the Lovász Bound. LoG, 2022.

P. K. Banerjee et al. Oversquashing in gnns through the lens of information contraction and graph expansion. AACCC, 2022.

R. Brüel-Gabrielsson, M. Yurochkin, and J. Solomon. Rewiring with positional encodings for graph neural networks. 2022

F. Di Giovanni et al. On over-squashing in message passing neural networks: The impact of width, depth, and topology. ICML 2023.

M. Black et al. Understanding Oversquashing in GNNs through the Lens of Effective Resistance. ICML 2023.

D. Gutteridge et al. Drew: Dynamically rewired message passing with delay. ICML 2023.

A. Velingker. Affinity-Aware Graph Networks. NeurIPS 2023.

F. Barbero. Locality-Aware Graph Rewiring in GNNs. ICML 2024.

J. Choi. PANDA: Expanded Width-Aware Message Passing Beyond Rewiring. ICML 2024

J. Southern. Understanding Virtual Nodes: Oversmoothing, Oversquashing, and Node Heterogeneity. 2024.

S. Geisler. Spatio-Spectral Graph Neural Networks. 2024.
questions: questions: Questions are listed along with each specific identified weaknesses in the previous section.
limitations: limitations: I summarize the limitations of the work, however, I refer to weaknesses section for a more detailed feedback.

* Relate work is not complete and not clearly explained and structured.
* The authors do not closely connect how virtual nodes solve oversquahing.
* Authors do not discuss limitations of the virtual nodes wrt over-squashing.
* All the important contributions are empirical, and the theoretical ones are not very novel.
* The core of the idea is not very novel, just a combination of previous work on virtual nodes an probabilistic message passing, but not a big contribution by itself. Also, they mention as contributions that their methods alliviates the limitations of GNNs, but there is no explanation why, just empirical results.

Comments: title: Please Engage in Discussion
Please Engage in Discussion
Rebuttal - References
Rebuttal - New comparisons and experiments
Rebuttal - P2 (cont. Methodology)
Rebuttal - References
Rebuttal - New experimental results
Rebuttal - References
comment: Thank you to the authors for addressing my concern with the oversmoothing measurement. It has alleviated my concerns. I will maintain my score.
We thank the reviewer for their suggestion.

While our method is not tailored to alleviate over-smoothing, we agree that having evidence that our model does not _increase_ the phenomena of over-smoothing is beneficial.

We have measured the Dirichlet Energy of the final layer on the heterophilic WebKB Cornell, Wisconsin, Texas datasets and on the roman-empire dataset from [PV23]. We compare IPR-MPNN with the base GINE model.

The results are averaged for 5 runs, and we also report the standard deviations:

| Model       | Cornell ↑             | Wisconsin ↑           | Texas ↑               | Roman-Empire ↑        |
|-------------|-----------------------|-----------------------|-----------------------|-----------------------|
| GINE (base) | 8.34 ± 4.74           | 5.09 ± 3.11           | 7.41 ± 4.57          | 19.36 ± 1.16          |
| IPR-MPNN    | **12.10 ± 4.77**          | **8.89 ± 2.67**           | **9.65 ± 1.90**           | **49.49 ± 2.41**          |

As can be observed, IPR-MPNNs have higher Dirichlet Energy on Cornell, Wisconsin, Texas and significantly higher on roman-empire, indicating that over-smoothing is alleviated for these datasets. 

Please note that IPR-MPNN also obtains much better overall results on these datasets (see previous comment).

We will add these results to the final version of the paper. Overall, the empirical evidence indicates that IPR-MPNNs _do not_ increase over-smoothing, but they rather slightly alleviate the phenomena.

Thank you for helping us improve our work. We kindly ask the reviewer to reconsider their assessment and increase their score if they find our response convincing. We are open to answering any remaining questions if time permits.
I appreciate the author’s detailed response to my questions and comments, as well as the additional experimental results. However, I still have concerns about oversmoothing. Incorporating a measurement based on Dirichlet Energy might provide further insights. After considering the feedback from other reviewers, I have decided to maintain my current score.
We thank the reviewer for their positive re-assessment of our work!

As we mentioned in the rebuttal, we will include the missing related work and the concurrent works of [SN24, GR24] in the final version of the paper. We agree that there are some strong connections between our work and [SN24, GR24], and we will add a discussion where we highlight the similarities and differences between the approaches.

We will also lower our claims by clarifying that we only have proof for higher expressiveness in the WL sense but no direct theoretical evidence that we alleviate OSM/OSQ. We will only highlight that the empirical evidence and performance indicate that this _might_ be the case. 

For further work, it would also be very interesting to investigate whether the theoretical results from [SN24] could be extended to our probabilistic method and if our sparsification via sampling strategy adds any theoretical benefits, as is the case for WL expressiveness. 

Once again, we thank the reviewer for their feedback and for engaging with us during the rebuttal discussions period!
Thanks for your thoughtful rebuttal. It looks like some of my questions/points were valid and some had to do with a misunderstanding of the work as it is currently presented. Thank you for taking the time to explain which was which. I also acknowledge some misunderstanding of some parts of the paper on my part. I reply to specific parts where I still do not fully agree with the authors:

**[Related work]** Regarding related work, I pointed out several (+10, most of them not concurrent and even *seminal*) works that are quite aligned and related to the topic the authors propose (rewiring, analysis of why and how rewiring works and virtual nodes). So I still think my point about missing related work was valid. Regarding the concurrent work, of course I agree and I am aware with the authors that these 2 works are concurrent. However, I think that these works, especially SN24, are of special interest since they provide theoretical insights and limitations and the interplay of virtual nodes and the over-problems. Even if this work is concurrent, the theoretical insights identified in this paper are crucial for this work. They find that a virtual node mitigates over-squashing only for some graphs. They discover that the average commute time is reduced when a virtual node is added, and also that the Jacobian between 2 nodes is independent of the distance between them if they are separated by more than 2 hops after a virtual node is added (only for some classes of graphs).


**[Theoretical explanations of VN-Overproblems]** I still maintain my position as the authors do not prove theoretically that their method solves over-quashing (as SN24 does), but they do an empirical analysis about it. Don't get me wrong, I don't think a deeper theoretical analysis is necessary for the proposed paper, but the authors claim several times that "the method reduces over-..." just based on experimental results. Considering that a theoretical analysis might be aimed for a follow-up paper, I suggest the authors to lower the claims about solving over-smoothing and over-squashing and frame it as a result of the empirical work. Good empirical work can stand on its own, especially in practically oriented papers like this one, but the statements need to be in tone with the paper and the community in which it is presented.


***After reading the explanations of my concerns, the discussion of related work, the discussion of introducing bottlenecks, and the extensive additional experiments (runtime and baselines) performed by the authors, I raise my score from 3 to 5 and advocate acceptance of the paper.***.
We once again thank the reviewer for their suggestions and questions!
Dear Reviewer 5s4E,

Thank you for your time and efforts throughout the review period. Please read the authors' rebuttal as soon as possible and indicate if their responses have addressed all your concerns. 

Best,

Area Chair
Dear Reviewer UsNW,

Thank you for your time and efforts throughout the review period. Please read the authors' rebuttal as soon as possible and indicate if their responses have addressed all your concerns. 

Best,

Area Chair
Thank you for your response. After reading the response and other reviewers' comments, I choose to maintain my score.
We once again thank the reviewer for their positive assessment of our work and for their comments, which helped us improve our paper!
Thank you very much for the detailed response to my questions. 

The discussion of the advantages of having sparse, learnable connections over fully-connected virtual nodes will indeed make the paper stronger and the results look good for large graphs. Adding the number of virtual nodes (as a fraction of total number of nodes) used will be useful here in understanding the scaling. 

Because of this I have raised my score and will push for acceptance.
[SN24]: Southern, J. et al. “Understanding Virtual Nodes: Oversmoothing, Oversquashing, and Node Heterogeneity”, arXiv 2024

[GR24]: Geisler, S. et al., “Spatio-Spectral Graph Neural Networks”, arXiv 2024

[GE23]: Gutteridge, B. et al., “DRew: Dynamically Rewired Message Passing with Delay”, ICML 2023

[BGN22]: Brüel-Gabrielsson, R. et al., “Rewiring with Positional Encodings for Graph Neural Networks”, TMLR 2023

[BA17]: Battaglia, P.W. et al., “Relational inductive biases, deep learning, and graph networks”, arXiv 2018

[BE23]: Banerjee, P.K. et al., “Oversquashing in GNNs through the lens of information contraction and graph expansion”, ALLERTON 2022

[KR22]: Karhadkar, K. et al., “FoSR: First-order spectral rewiring for addressing over-squashing in GNNs”, ICLR 2023

[ARZ22]: Arnaiz-Rodriguez, A. et al., “DiffWire: Inductive Graph Rewiring via the Lovász Bound”, LoG 2022

[VR23]: Velingker, A. et al., “Affinity-Aware Graph Networks”, NeurIPS 2023

[BO24]: Barbero, F. et al., “Locality-Aware Graph Rewiring in GNNs”, ICML 2024

[QN24]: Qian, C. et al., “Probabilistically-Rewired Message-Passing Neural Networks”, ICLR 2024

[PM17]: Pham, T. et al., “Graph Classification via Deep Learning with Virtual Nodes”, arXiv 2017

[LI17]: Li, J. et al., “Learning graph-level representations for drug discovery”, arXiv 2017

[CI23]: Cai, C. et al., “On the connection between MPNN and Graph Transformers”, ICML 2023

[GR19]: Gasteiger, J., Weißenberger, S., and Günnemann, S. "Diffusion improves graph learning." NeurIPS 2019.

[PI20]: Pei, H., et al. "Geom-gcn: Geometric graph convolutional networks." ICLR 2020.

[TG21]: Topping, J., et al. "Understanding over-squashing and bottlenecks on graphs via curvature." ICLR 2022.

[RK22]: Rampášek, L., et al. "Recipe for a General, Powerful, Scalable Graph Transformer." NeurIPS 2022.

[YG21]: Ying, C., et al. "Do Transformers Really Perform Bad for Graph Representation?" NeurIPS 2021.

[NS24-1]: NeurIPS 2024 Reviewer Guidelines

[NS24-2]: NeurIPS 2024 Call For Papers
---
**Comparison with DiffWire [ARZ22] on molecular datasets:**

| Datasets | GAP (R) [ARZ22] 	| GAP (N) [ARZ22] 	| IPR-MPNN 	|
|----------|-------------|-------------|--------------|
| MUTAG	&uarr;| 86.9 ± 4.0 | 86.9 ± 4.0  | **98.0 ± 3.4**   |
| PROTEINS &uarr;| 75.0 ± 3.0 | 75.3 ± 2.1 | **85.4 ± 4.4**   |

---
**Comparison with various methods (incl. DiffWire, SDRF) on the heterophilic WebKB datasets**

| Methods      | Cornell &uarr;   | Texas &uarr; | Wisconsin &uarr; |
|-------------------------------|-------------------------|-------------| --- |
| GINE                          	| 0.448±0.073             | 0.650±0.068 	| 0.517±0.054 | 
| DIGL [GR19] 			| 0.582±0.005             | 0.620±0.003 	| 0.495±0.003|
| Geom-GCN [PI20]  		| 0.608±N/A                | 0.676±N/A    	| 0.641±N/A |
| SDRF [TG21]  			| 0.546±0.004             | 0.644±0.004 	| 0.555±0.003 |
| DiffWire [ARZ22]  			| 0.690±0.044		| N/A		| 0.791±0.021 |
| GPS [RK22] 			| 0.718±0.024 		| 0.773±0.013  | 0.798±0.090 |
| Graphormer [YG21] 		| 0.683±0.017		| 0.767±0.017	| 0.770±0.019|
| IPR-MPNN 			| __0.764±0.056__ 	| __0.808±0.052__ | __0.804±0.052__ |

---

**Comparison with S2GCN [GR2024] and GatedGCN+VN$_G$ [SN2024] on the long-range peptides datasets.**

| Model                      	| Func  &uarr;              	| Struct &darr;            	|
|--------------------------------|-------------------------|------------------------|
| S2GCN [GR2024]              	| 0.7275 ± 0.0066     	| 0.2467 ± 0.0019    	|
| S2GCN+PE [GR2024]           	| **0.7311 ± 0.0066**     	| 0.2447 ± 0.0019    	|
| GatedGCN+PE+VN$_G$ [SN2024, best] | 0.6822 ± 0.0052     	| 0.2458 ± 0.0006    	|
| IPR-MPNN                   	| 0.7210 ± 0.0039     	| **0.2422 ± 0.0007**    	|

---

**Runtime comparison with the Drew rewiring method [GE23].**

| Feature  | IPR-MPNN     	| Drew [GE23]        	|
|----------|------------------|------------------|
| Model	| IPR-MPNN     	| Drew         	|
| Params   | 536k         	| 522k         	|
| Mem  	| 1.9GB        	| 1.8GB        	|
| s/train  | 2.98s ±0.02  	| 3.20s ±0.03  	|
| s/val	| 0.27s ±0.00  	| 0.36s ±0.00  	|
* **Connections with [GR24, SN24]: We want to reiterate that [GR24, SN24] were not publicly available at the time of submission.** Investigating the connections between IPR-MPNNs and these works would be very interesting. [GR24] is similar to IPR-MPNNs in its virtual node interpretation but differs fundamentally as it uses spectral layers to connect virtual nodes to the base graph, while we use a probabilistic, differentiable k-subset sampling approach to connect the base graph to the virtual graph in a data-driven manner. Our method is also similar to [SN24] since they also employ a form of heterogeneous message-passing. However, they do not use multiple virtual nodes and do not sparsify the connections between the original graph and the virtual nodes. Therefore, the theory in [SN24] is not directly applicable to our probabilistic approach. 

We will update the manuscript with a discussion regarding our connections with the two papers.

---

* **Quadratic worst-case complexity:** We acknowledge in our Supplementary that the worst-case complexity of our model is quadratic. However, it is quadratic only in the number of virtual nodes, thereby quadratic in the number of original nodes only if the number of virtual nodes is greater or equal to the number of original nodes. **If the original graph is complete, MPNNs also have quadratic complexity in the number of nodes. Any rewiring method that might lead to a complete graph is potentially quadratic.** This is an improbable scenario for both general rewiring methods and IPR-MPNNs. We demonstrate that using up to 8 virtual nodes provides state-of-the-art performance (Table A5) while remaining sub-quadratic in the number of original nodes. We also analyze performance for up to 30 virtual nodes in Table A9, showing runtime and memory improvements over Graph Transformers even in the worst-case scenario.

---

* **Virtual nodes are bottlenecks:** The primary motivation for our probabilistic layer that selectively connects the initial graph to the virtual nodes is to alleviate potential information bottlenecks between the base and virtual graphs. This acts as a sparsification operation, reducing the information transmitted to individual virtual nodes. This is different from [SN24] since they analyze a single virtual node connected to the entire original graph. Our hierarchical, heterogeneous message-passing scheme allows for bigger hidden dimensions for the virtual graph, which can also help alleviate bottlenecks. Empirically, we show that a low number of virtual nodes is sufficient for passing long-range information in most practical scenarios. For example, in the Trees-NeighborsMatch results (Figure A4), IPR-MPNNs perform as well as a Transformer up to a depth of 6, using only two virtual nodes. 

We thank the reviewer for the question - this is a very good point, and we will add a more nuanced discussion about possible bottlenecks in the virtual graph in our next paper revision.

---

* **No comparison with [ARZ22]:** We compare with publicly available results of various rewiring approaches across multiple datasets: Drew, SPN, and PR-MPNN on QM9 (Table 2); Drew, PR-MPNN, and AMP on LRGB (Table 3); PR-MPNN on ZINC and OGB-Molhiv (Table 4); and FoSR (Table A6). Re-running all possible baselines is very expensive and time-consuming. [ARZ23] contains some results that we can compare with, and we have also performed new experiments with IPR-MPNNs on the heterophilic datasets reported in their paper.

As can be seen in the next official comment, we obtain significantly better results than DiffWire on the molecular and heterophilic datasets.

---

* **No comparisons with [GR24, SN24]:** The authors of the two parallel works, which became available after the NeurIPS submission deadline, also report results on peptides. As can be seen in the next official comment, our method obtains overall better scores on the Peptides datasets when compared with [SN24], and is competitive with [GR24], obtaining slightly worse results on func, and slightly better results on struct.

We will update the tables for the revision, including results from [ARZ22, SN24, GR24] in our paper.

---

* **No runtime comparison with other rewiring methods:** We compare with one rewiring method (PR-MPNN) in our extended runtime analysis (Table A9 in the Supplementary). Additionally, we include comparisons with the base GINE in both tables as a lower bound for required resources. In all scenarios, IPR-MPNNs' overhead over the base GINE model is generally low (see Table 1 and Table A9). 

We have also performed an additional experiment comparing IPR-MPNNs with Drew [GE23]. We fix the same number of message-passing layers with a similar number of parameters, and we run the Drew model from the code provided by the authors on the same hardware. The two methods are comparable, with IPR-MPNNs being slightly faster, with slightly higher memory requirements (please see the next official comment).
[GR19]: Gasteiger, J., Weißenberger, S., and Günnemann, S. "Diffusion improves graph learning." NeurIPS 2019.

[PI20]: Pei, H., et al. "Geom-gcn: Geometric graph convolutional networks." ICLR 2020.

[TG21]: Topping, J., et al. "Understanding over-squashing and bottlenecks on graphs via curvature." ICLR 2022.

[ARZ22]: Arnaiz-Rodríguez, A., et al. "Diffwire: Inductive graph rewiring via the Lovász bound." LoG 2022.

[RK22]: Rampášek, L., et al. "Recipe for a General, Powerful, Scalable Graph Transformer." NeurIPS 2022.

[YG21]: Ying, C., et al. "Do Transformers Really Perform Bad for Graph Representation?" NeurIPS 2021.

[PV23]: Platonov, O., et al. "A critical look at the evaluation of GNNs under heterophily: Are we really making progress?", ICLR 2023
**Q1, Q2: New experimental results on the WebKB heterophilic datasets and on the heterophilic datasets proposed in [PV23]:**


| Methods      | Cornell &uarr;   | Texas &uarr; | Wisconsin &uarr; |
|-------------------------------|-------------------------|-------------| --- |
| GINE                          	| 0.448±0.073             | 0.650±0.068 	| 0.517±0.054 | 
| DIGL [GR19] 			| 0.582±0.005             | 0.620±0.003 	| 0.495±0.003|
| Geom-GCN [PI20]  		| 0.608±N/A                | 0.676±N/A    	| 0.641±N/A |
| SDRF [TG21]  			| 0.546±0.004             | 0.644±0.004 	| 0.555±0.003 |
| DiffWire [ARZ22]  			| 0.690±0.044		| N/A		| 0.791±0.021 |
| GPS [RK22] 			| 0.718±0.024 		| 0.773±0.013  | 0.798±0.090 |
| Graphormer [YG21] 		| 0.683±0.017		| 0.767±0.017	| 0.770±0.019|
| IPR-MPNN 			| __0.764±0.056__ 	| __0.808±0.052__ | __0.804±0.052__ |


| Model        	| Roman-empire &uarr;           	| Tolokers &uarr;               	| Minesweeper  &uarr;           	| Amazon-ratings &uarr;         	|
|------------------|-----------------------------|-----------------------------|-----------------------------|-----------------------------|
| GINE (base)  	| 0.476±0.006             	| 0.807±0.006             	| 0.799±0.002             	| **0.488±0.006**         	|
| IPR-MPNN (ours)  | **0.839±0.006**         	| **0.820±0.008**         	| **0.887±0.006**         	| 0.480±0.007             	|

---

**Q4: Reduced runtime report, also available in Table A9 of the Appendix:**



| Model  	| #Params | V. Nodes | Samples | Train s/ep      	| Val s/ep        	| Mem. usage |
|------------|---------|----------|---------|---------------------|---------------------|------------|
| GINE   	| 502k	| -    	| -   	| 3.19 ± 0.03     	| 0.20 ± 0.01     	| 0.5GB  	|
| GraphGPS   | 558k	| -    	| -   	| 17.02 ± 0.70    	| 0.65 ± 0.06     	| 6.6GB  	|
| IPR-MPNN   | 548k	| 2    	| 1   	| 7.31 ± 0.08     	| 0.35 ± 0.01     	| 0.8GB  	|
| IPR-MPNN   | 548k	| 4    	| 2   	| 7.37 ± 0.08     	| 0.35 ± 0.01     	| 0.8GB  	|
| IPR-MPNN   | 549k	| 30   	| 20  	| 9.41 ± 0.38     	| 0.43 ± 0.01     	| 1.2GB  	|


---

**Q6: Table comparing a virtual node connected to the entire graph (1VN - FC) with IPR-MPNN with two virtual nodes, one sample (2VN1S), two samples (2VN2S) and four samples (2VN4S)**

| Model           	| ZINC &darr;            	| ogb-molhiv &uarr;      	| peptides-func &uarr;  	| peptides-struct 	&darr;|
|---------------------|----------------------|----------------------|---------------------|---------------------|
| 1VN - FC        	| 0.074 ± 0.002    	| 0.753 ± 0.011    	| 0.7039 ± 0.0046 	| 0.2435 ± 0.0007 	|
| IPR-MPNN 2VN1S  	| 0.072 ± 0.004    	| 0.762 ± 0.014    	| 0.7146 ± 0.0055 	| 0.2472 ± 0.0014 	|
| IPR-MPNN 2VN2S  	| **0.067 ± 0.001**    	| **0.788 ± 0.006**    	| **0.7210 ± 0.0039** 	| **0.2422 ± 0.0007** 	|


| Model	| peptides-func &uarr;  	|
|----------|---------------------|
| 1VN - FC | 0.7039 ± 0.0046 	|
| 2VN1S	| 0.7146 ± 0.0055 	|
| 2VN2S	| **0.7210 ± 0.0039** 	|
| 2VN4S	| 0.7145 ± 0.0020 	|
[GR19]: Gasteiger, J., Weißenberger, S., and Günnemann, S. "Diffusion improves graph learning." NeurIPS 2019.

[PI20]: Pei, H., et al. "Geom-gcn: Geometric graph convolutional networks." ICLR 2020.

[TG21]: Topping, J., et al. "Understanding over-squashing and bottlenecks on graphs via curvature." ICLR 2022.

[ARZ22]: Arnaiz-Rodríguez, A., et al. "Diffwire: Inductive graph rewiring via the Lovász bound." LoG 2022.

[RK22]: Rampášek, L., et al. "Recipe for a General, Powerful, Scalable Graph Transformer." NeurIPS 2022.

[YG21]: Ying, C., et al. "Do Transformers Really Perform Bad for Graph Representation?" NeurIPS 2021.

[PV23]: Platonov, O., et al. "A critical look at the evaluation of GNNs under heterophily: Are we really making progress?", ICLR 2023