Geometry-Aware Edge Pooling for Graph Neural Networks

TL;DR

This paper introduces geometry-aware edge pooling layers for GNNs that preserve graph structure and diversity during pooling, leading to improved performance and interpretability across various graph classification tasks.

Contribution

The paper presents novel structure-aware pooling methods leveraging diffusion geometry and diversity measures, enhancing graph preservation and performance in GNNs.

Findings

Achieves top performance across diverse graph classification tasks.

Preserves key spectral properties of input graphs.

Maintains high accuracy across different pooling ratios.

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 discarding fundamental graph structures, thus reducing interpretability. This leads to unreliable performance across dataset types, downstream tasks and pooling ratios. Addressing these concerns, we propose novel graph pooling layers for structure-aware pooling via edge collapses. Our methods leverage diffusion geometry and iteratively reduce a graph's size while preserving both its metric structure and its 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 top 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.