Gumbel-MPNN: Graph Rewiring with Gumbel-Softmax

Marcel Hoffmann, Lukas Galke Poech, A. Scherp

Abstract

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.

Benchmark Results

Tested on 7 datasets 18 baseline values reported too low

Tables show presented results in global context (best 3 papers + GCN + MLP).

Classic

CiteSeer

Global leaderboard →

#	Model	Accuracy ↑	Paper
1	CNA code ↗ GNN global	0.9575 ±0.0058	Graph Neural Networks Need Cluster-Normalize-Activate Modules code ↗
2	IS-GIB code ↗ Hyb global	0.9390 ±0.0187	Individual and Structural Graph Information Bottlenecks for Out-of-Distribution Generalization code ↗
3	EERM code ↗ GNN global	0.9112 ±0.0145	Individual and Structural Graph Information Bottlenecks for Out-of-Distribution Generalization code ↗
35	GAT GNN baseline actual #35 vs paper #124 89 places worse	0.8080 ±0.0026 best reported	actual #35 vs paper #124 89 places worse paper 0.7643 ±0.0010 best at pub: Feature Expansion for Graph Neural Networks (ICML 2023) paper 0.7643 ±0.0010 Δ at pub 0.0437
44	GCN GNN baseline actual #44 vs paper #131 87 places worse	0.8033 ±0.0088 best reported	actual #44 vs paper #131 87 places worse paper 0.7610 ±0.0015 best at pub: NTFormer: A Composite Node Tokenized Graph Transformer for Node Classification (2024) paper 0.7610 ±0.0015 Δ at pub 0.0423
61	GraphSAGE GNN baseline actual #61 vs paper #130 69 places worse	0.7923 ±0.0053 best reported	actual #61 vs paper #130 69 places worse paper 0.7612 ±0.0014 best at pub: Gophormer: Ego-Graph Transformer for Node Classification (2021) paper 0.7612 ±0.0014 Δ at pub 0.0311
131	SDRF with GCN Hyb baseline	0.7612 ±0.0014	This paper
137	GCN with DropEdge GNN baseline	0.7599 ±0.0015	This paper
155	Gumbel-MPNN Hyb ★ proposed	0.7554 ±0.0017	This paper
259	SAGE code ↗ GNN global	0.7361 ±0.0190	Efficient Tuning and Inference for Large Language Models on Textual Graphs code ↗
350	MLP Trad baseline actual #350 vs paper #350	0.7289 ±0.0010	actual #350 vs paper #350 paper 0.7289 ±0.0010 prev best: On the Equivalence of Graph Convolution and Mixup (TMLR 2023) Improvement within reported uncertainty; not highlighted. paper 0.7289 ±0.0010 Δ at pub 0.0099

Heterophilic Graphs

Actor

Global leaderboard →

#	Model	Accuracy ↑	Paper
1	DiGGR Hyb global	0.4535 ±0.0353	Disentangled Generative Graph Representation Learning
2	TRIGON code ↗ Hyb global	0.4402 ±0.0028	Dynamic Triangulation-Based Graph Rewiring for Graph Neural Networks code ↗
3	DoG code ↗ ext data GNN global	0.4320 ±0.0120	Diffusion on Graph: Augmentation of Graph Structure for Node Classification code ↗
42	MLP Trad baseline actual #42 vs paper #161 119 places worse	0.4026 ±0.0108 best reported	actual #42 vs paper #161 119 places worse paper 0.3650 ±0.0013 best at pub: Heterophily-Aware Graph Attention Network (2023) paper 0.3650 ±0.0013 Δ at pub 0.0376
68	GraphSAGE GNN baseline actual #68 vs paper #145 77 places worse	0.3899 ±0.0085 best reported	actual #68 vs paper #145 77 places worse paper 0.3679 ±0.0010 best at pub: A Quasi-Wasserstein Loss for Learning Graph Neural Networks (WWW 2023) paper 0.3679 ±0.0010 Δ at pub 0.0220
128	Gumbel-MPNN Hyb ★ proposed	0.3713 ±0.0010	This paper
135	GCN GNN baseline actual #135 vs paper #223 88 places worse	0.3690 best reported	actual #135 vs paper #223 88 places worse paper 0.3521 ±0.0007 best at pub: Beyond Low-frequency Information in Graph Convolutional Networks (AAAI 2021) paper 0.3521 ±0.0007 Δ at pub 0.0169
182	GAT GNN baseline actual #182 vs paper #236 54 places worse	0.3605 ±0.0035 best reported	actual #182 vs paper #236 54 places worse paper 0.3491 ±0.0012 best at pub: Hierarchical Graph Transformer with Adaptive Node Sampling (NeurIPS 2022) paper 0.3491 ±0.0012 Δ at pub 0.0114
226	SDRF with GCN Hyb baseline	0.3515 ±0.0005	This paper
266	GIN code ↗ GNN global	0.3410 ±0.0030	What Can We Learn From MIMO Graph Convolutions? code ↗
279	GCN with DropEdge GNN baseline	0.3359 ±0.0010	This paper

