# Gumbel-MPNN: Graph Rewiring with Gumbel-Softmax

arXiv: https://arxiv.org/abs/2508.17531 · HTML: https://mlsanity.com/paper/2508.17531/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 43 of its 54 references that are in the map.
- First posted on arXiv 24 Aug 2025.

## Builds on (most cited first, 30 of 43)

- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Pitfalls of Graph Neural Network Evaluation](https://mlsanity.com/paper/1811.05868/) (2018)
- [Representation Learning on Graphs with Jumping Knowledge Networks](https://mlsanity.com/paper/1806.03536/) (2018)
- [Simple and Deep Graph Convolutional Networks](https://mlsanity.com/paper/2007.02133/) (2020)
- [Geom-GCN: Geometric Graph Convolutional Networks](https://mlsanity.com/paper/2002.05287/) (2020)
- [DropEdge: Towards Deep Graph Convolutional Networks on Node Classification](https://mlsanity.com/paper/1907.10903/) (2019)
- [Adaptive Universal Generalized PageRank Graph Neural Network](https://mlsanity.com/paper/2006.07988/) (2020)
- [Multi-scale Attributed Node Embedding](https://mlsanity.com/paper/1909.13021/) (2019)
- [Understanding over-squashing and bottlenecks on graphs via curvature](https://mlsanity.com/paper/2111.14522/) (2021)
- [Beyond Low-frequency Information in Graph Convolutional Networks](https://mlsanity.com/paper/2101.00797/) (2021)
- [A critical look at the evaluation of GNNs under heterophily: Are we really making progress?](https://mlsanity.com/paper/2302.11640/) (2023)
- [Learning Discrete Structures for Graph Neural Networks](https://mlsanity.com/paper/1903.11960/) (2019)
- [Neural Sheaf Diffusion: A Topological Perspective on Heterophily and Oversmoothing in GNNs](https://mlsanity.com/paper/2202.04579/) (2022)
- [Graph Neural Networks with Heterophily](https://mlsanity.com/paper/2009.13566/) (2020)
- [Two Sides of the Same Coin: Heterophily and Oversmoothing in Graph Convolutional Neural Networks](https://mlsanity.com/paper/2102.06462/) (2021)
- [Is Homophily a Necessity for Graph Neural Networks?](https://mlsanity.com/paper/2106.06134/) (2021)
- [Finding Global Homophily in Graph Neural Networks When Meeting Heterophily](https://mlsanity.com/paper/2205.07308/) (2022)
- [How Powerful are Spectral Graph Neural Networks](https://mlsanity.com/paper/2205.11172/) (2022)
- [Revisiting Heterophily For Graph Neural Networks](https://mlsanity.com/paper/2210.07606/) (2022)
- [Breaking the Limit of Graph Neural Networks by Improving the Assortativity of Graphs with Local Mixing Patterns](https://mlsanity.com/paper/2106.06586/) (2021)
- [Edge Directionality Improves Learning on Heterophilic Graphs](https://mlsanity.com/paper/2305.10498/) (2023)
- [Differentiable Graph Module (DGM) for Graph Convolutional Networks](https://mlsanity.com/paper/2002.04999/) (2020)
- [Classic GNNs are Strong Baselines: Reassessing GNNs for Node Classification](https://mlsanity.com/paper/2406.08993/) (2024)
- [GBK-GNN: Gated Bi-Kernel Graph Neural Networks for Modeling Both Homophily and Heterophily](https://mlsanity.com/paper/2110.15777/) (2021)
- [DRew: Dynamically Rewired Message Passing with Delay](https://mlsanity.com/paper/2305.08018/) (2023)
- [Towards Robust Graph Neural Networks for Noisy Graphs with Sparse Labels](https://mlsanity.com/paper/2201.00232/) (2022)

## Abstract (arXiv)

Graph homophily has been considered an essential property for message-passing neural networks (MPNN) in node classification. Recent findings suggest that performance is more closely tied to the consistency of neighborhood class distributions. We demonstrate that the MPNN performance depends on the number of components of the overall neighborhood distribution within a class. By breaking down the classes into their neighborhood distribution components, we increase measures of neighborhood distribution informativeness but do not observe an improvement in MPNN performance. We propose a Gumbel-Softmax-based rewiring method that reduces deviations in neighborhood distributions. Our results show that our new method enhances neighborhood informativeness, handles long-range dependencies, mitigates oversquashing, and increases the classification performance of the MPNN. The code is available at https://github.com/Bobowner/Gumbel-Softmax-MPNN.