Graph Pooling via Ricci Flow
Amy Feng, Melanie Weber
TL;DR
The paper addresses the need for effective pooling in graph neural networks that accounts for both graph topology and node attributes. It introduces ORC-Pool, a trainable pooling operator that leverages Ollivier's discrete Ricci curvature and a discrete Ricci-flow to perform curvature-based coarsening, integrated as a pooling layer within MPGNNs. The authors provide theoretical properties (permutation-invariance and expressivity) and demonstrate empirical gains in node clustering and graph classification, while also proposing scalable curvature approximations. The approach expands curvature-based clustering to attributed graphs, enabling multi-scale, geometry-informed pooling with practical implications for deep GNN architectures. Overall, ORC-Pool offers a principled, geometry-driven mechanism to coarsen attributed graphs while preserving and enhancing representational power.
Abstract
Graph Machine Learning often involves the clustering of nodes based on similarity structure encoded in the graph's topology and the nodes' attributes. On homophilous graphs, the integration of pooling layers has been shown to enhance the performance of Graph Neural Networks by accounting for inherent multi-scale structure. Here, similar nodes are grouped together to coarsen the graph and reduce the input size in subsequent layers in deeper architectures. In both settings, the underlying clustering approach can be implemented via graph pooling operators, which often rely on classical tools from Graph Theory. In this work, we introduce a graph pooling operator (ORC-Pool), which utilizes a characterization of the graph's geometry via Ollivier's discrete Ricci curvature and an associated geometric flow. Previous Ricci flow based clustering approaches have shown great promise across several domains, but are by construction unable to account for similarity structure encoded in the node attributes. However, in many ML applications, such information is vital for downstream tasks. ORC-Pool extends such clustering approaches to attributed graphs, allowing for the integration of geometric coarsening into Graph Neural Networks as a pooling layer.