Heterophily Benchmark

Roman-empire

Global leaderboard →

#	Model	Accuracy ↑	Paper
1	Dir-Poly code ↗ GNN global	0.9451 ±0.0022	Graph Transductive Sharpening: Leveraging Unlabeled Predictions in Node Classification code ↗
2	D-CDE code ↗ GNN global	0.9387 ±0.0041	Distributed-Order Fractional Graph Operating Network code ↗
3	LargeScaleNet code ↗ Hyb global	0.9358 ±0.0024	Scale-aware Message Passing For Graph Node Classification code ↗
26	GraphSAGE GNN baseline actual #26 vs paper #82 56 places worse	0.9106 ±0.0027 best reported	actual #26 vs paper #82 56 places worse paper 0.8178 ±0.0006 best at pub: Directed Semi-Simplicial Learning with Applications to Brain Activity Decoding (2025) paper 0.8178 ±0.0006 Δ at pub 0.0928
64	SAGE GNN global	0.8574 ±0.0067	TANGO: Graph Neural Dynamics via Learned Energy and Tangential Flows
77	GCN GNN baseline actual #77 vs paper #97 20 places worse	0.8272 ±0.0082 best reported	actual #77 vs paper #97 20 places worse paper 0.7882 ±0.0009 best at pub: HyperAggregation: Aggregating over Graph Edges with Hypernetworks (2024) paper 0.7882 ±0.0009 Δ at pub 0.0390
79	Gumbel-MPNN Hyb ★ proposed	0.8223 ±0.0007	This paper
83	GAT GNN baseline actual #83 vs paper #83	0.8145 ±0.0009	actual #83 vs paper #83 paper 0.8145 ±0.0009 prev best: Graph Neural Networks for Graphs with Heterophily: A Survey (2022) paper 0.8145 ±0.0009 Δ at pub 0.0058
85	SDRF with GCN Hyb baseline	0.8060 ±0.0007	This paper
101	GCN with DropEdge GNN baseline	0.7826 ±0.0011	This paper
135	MLP Trad baseline actual #135 vs paper #140 5 places worse	0.6665 ±0.0047 best reported	actual #135 vs paper #140 5 places worse paper 0.6588 ±0.0004 best at pub: HyperAggregation: Aggregating over Graph Edges with Hypernetworks (2024) paper 0.6588 ±0.0004 Δ at pub 0.0077

Amazon-ratings

Global leaderboard →

#	Model	Accuracy ↑	Paper
1	UniGAP global	0.6375 ±0.0110	UniGAP: A Universal and Adaptive Graph Upsampling Approach to Mitigate Over-Smoothing in Node …
2	H2GCN + HalfHop GNN global	0.5961 ±0.0100	UniGAP: A Universal and Adaptive Graph Upsampling Approach to Mitigate Over-Smoothing in Node …
3	GraphSAGE + HalfHop GNN global	0.5897 ±0.0150	UniGAP: A Universal and Adaptive Graph Upsampling Approach to Mitigate Over-Smoothing in Node …
12	SAGE code ↗ GNN global	0.5569 ±0.0023	Graph Transductive Sharpening: Leveraging Unlabeled Predictions in Node Classification code ↗
14	GAT GNN baseline actual #14 vs paper #84 70 places worse	0.5056 ±0.0006	actual #14 vs paper #84 70 places worse paper 0.5056 ±0.0006 prev best: Piecewise Constant Spectral Graph Neural Network (TMLR 2025) best today: Graph Transductive Sharpening: Leveraging Unlabeled Predictions in Node Classification (2026) Improvement within reported uncertainty; not highlighted. paper 0.5056 ±0.0006 Δ at pub 0.0028 Δ today 0.0503
15	GraphSAGE GNN baseline actual #15 vs paper #83 68 places worse	0.5540 ±0.0020 best reported	actual #15 vs paper #83 68 places worse paper 0.5061 ±0.0010 best at pub: UniGAP: A Universal and Adaptive Graph Upsampling Approach to Mitigate Over-Smoothing in Node Classification Tasks (2024) paper 0.5061 ±0.0010 Δ at pub 0.0479
24	GCN GNN baseline actual #24 vs paper #98 74 places worse	0.5411 ±0.0034 best reported	actual #24 vs paper #98 74 places worse paper 0.4965 ±0.0013 best at pub: UniGAP: A Universal and Adaptive Graph Upsampling Approach to Mitigate Over-Smoothing in Node Classification Tasks (2024) best today: Graph Transductive Sharpening: Leveraging Unlabeled Predictions in Node Classification (2026) paper 0.4965 ±0.0013 Δ at pub 0.0415 Δ today 0.0446
84	MLP Trad baseline actual #84 vs paper #134 50 places worse	0.5058 best reported	actual #84 vs paper #134 50 places worse paper 0.4668 ±0.0013 best at pub: Understanding Heterophily for Graph Neural Networks (ICML 2024) paper 0.4668 ±0.0013 Δ at pub 0.0390
91	Gumbel-MPNN Hyb ★ proposed	0.5008 ±0.0006	This paper
96	GCN with DropEdge GNN baseline	0.4984 ±0.0013	This paper
158	SDRF with GCN Hyb baseline	0.4403 ±0.0018	This paper

