Gumbel-MPNN: Graph Rewiring with Gumbel-Softmax

TL;DR

This paper introduces Gumbel-MPNN, a graph rewiring technique using Gumbel-Softmax that improves message-passing neural networks by enhancing neighborhood distribution informativeness and addressing issues like oversquashing.

Contribution

It proposes a novel Gumbel-Softmax-based rewiring method that reduces neighborhood distribution deviations and improves MPNN performance on node classification tasks.

Findings

Enhances neighborhood informativeness

Handles long-range dependencies effectively

Mitigates oversquashing in MPNNs

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 https://github.com/Bobowner/Gumbel-Softmax-MPNN.