Gumbel-MPNN: Graph Rewiring with Gumbel-Softmax
This work addresses performance limitations in graph neural networks for node classification, though it appears incremental as it builds on existing rewiring methods with a specific technique.
The paper tackles the problem of improving message-passing neural networks (MPNNs) for node classification by addressing inconsistencies in neighborhood distributions, proposing a Gumbel-Softmax-based rewiring method that enhances neighborhood informativeness and increases classification performance.
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.