Research map / Graph learning / GNN architectures
Probabilistically Rewired Message-Passing Neural Networks
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes PR-MPNNs, which learn a k-subset edge distribution (via exact differentiable SIMPLE sampling) to rewire input graphs for MPNNs, with theory on expressivity and experiments on over-squashing, expressivity benchmarks, and molecular/TU datasets.
Problem
MPNNs operate on a fixed input graph, suffer from over-squashing, under-reaching and limited expressive power; existing rewiring methods are heuristic or purely random and their effect on expressivity is unclear.
Approach
Learn priors over candidate edges with an upstream MPNN, sample exactly-k edge subsets from a k-subset distribution using the SIMPLE gradient estimator, aggregate the sampled adjacency matrices into a rewired graph, and feed it to a downstream MPNN trained end-to-end.
Key results
- On TUDataset, PR-MPNN reports the best mean accuracy on Mutag (98.4), PTC_MR (74.3), Proteins (80.7) and NCI109 (84.6), and second on NCI1 (85.6) behind the WL kernel.
- On Exp, PR-MPNN reaches 1.000 accuracy versus 0.511 for GIN; on CSL without positional encodings it reaches 0.998 versus 0.100 for GIN.
- On Zinc without edge features, PR-MPNN with SIMPLE reaches MAE 0.139 vs 0.258 for the base GIN model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use probabilistic graph rewiring with the Simple gradient estimator instead of the base GIN | GIN · Zinc - Edge | MAE: 0.258 ±0.006 -> 0.139 ± 0.001 | Helped |
| Use the Simple gradient estimator rather than Gumbel SoftSub-ST or I-MLE (ablation) | PR-MPNN · Zinc - Edge | MAE: 0.153 ±0.003; 0.151 ±0.001 -> 0.139 ±0.001 | Helped |
| Use probabilistic graph rewiring with the Simple gradient estimator instead of the base GIN | GIN · Zinc + Edge | MAE: 0.207 ±0.006 -> 0.085 ±0.002 | Helped |
| Use the Simple gradient estimator rather than Gumbel SoftSub-ST or I-MLE (ablation) | PR-MPNN · Zinc + Edge | MAE: 0.103 ±0.008; 0.104 ±0.008 -> 0.085 ±0.002 | Mixed |
| Use probabilistic graph rewiring with the Simple gradient estimator instead of the base GIN | GIN · ogbg-molhiv + Edge | ROC-AUC: 0.775 ±0.011 -> 0.795 ±0.009 | Helped |
| Use the Simple gradient estimator rather than Gumbel SoftSub-ST or I-MLE (ablation) | PR-MPNN · ogbg-molhiv + Edge | ROC-AUC: 0.760 ±0.025; 0.774 ±0.015 -> 0.795 ±0.009 | Helped |
| Use probabilistic graph rewiring with the Simple gradient estimator instead of the base GIN | GIN · Alchemy + Edge | MAE: 11.12 ±0.690 -> 6.447 ±0.057 | Helped |
| Use the Simple gradient estimator rather than Gumbel SoftSub-ST or I-MLE (ablation) | PR-MPNN · Alchemy + Edge | MAE: 6.858 ±0.090; 6.692 ±0.061 -> 6.447 ±0.057 | Helped |
| Add positional encodings (ablation) | PR-MPNN · CSL | accuracy: 0.998 ± 0.008 -> 1.000 ± 0.000 | No clear effect |
| Use PR-MPNN without positional encodings rather than GIN without positional encodings | GIN · CSL | accuracy: 0.100 ± 0.000 -> 0.998 ± 0.008 | Helped |
| Use PR-MPNN rather than GIN | GIN · Exp | accuracy: 0.511 ± 0.021 -> 1.000 ± 0.000 | Helped |
| Use PR-MPNN rather than the strongest listed baseline | GIN · Mutag | accuracy: 94.4 ±3.7 -> 98.4 ±2.4 | Helped |
Surprising findings
- Uniform random edge removal can be as good as learned sampling on some graph pairs, per the negative proposition on regular graphs.
Limitations
- Worst-case O(n^2) complexity when considering all candidate edges, requiring heuristic subset selection of candidate edges.
- GIN aggregation of the downstream model is a limiting factor on heterophilic node-level data, where graph transformers outperform PR-MPNN except on Texas.
Experimental setup
Datasets: Trees-NeighborsMatch Trees-LeafCount Exp CSL 4-Cycles Zinc Alchemy ogbg-molhiv QM9 LRGB TUDataset (Mutag, PTC_MR, Proteins, NCI1, NCI109, IMDB-B, IMDB-M) Cornell
Benchmarks: Trees-NeighborsMatch Trees-LeafCount Exp CSL 4-Cycles Zinc Alchemy ogbg-molhiv TUDataset Cornell Wisconsin Texas
Metrics: test accuracy MAE ROC-AUC 10-fold CV accuracy
Baselines: GPS SAT Gumbel SoftSub-ST I-MLE DropGNN FoSR CIN CIN++ CAN GIN ID-GNN OSAN
Builds on: SIMPLE GIN DropGNN DropEdge GINE
Closest papers by meaning
- Probabilistic Graph Rewiring via Virtual Nodes (NeurIPS 2024 · similarity 0.94)
- Locality-Aware Graph Rewiring in GNNs (ICLR 2024 · similarity 0.89)
- Improving the Effective Receptive Field of Message-Passing Neural Networks (ICML 2025 · similarity 0.88)
- PANDA: Expanded Width-Aware Message Passing Beyond Rewiring (ICML 2024 · similarity 0.88)
- Accurate and Scalable Graph Neural Networks via Message Invariance (ICLR 2025 · similarity 0.87)
- Graph Cascades: Contagion-Based Mesoscopic Rewiring for Structure-Aware Graph Machine Learning (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.