DiffWire: Inductive Graph Rewiring via the Lovász Bound
Graph Neural Networks (GNNs) have been shown to achieve competitive results
to tackle graph-related tasks, such as node and graph classification, link
prediction and node and graph clustering in a variety of domains. Most GNNs use
a message passing framework and hence are called MPNNs. Despite their promising
results, MPNNs have been reported to suffer from over-smoothing, over-squashing
and under-reaching. Graph rewiring and graph pooling have been proposed in the
literature as solutions to address these limitations. However, most
state-of-the-art graph rewiring methods fail to preserve the global topology of
the graph, are not differentiable (inductive) and require the tuning of
hyper-parameters. In this paper, we propose DiffWire, a novel framework for
graph rewiring in MPNNs that is principled, fully differentiable and
parameter-free by leveraging the Lov\'asz bound. Our approach provides a
unified theory for graph rewiring by proposing two new, complementary layers in
MPNNs: first, CTLayer, a layer that learns the commute times and uses them as a
relevance function for edge re-weighting; second, GAPLayer, a layer to optimize
the spectral gap, depending on the nature of the network and the task at hand.
We empirically validate the value of our proposed approach and each of these
layers separately with benchmark datasets for graph classification. DiffWire
brings together the learnability of commute times to related definitions of
curvature, opening the door to the development of more expressive MPNNs.