Geometry-Aware Edge Pooling for Graph Neural Networks

TL;DR

This work tackles the challenge of graph pooling in GNNs by introducing two geometry-aware edge-pooling layers, MagEdgePool and SpreadEdgePool, that preserve diffusion-based geometry and structural diversity through edge contractions guided by magnitude or spread. The authors establish key theoretical properties, including additivity on disjoint graphs, isomorphism invariance, and monotonicity under 1-Lipschitz contractions, and analyze expressivity and computational costs. Empirically, the methods achieve top or competitive performance on graph classification and regression tasks, while preserving spectral properties and maintaining stable accuracy across pooling ratios; SpreadEdgePool offers notable computational efficiency. Overall, the approach provides interpretable, non-trainable pooling that robustly encodes graphs’ geometry, with practical scalability for moderate-sized graphs and potential extensions to larger graphs and alternative distance measures.

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.

Figures (12)

• Figure 1: Examples of graphs pooled to approximately half their original size compared across pooling layers. Our proposed methods, MagEdgePool and SpreadEdgePool, respect the original graphs' geometry during pooling. Alternative approaches tend to obscure adjacency relationships to varying extents by creating counter-intuitive edges (Graclus, NMF), disconnecting entire portions of the graphs (TopK, SAGPool), or returning dense representations that do not preserve any geometric structure (DiffPool, MinCut).
• Figure 2: Illustrating our proposed pooling method, MagEdgePool, on a graph from the ENZYMES dataset across varying pooling ratios. Each edge is coloured by its magnitude difference, which measures the impact its contraction would have on the graph's structural diversity. Edges with low magnitude differences are most redundant for the graph's geometry and are collapsed first.
• Figure 3: Structure preservation for all graphs in the NCI1 dataset across varying pooling ratios. Left: The spectral distance between the normalised Laplacians of the original and the pooled graphs. Right: The relative difference in magnitude, which summarises the proportional difference in structural diversity after pooling. Violin plots show the variability across graphs at pooling ratio 0.5.
• Figure 4: Classification performance across varying the pooling ratio for different pooling layers. Pooling is applied as part of a GIN architecture. Results are shown for the ENZYME and NCI1 datasets. Lines show the mean and shaded areas the standard deviation of the test accuracy.
• Figure S.5: Comparison between magnitude and spread computed from diffusion distances for all graphs in three graph datasets, NCI1, ENZYMES and IMDB-Multi.

Theorems & Definitions (23)

• Theorem 1
• Theorem 2
• Theorem 3
• Theorem 4
• Theorem 5
• Theorem 5
• proof
• Theorem 5
• proof
• Theorem 5