Geometry-Aware Edge Pooling for Graph Neural Networks

Katharina Limbeck, Lydia Mezrag, Guy Wolf, Bastian Rieck

Abstract:

Graph Neural Networks (GNNs) have shown significant success for graph-based
tasks. Motivated by the prevalence of large datasets in real-world applications,
pooling layers are crucial components of GNNs. By reducing the size of input
graphs, pooling enables faster training and potentially better generalisation. However, existing pooling operations often optimise for the learning task at the expense
of fundamental graph structures and interpretability. This leads to unreliable
performance across varying dataset types, downstream tasks and pooling ratios.
Addressing these concerns, we propose novel graph pooling layers for structureaware pooling via edge collapses. Our methods leverage diffusion geometry and
iteratively reduce a graph’s size while preserving both its metric structure and
structural diversity. We guide pooling using magnitude, an isometry-invariant
diversity measure, which permits us to control the fidelity of the pooling process.
Further, we use the spread of a metric space as a faster and more stable alternative
ensuring computational efficiency. Empirical results demonstrate that our methods
(i) achieve superior performance compared to alternative pooling layers across a
range of diverse graph classification tasks, (ii) preserve key spectral properties of
the input graphs, and (iii) retain high accuracy across varying pooling ratios.