Table of Contents
Fetching ...

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.

Graph Pooling via Ricci Flow

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.
Paper Structure (65 sections, 11 theorems, 22 equations, 8 figures, 16 tables)

This paper contains 65 sections, 11 theorems, 22 equations, 8 figures, 16 tables.

Key Result

Corollary 1

Consider a simple architecture with a block of MP base layers, followed by a ORC-Pool layer. Let $G_1, G_2$ denote two 1-WL-distinguishable graphs with node attributes $X_1, X_2$. Further let $X_1' \neq X_2'$ denote the node representations learned by the block of MP layers. Then the coarsened graph

Figures (8)

  • Figure 1: Clustering based on connectivity vs. node attributes only.
  • Figure 2: Discrete Ricci curvature reveals coarse structure. Left: Relation between curvature and trajectories of diffusion processes starting at similar (green) and dissimilar (red) nodes. Right: Dumbbell graph with uniform edge weights (at initialization) and with curvature-adjusted edge weights (darker colors represent lower curvature).
  • Figure 3: Proposed geometric pooling operator (ORC-Pool), which utilizes a curvature-based, geometric selection function (Sel) to identify supernodes and superedges (Red), which are then reconnected to generate the pooled graph (Con).
  • Figure 4: $G_{a,b}$ ($a=b=3$)
  • Figure 5: Evolution of edge weights under Ricci flow in the Stochastic Block Model.
  • ...and 3 more figures

Theorems & Definitions (16)

  • Corollary 1
  • Definition 1
  • Lemma 1: informal, ni2019community
  • Theorem 1: informal
  • Theorem 2: Lower bound (tian2023curvature, Thm. 4.6)
  • Theorem 3: Upper bound (tian2023curvature, Thm. 4.6)
  • Corollary 2: Expressivity of ORC-Pool
  • proof
  • Theorem 4: bianchi2023expressive, Thm. 1
  • Lemma 2
  • ...and 6 more