Minesweeper

Global leaderboard →

#	Model	ROC-AUC ↑	Paper
1	BuNN code ↗ Hyb global	0.9899 ±0.0016	Bundle Neural Networks for message diffusion on graphs code ↗
2	DC-GNN GNN global	0.9850 ±0.0021	Differentiable Cluster Graph Neural Network
3	Tango -GPS global	0.9839 ±0.0054	TANGO: Graph Neural Dynamics via Learned Energy and Tangential Flows
36	GAT GNN baseline actual #36 vs paper #55 19 places worse	0.9238 ±0.0009	actual #36 vs paper #55 19 places worse paper 0.9238 ±0.0009 prev best: Graph Neural Networks for Graphs with Heterophily: A Survey (2022) best today: Relieving the Over-Aggregating Effect in Graph Transformers (NeurIPS 2025) Improvement within reported uncertainty; not highlighted. paper 0.9238 ±0.0009 Δ at pub 0.0037 Δ today 0.0153
42	GraphSAGE GNN baseline actual #42 vs paper #42	0.9358 ±0.0005	actual #42 vs paper #42 paper 0.9358 ±0.0005 prev best: Polynormer: Polynomial-Expressive Graph Transformer in Linear Time (ICLR 2024) Improvement within reported uncertainty; not highlighted. paper 0.9358 ±0.0005 Δ at pub 0.0007
46	SAGE GNN global	0.9351 ±0.0057	Beyond Oversquashing: Understanding Signal Propagation in GNNs Via Observables
68	GCN GNN baseline actual #68 vs paper #68	0.9150 ±0.0006	actual #68 vs paper #68 paper 0.9150 ±0.0006 prev best: Graph Neural Networks for Graphs with Heterophily: A Survey (2022) paper 0.9150 ±0.0006 Δ at pub 0.0175
70	Gumbel-MPNN Hyb ★ proposed	0.9118 ±0.0013	This paper
86	GCN with DropEdge GNN baseline	0.8995 ±0.0007	This paper
87	SDRF with GCN Hyb baseline	0.8993 ±0.0006	This paper
156	MLP Trad baseline actual #156 vs paper #160 4 places worse	0.5099 ±0.0147 best reported	actual #156 vs paper #160 4 places worse paper 0.5038 ±0.0012 best at pub: Evolving Computation Graphs (ICML 2023) best today: Graph Homophily Booster: Rethinking the Role of Discrete Features on Heterophilic Graphs (2025) Difference within reported uncertainty; not counted as underperformance. paper 0.5038 ±0.0012 Δ at pub 0.0051 Δ today 0.0061

Tolokers

Global leaderboard →

