Research map / Graph learning / GNN architectures
Probabilistic Graph Rewiring via Virtual Nodes
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 IPR-MPNN, which implicitly rewires graphs by adding a few virtual nodes and sampling exactly-k node-to-virtual-node edges end-to-end via differentiable k-subset sampling (SIMPLE), achieving linear-in-edges complexity.
Problem
MPNNs suffer from over-squashing and under-reaching, and existing rewiring methods or graph transformers either use heuristics or have quadratic complexity in the number of nodes.
Approach
Add a small number of virtual nodes and learn a distribution over exactly-k edges from each original node to virtual nodes (upstream MPNN gives priors; SIMPLE gives exact sampling and marginal-based gradients), then run message passing across original nodes, virtual nodes, and back, so the rewiring costs O(n·m) instead of n^2 pairs.
Key results
- On Peptides-struct IPR-MPNN reaches 0.2422 MAE versus 0.2460 for GRIT and 0.2431 for AMP, and on PCQM-Contact it is best on all three filtering variants (e.g., 0.4756 PCQM(3)).
- On QM9 IPR-MPNN obtains the best result on 12 of 13 targets, e.g., alpha 1.36 vs 1.63 for Drew-GIN.
- Training time per epoch on Peptides-struct is 2.98s vs 7.81s for GPS, with 22.2GB memory for GPS vs 1.9GB for IPR-MPNN.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add implicit probabilistic virtual-node rewiring instead of using a standard graph baseline | Graph neural networks · PCQM-Contact PCQM(1) | average precision: 0.3637 ±0.0020 -> 0.3670 ±0.0082 | Helped |
| Add implicit probabilistic virtual-node rewiring instead of using a standard graph baseline | Graph neural networks · Peptides-struct | MAE: 0.2431 ±0.0004 -> 0.2422 ±0.0007 | Helped |
| Add implicit probabilistic virtual-node rewiring | GINE · Roman-Empire | accuracy: 0.476 ±0.006 -> 0.839 ±0.006 | Helped |
| Add implicit probabilistic virtual-node rewiring | GINE · Amazon-ratings | accuracy: 0.488 ±0.006 -> 0.480 ±0.007 | Hurt |
| Use two sampled virtual nodes instead of one fully connected virtual node (ablation) | IPR-MPNN · ZINC | MAE: 0.074 ±0.002 -> 0.067 ±0.001 | Helped |
| Use two samples rather than one (ablation) | IPR-MPNN · ogb-molhiv | ROC-AUC: 0.762 ±0.014 -> 0.788 ±0.006 | Helped |
| Use two samples rather than one (ablation) | IPR-MPNN · Peptides-func | average precision: 0.7146 ±0.0055 -> 0.7210 ±0.0039 | Helped |
| Increase samples from two to four (ablation) | IPR-MPNN · Peptides-func | average precision: 0.7210 ±0.0039 -> 0.7145 ±0.0020 | Hurt |
| Use two samples rather than one (ablation) | IPR-MPNN · Peptides-struct | MAE: 0.2472 ±0.0014 -> 0.2422 ±0.0007 | Helped |
| Use implicit virtual-node rewiring instead of a standard GIN | GIN · Exp | accuracy: 0.511 ± 0.021 -> 1.000 ± 0.000 | Helped |
| Use implicit virtual-node rewiring instead of a standard GIN | GIN · Csl | accuracy: 0.100 ± 0.000 -> 0.987 ± 0.013 | Helped |
Surprising findings
- On Amazon-Ratings IPR-MPNN (0.480) underperforms the base GINE (0.488), and on Tolokers it is only marginally better.
- On Trees-LeafCount and Trees-NeighboursMatch the model reaches perfect accuracy with a single-layer downstream network for depth four, showing under-reaching is alleviated.
Limitations
- Complexity is stated as O(|E| + n·m + m^2) with m small, so benefits depend on keeping the number of virtual nodes small.
- Several tables show mixed results (e.g., Mutag, Amazon-Ratings), and the sampling uses exact k-subset with multiple samples which adds cost.
- The expressiveness results are partial: they preserve only most partial isomorphisms and are probabilistic.
Experimental setup
Datasets: QM9 ZINC OGBG-Molhiv TUDataset Peptides-func Peptides-struct PCQM-Contact Trees-NeighboursMatch Trees-LeafCount EXP CSL Cornell
Benchmarks: QM9 ZINC OGBG-Molhiv TUDataset Peptides-func Peptides-struct PCQM-Contact Trees-NeighboursMatch Trees-LeafCount EXP CSL Cornell
Metrics: MAE accuracy average precision training time per epoch memory usage layer-wise sensitivity total effective resistance
Baselines: GINE GPS PR-MPNN Drew SAT S2GCN AMP NBA GRIT Exphormer
Builds on: PR-MPNN SIMPLE GIN GINE GPS
Closest papers by meaning
- Probabilistically Rewired Message-Passing Neural Networks (ICLR 2024 · similarity 0.94)
- Locality-Aware Graph Rewiring in GNNs (ICLR 2024 · similarity 0.91)
- PANDA: Expanded Width-Aware Message Passing Beyond Rewiring (ICML 2024 · similarity 0.89)
- Towards Dynamic Message Passing on Graphs (NeurIPS 2024 · similarity 0.88)
- Improving the Effective Receptive Field of Message-Passing Neural Networks (ICML 2025 · similarity 0.88)
- Mitigating Over-Squashing in Graph Neural Networks by Spectrum-Preserving Sparsification (ICML 2025 · similarity 0.88)
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.