[Paper Review] Probabilistically Rewired Message-Passing Neural Networks

Chendi Qian, Andrei Manolache+5 more|arXiv (Cornell University)|Oct 3, 2023

Advanced Graph Neural Networks4 citations

TL;DR

This paper proposes Probabilistically Rewired Message-Passing Neural Networks (PR-MPNNs), a differentiable, principled framework that learns to add or remove edges in a graph based on task-relevant priors using exact, differentiable k-subset sampling. The method enhances expressive power, mitigates over-squashing and under-reaching, and achieves state-of-the-art or competitive performance on molecular and citation graph benchmarks, outperforming both standard MPNNs and graph transformers in key settings.

ABSTRACT

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

Motivation & Objective

• To address the limitations of message-passing GNNs, such as over-squashing and limited expressive power due to fixed graph structures.
• To develop a principled, differentiable approach for learning graph structure modifications that are tailored to the prediction task.
• To overcome the shortcomings of heuristic rewiring methods that do not adapt to data distribution or introduce random edges without justification.
• To theoretically analyze how probabilistic rewiring enhances model expressivity and identify conditions under which it outperforms random rewiring.
• To empirically validate that PR-MPNNs achieve superior or competitive performance compared to MPNNs and graph transformers on real-world datasets.

Proposed method

• PR-MPNNs use an upstream model to learn a probability distribution over candidate edges, parameterized as a k-subset distribution conditioned on exactly k edges.
• The framework samples multiple k-edge adjacency matrices from this learned distribution, which are then aggregated (e.g., via subtraction) to form a final, differentiable graph structure.
• Backward gradients are computed via exact, differentiable marginals of the k-subset distribution, reducing bias and variance in training.
• The method leverages recent advances in exact and differentiable k-subset sampling to enable end-to-end training of the edge selection mechanism.
• The final graph structure is used as input to a downstream MPNN for prediction, enabling joint optimization of both structure and representation learning.
• The approach is designed to be scalable and differentiable, allowing for effective backpropagation through the stochastic edge selection process.

Experimental results

Research questions

• RQ1 Under what conditions does probabilistic rewiring enhance the expressive power of message-passing GNNs compared to fixed or randomly rewired graphs?
• RQ2 How does the differentiable k-subset sampling mechanism enable effective and low-variance training of edge selection policies?
• RQ3 To what extent does PR-MPNN mitigate over-squashing and under-reaching in graph neural networks?
• RQ4 Can PR-MPNNs achieve performance competitive with or superior to graph transformers and state-of-the-art MPNNs on real-world graph benchmarks?
• RQ5 How does the learned edge selection policy adapt to the underlying data distribution and prediction task?

Key findings

• On the Peptides-struct dataset from the LRGB benchmark, PR-MPNN achieved the best MAE score of 0.2477 ± 0.0005, outperforming DRew-GIN (0.2799 ± 0.0016) and other baselines.
• On the heterophilic and transductive WebKB datasets, PR-MPNN achieved the highest accuracy on Cornell (0.659 ± 0.040), Texas (0.827 ± 0.032), and Wisconsin (0.750 ± 0.015), surpassing all MPNN baselines and competing with graph transformers.
• PR-MPNN significantly reduced over-squashing and under-reaching by learning to add long-range edges that improve information flow, as evidenced by improved performance on tasks requiring global structure.
• The method achieved a 1.98× improvement in mean absolute error (MAE) on the gap property prediction task compared to the base GIN model (1.70 ± 0.01 vs. 3.37 ± 0.03).
• On the ZPVE property, PR-MPNN achieved a 14.13× improvement over the base GIN model (4.73 ± 0.08 vs. 66.87 ± 1.45), demonstrating strong mitigation of information compression.
• Theoretical analysis confirmed that PR-MPNNs can achieve higher expressive power than standard MPNNs under precise conditions related to the distribution of learned edges and their alignment with task-relevant structures.