#	Model	ROC-AUC ↑	Paper
1	Polynormer-r code ↗ GT global	0.8591 ±0.0074	Polynormer: Polynomial-Expressive Graph Transformer in Linear Time code ↗
2	DC-GNN GNN global	0.8588 ±0.0081	Differentiable Cluster Graph Neural Network
3	TANGO-GPS Hyb global	0.8566 ±0.0101	TANGO: Graph Neural Dynamics via Learned Energy and Tangential Flows
28	GAT GNN baseline actual #28 vs paper #85 57 places worse	0.8443 ±0.0100 best reported	actual #28 vs paper #85 57 places worse paper 0.8099 ±0.0017 best at pub: TopoBench: A Framework for Benchmarking Topological Deep Learning (2024) paper 0.8099 ±0.0017 Δ at pub 0.0344
50	GCN GNN baseline actual #50 vs paper #88 38 places worse	0.8364 ±0.0067 best reported	actual #50 vs paper #88 38 places worse paper 0.8075 ±0.0012 best at pub: Graph Neural Networks for Graphs with Heterophily: A Survey (2022) best today: Beyond Oversquashing: Understanding Signal Propagation in GNNs Via Observables (2026) paper 0.8075 ±0.0012 Δ at pub 0.0289 Δ today 0.0289
77	GraphSAGE GNN baseline actual #77 vs paper #93 16 places worse	0.8243 ±0.0044 best reported	actual #77 vs paper #93 16 places worse paper 0.8016 ±0.0012 best at pub: Polynormer: Polynomial-Expressive Graph Transformer in Linear Time (ICLR 2024) best today: Relieving the Over-Aggregating Effect in Graph Transformers (NeurIPS 2025) paper 0.8016 ±0.0012 Δ at pub 0.0227 Δ today 0.0227
76	SAGE GNN global	0.8243 ±0.0044	Beyond Oversquashing: Understanding Signal Propagation in GNNs Via Observables
81	Gumbel-MPNN Hyb ★ proposed	0.8156 ±0.0041	This paper
82	GCN with DropEdge GNN baseline	0.8149 ±0.0011	This paper
86	SDRF with GCN Hyb baseline	0.8090 ±0.0011	This paper
139	MLP Trad baseline actual #139 vs paper #139	0.7362 ±0.0020	actual #139 vs paper #139 paper 0.7362 ±0.0020 prev best: Evolving Computation Graphs (ICML 2023) Improvement within reported uncertainty; not highlighted. paper 0.7362 ±0.0020 Δ at pub 0.0067

Questions

Global leaderboard →

#	Model	ROC-AUC ↑	Paper
1	NID_GAT code ↗ global	0.9703 ±0.0002	Node Identifiers: Compact, Discrete Representations for Efficient Graph Learning code ↗
2	NID_GCN code ↗ global	0.9685 ±0.0010	Node Identifiers: Compact, Discrete Representations for Efficient Graph Learning code ↗
3	UniGAP global	0.8421 ±0.0080	UniGAP: A Universal and Adaptive Graph Upsampling Approach to Mitigate Over-Smoothing in Node …
16	GCN GNN baseline actual #16 vs paper #62 46 places worse	0.7902 ±0.0060 best reported	actual #16 vs paper #62 46 places worse paper 0.7607 ±0.0024 best at pub: UniGAP: A Universal and Adaptive Graph Upsampling Approach to Mitigate Over-Smoothing in Node Classification Tasks (2024) paper 0.7607 ±0.0024 Δ at pub 0.0295
26	GAT GNN baseline actual #26 vs paper #88 62 places worse	0.7835 ±0.0116 best reported	actual #26 vs paper #88 62 places worse paper 0.7254 ±0.0010 best at pub: Node Identifiers: Compact, Discrete Representations for Efficient Graph Learning (ICLR 2024) paper 0.7254 ±0.0010 Δ at pub 0.0581
46	GraphSAGE GNN baseline actual #46 vs paper #73 27 places worse	0.7721 ±0.0130 best reported	actual #46 vs paper #73 27 places worse paper 0.7538 ±0.0009 best at pub: UniGAP: A Universal and Adaptive Graph Upsampling Approach to Mitigate Over-Smoothing in Node Classification Tasks (2024) Difference within reported uncertainty; not counted as underperformance. paper 0.7538 ±0.0009 Δ at pub 0.0183
55	Gumbel-MPNN Hyb ★ proposed	0.7661 ±0.0012	This paper
58	SAGE GNN global	0.7644 ±0.0062	Beyond Oversquashing: Understanding Signal Propagation in GNNs Via Observables
96	MLP Trad baseline actual #96 vs paper #96	0.7082 ±0.0010	actual #96 vs paper #96 paper 0.7082 ±0.0010 prev best: Evolving Computation Graphs (ICML 2023) Improvement within reported uncertainty; not highlighted. paper 0.7082 ±0.0010 Δ at pub 0.0048
143	SDRF with GCN Hyb baseline	0.5716 ±0.0014	This paper
146	GCN with DropEdge GNN baseline	0.5592 ±0.0011	This paper

SOTA performance (not significantly different from the best) Proposed method (this paper) Baseline / comparison (this paper) Global context (other papers) Best at pub — reported result worse than the best known when published (best reported in green) Best at pub — reported value improved over the best known when published Best today — OK at publication; today's leaderboard best is elsewhere