Graph Pooling via Ricci Flow
Authors: Amy Feng
Melanie Weber
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (12/2024)
Graph Pooling via Ricci Flow
Amy Feng
afengS@college.havvrrr edu
Horvard College
Melanie Weber
wwweerr@sea..harvarree edu
Horvard University
Reviewed on OpenReview: https: // openreview. net/ forum? dd=apBipWWPuk
Abstract
Graph Machine Learning often involves the clustering of nodes based on similarity structure encoded 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. 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 '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 pooling layer.
1 Introduction
Multi-scale structure is ubiquitous in relational data across domains; examples include complex molecules in computational biology, systems of interacting particles in physics, as well as complex financial and social systems. Many graph learning tasks, such as clustoring and coarsening, rely on such structure. Graph Neural Networks (GNNs) account for multi-scale structure via pooling layers, which form crucial building blocks many state of the art architectures.
Node clustering often utilizes implicit or explicit pooling operation, which decomposes given graph into densely connected subgraphs by grouping similar nodes. Applied at different scales, aggregating information across subgraphs allows for uncovering multi-scale structure (coarsening). A number of classical algorithms implement such operations, including spectral clustering 22 M the Louvain algorithm (IT and Graclus 14
Recently, clustering approaches based on graph curvature (41
50):
59
19] have received interest. They utilize discrete Ricci curvature 4
23]
to characterize the local geometry of graph by aggregating geometric information across 2-hop node neighborhoods. An associated curvature 22 characterizes graph geometry at more global scale, uncovering its coarse geometry.
Figure 1: Clustering based on connectivity node attributes only.
flow
A crucial shortcoming of classical node clustering approaches the case aftributed graphs is their inability capture structural information encoded the node attributes, as they evaluate only similarity structure encoded in the graph's i's topology. However, such information is often vital for downstream tasks. Applying

---

Published in Transactions on Machine Learning Research (12/2024)
separate clustering approach to the (typically vector-valued) node attributes does not resolve this issue, it is blind to similarity structure encoded in the graph's 's topology, which is often equally vital for downstream tasks (Fig. D Hence, to capture meaningful clusters in attributed graphs, pooling operators need to evaluate both graph topology and node attributes. Recent literature has introduced a plethora of pooling operators for altributed graphs
many which are based on the classical algorithuns listed above. In this work, we extend, for the first time, clustering approaches based on Ricci flow to attributed graphs.
Related Work. The design of pooling layers has historically build on clustering algorithms, such GRACLUS (4 MINCUTPOOL (8. A.A plethora of pooling operators has been proposed in recent literature, inspired by spectral clustering
32 matrix factorization (1), hierarchical clustering (65), modularity (39) and multi-set encoding (5), among others. The Secrect-Reuue-CCe Comanct (short: SRC) framework (28) serves as a unifying framework and benchmark for pooling operators. Most closely related to our work is a pooling approach proposed by (47), which performs edge cuts guided by Ricei curvature, instead a curvature- adjustment computed via Ricci flow. Among other differences, their approach does not incorporate node attributes, which can lead to reduced performance on attributed graphs. We provide a detailed comparison of the two approaches (conceptual and experimental) in Apx. C
Summary of contributions. The main contributions of the paper are as follows:
1. We introduce ORC-PooL, a trainable pooling operator, which utilizes discrete Ricci curvature and an associated geometric flow to identify salient multi-scale structure graphs. ORC-PooL groups nodes according to similarity structure encoded both in the graph's topology, as well as its nodes' attributes, presenting the first extension of Rieci flow based clustering to attributed graphs.
2. We further introduce a pooling layer, which allows for incorporating ORC-PooL into Passsgg-Paasiigg Graph Neural Networks.
3. We perform range of computational experiments, which demonstrate the utility of ORC-Poor. layers through improvements in node- and graph-level tasks. We complement our empirical results with a discussion of the structural properties of ORC-POOL.
2 Background and Notation
Let G = (V, E} denote a graph, V ([V)= N) the set of nodes and E = V V the set of edges; G is endowed with the usual shortest-path metric dg. We assume that is attributed and denote node attributes as X x/ We further assume that G weighted; i.c., its edges are endowed with scalar attributes, given by a weight function
+ R+, Below, we recall basic Graph Neural Networks concepts and introduce the graph curvature notions and associated curvature flow utilized in this work.
2.1 Graph Neural Networks with Pooling
Message-passing Graph Neural Networks. The blueprint of many state of the art architectures are Message- GNNs (MPGNNs) 2 BD), which learn node embeddings via iterative imessage-passingg (MP) scheme: Node representations are iteratively updated Node attributes in the input graph determine the node representations at initialization. Formally, MPGNNs (t+1) and an update function fup (1 O....,, L-1 denoting the layer):
function of their neighbors' representations.
consist of a message m
(1+1) mf+1) 1) (1)
Agg (n..), (h(() | u E.N.)))
(++1) = fup (1()), mmmm (+1)).
fAgg fup may be implemented via MLPs with trainable parameters. Popular examples of MPGNNs are GCN (4) GIN (3) GraphSage 30 and GAT (66)
Graph pooling operators. State of the art GNN architectures often combine MP base layers with pooling layers. The selection, reduction, and connection (short: SRC) framework (28) formalizes pooling operators
2

---

Published in Transactions on Machine Learning Research (12/2024)
maps POOl: G G = (XP EE) composed of three functions that act the nodes, node attributes and edges of G, respectively (Fig. 3):
1. A selection function, which identifies a set of nodes that are merged to supernodes (V- S (I,...,,NN) NN): SEL: (I....;N) +>
(x))... The selection function lies at the heart of the pooling operator. Below, we propose a geometric selection function, which utilizes the graph's geometry for node-level clustering.
2. A reduction function, which computes the attributes of supernodes V% by aggregating the attributes of its nodes: RED: : = (..):...))kExxx + {2
3. A connection function, which determines the connectivity of supernodes and (ro-)an Jassigns edge attributes: CON :
((). Vi), <x))
SRC is unifying framework for commonly used pooling operatons and has been utilized in graph machine learning libraries such as SPEKTRAI (27) and PYTORCH GROMETRIC([].
2.2 Graph Curvature
Classically, Ricci curvature establishes a connection between the local dispersion of geodesics and the local curvature of manifold, deriving from crucial connection between volume growth rates and curvature: Negative curvature characterizes exponential fast volume growth, while positive curvature indicates contraction. Several analogous notions have been proposed for discrete spaces, including by Ollivier (43), Forman (23) and Erbar-Maas (8) Common among all notions is the observation that edges, which encode long-range dependencies, have low (negative) curvature, while edges that form local connections have high curvature, allowing for characterization of local and global connectivity patterns in a graph. Here, we focus on Ollivier's curvature, which we introduce below.
2.2.1 Ollivier's Ricci curvature
Ollivier 43) introduces a Ricci curvature, which relates the curvature along a geodesic between nearby points 2, on manifold M with the transportation distance between their neighborhoods. Recall that the Wasscrstein-1 distance between two probability distributions M1. is given by
/ (dm(z,))) (a( v) dx dy.
W2() = inf pertti
(1)
where dm v) denotes the geodesic distance, (x) = M dm(r, 2) S Se the e-neighborhood of and M2) the set of measures with marginals U1. Further let 2e (2) denote measure on this neighborhood. Then Ollivier's 's Ricci curvature (ORC) is given by KM((2 v) = = W(a
Analogous to the continuous case, ORC can be defined in discrete spaces: Let B denote matrix with entries by = e- inverse proportional to the edge weights and DB denote a matrix with entries d; = E, bij- Consider a diffusion process on the graph G, i.e, we place a point mass 6, at a node v and consider the diffusion (t) = & (D (12). The distribution generated by (t) defines measure the t-neighborhood of v. We define a discrete analog of ORC [33 26) for an edge e = (u, v) via the transportation distances between neighborhood measures defined by diflusion processes starting at its adjacent vertices, i.e.,
dmla.y)
W (p.((()).p.()) de(u, v)
(2)
How does this quantity relate to the coarse structure of the graph? In most graphs, similar nodes form densely connected subgraphs (homophily principle). As result, if are similar, then the diffusion processes will stay nearby with high probability, exploring the densely connected subgraph. In contrast, if u, are dissimilar, e.g., belong to separate clusters, then the diffusion processes are likely to explore separate clusters, drawing apart quickly. The transportation distance (p. Pt.), (t) much higher in the second case than the first. Thus, edges that connect dissimilar nodes have low (usually negative) curvature, whereas edges that connect similar nodes have high curvature. Hence, bridges between communities may be identified via low curvature
3

---

Published in Transactions on Machine Learning Research (12/2024)
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-adjuste edge weights (darker colors represent lower curvature).
(see Fig. 2 To ensure computational feasibility, we restrict ourselves to a first-order approximation of the diffusion process, Ps = al + (1 - D B, which defines a mcasure on the 1-hop node neighborhoods. With this choice, we recover the classical discrete Ollivier-Ricci curvature 43 36
We note that the pooling operations introduced below naturally extend to curvature computed over t-hop neighborhoods. Varying the neighborhood radius could provide an additional avenue for incorporating multi-scale structure into the curvature computation, which could be valuable in practise, albeit at higher computational cost. Notice that the curvature notion introduced above does not account for node attributes. We may define edge weights that encode a similarity measure on the node attributes, e.g., Wij m Ex (1)z2f)).
2.2.2 Curvature-based Clustering
Ricci Flow. In continuous spaces, an associated geometric flow (Ricci flow) is of great importance, as it reveals deep connections between the geometry and topology of the manifold (geometrization conjecture). As loose analogy, a geometric flow associated with discrete Ricei curvature can be defined, which has been related the community structure (II) and coarse geometry (II) of graphs. Ollivier (42) proposes a curvature flow
d do(u, v)( = (t) do (u, (t) ((u,v) E E),
dt
associated with discrete ORC along edges (u, v). discretization of this flow gives combinatorial evolution equation, which evolves edge weights according to the local geometry of the graph: Consider a family of weighted graphs G = (V, E, WJ, which is constructed from an input graph G Go by evolving its edge weights (setting dt = 1)
w f- (1 Kuu)ddo v)
(u, E)
(3)
where Kuv, do(u, v) are computed on G Eq. may be viewed as a discrete analogue of continuous Rieci flow. To control the scale of the edge weights, we normalize after each iteration.
Geometric clustering algorithms. Geometric clustering approaches based on ORC for non-attributed graphs have previously been considered, e.g. (50) 41 54). These algorithms partition-based, i.e., remove edges to decompose the graph into subgraphs, and come in two flavors: The first takes a local perspective, cutting edges with ORC below given threshold (c.g., (50)). The second evolves edge weights under Ricci flow (Eq. 3) encoding local and more global geometric information into the edge weights, before removing edges with weight below a certain threshold (e.g., (II)). Both approaches exploit the observation that ORC flow "highlights' the multi-scale structure of the graph (o.g., community structure). We emphasize that curvature, by construction, evaluates only the graph's connectivity, but cannot account for node attributes (apart from carefully chosen initialization of edge weights, above).
4

---

Published in Transactions on Machine Learning Research (12/2024)
2.2.3 Curvature Approximation
Computing ORC as defined above involves solving an optimal transport problem (i.e., computing the Wasserstein-1 distances between measures on the 1-hop neighborhoods of the adjacent vertices). In the discrete setting, this corresponds to computing the earth mover's distance, which, in the worst case, has complexity O(EE)m²) (where m denotes the maximum degree of nodes in G). A commonly used approximation via Sinkhorn's 's algorithm has reduced complexity of F()Fmmm², which can stiil be prohibitively expensive on largo-scale graphs. Recently, Tian et al. (54) introduced combinatorial approximation of ORC of the form Ruv :- 3 (m + mlor), where KdP and klo are combinatorial upper and lower bounds, which can be computed in O(E m)). More details, well as the formal statement of the bounds can be found in Apx. [ In large or dense graphs, implementing the above introduced Ricci flow via this approximation can significantly reduce runtime 54).
3 Geometric Coarsening and Pooling in attributed graphs
Characterizing the coarse geometry of a graph is invaluable for graph learning tasks. In this paper, we introduce a geometric pooling operator (ORC-PoOL, Fig. 3) which may serve as a standalone graph coarsening approach, as well as a pooling layer that can be integrated into GNN architectures. ORC-POOL allows for evaluating similarity structure encoded in the graph's geometry via ORC flow, while also accounting for similarity structure encoded in the node attributes. We follow the SRC framework (28) discussed above to formalize our approach.
3.1 Geometric Graph Coarsening
Geometric selection function. To define a pooling operator that elfectively coarsens graph, the selection function needs to identify cluster assignment that preserves the overall graph topology (community structure, etc.). We propose a geometric selection function, which computes the cluster assignment guided by the relation between curvature and graph topology (sec. 2.2 We encode curvature information into a matrix Cr = (ct) with entries Cij = Way given by the evolved edge weights after T iterations of ORC flow. The matrix C' can be seen as a curvature-adjusted adjacency matrix, which assigns an importance score to each edge reflecting its structural role: Under ORC flow (Eq. 3 structurally important edges are assigned a high score, whereas edges between similar nodes receive low score. This encodes global similarity structure, uncovering the coarse geometry of the graph (see discussion above and Fig. 2 Edges whose weight decreases under ORC flow contract (positive curvature), moving the adjacent, similar nodes closer together. On the other hand, edges whose weight increases under ORC flow expand (negative curvature), moving the adjacent dissimilar nodes further apart. This naturally induces coarsening of the graph. Curvature-based coarsening utilizes this emerging structure by cutting edges with a high weight (W. > A) and then merging the remaining connected components into supernodes. The threshold A is a crucial hyperaarameter, which is often difficult to choose in practice. Ni et al. (4I learn the threshold by optimizing the modularity of the resulting decomposition. Modularity optimization is NP-hard and hence an exact solution is intractable. They approximate the solution via an expensive parameter search. Since this subroutine is both expensive and non-differentiable, it is not suitable for integration into a trainable architecture. Hence, we need to employ a different auxiliary loss for identifying edges to cut, which we now describe.
Graph Cuts. Let Se {0, 1 N denote the assignment matrix computed by the selection function, its entries specify the assignment. of N vertices to K supernodes, i.e., 8i = 1, if vertex i is Vs and 8:s = otherwise. The problem of computing S by removing edges relates to the mincut problem, a classical problem in Combinatorics. In a seminal paper, Shi and Malik (49) show that the min-cut problem can be written optimization problem with respect to the number of edges within and between putative clusters (here, supernodes):
K Enjen linj
max K
Σ k=
Even ling
ling"
(4)
5

---

Published in Transactions on Machine Learning Research (12/2024)
G EV
G = (V,8)
SEL
RED
CON
& E:
Figure Proposed geometric pooling operator (ORC-Poot), which utilizes a curvature-basel geometric selection function (SEL) to identify supernodes and superedges (RED), which are then reconnected to generate the pooled graph (CON).
Here, Lnn denotes the edge indicator function. In the Graph ML literature, an equivalent formulation of due to Dhillon et al. 13 has been widely adapted,
St AS Σ
s.t. IN.
(5)
DS
where A denotes the graph's (unweighted) adjacency matrix, D= diag A) its degree matrix and Sy the kth row of S. While this problem itself is NP-hard, it can be efficiently solved via the relaxation (52
K
Q' = argiaax QERNXK K
AQ R=11
s.t. Q = D DS)
(6)
QUQ-IK
Using classical results from spectral graph theory, one can show that the optimizer takes the form Q UKV, where UK E N is formed by the top K eigenvectors of the normalized adjacency A and V E O(K) is an orthogonal transform. To recover the cluster assignment S, one can apply k-means clustering to the rows of Q' (57). Approaches for efficiently computing minimal graph cuts have been adapted for the design of effective graph pooling operators (9] BI). We have argued above that computing a curvature-adjustmeen emphasizes the underlying community structure (e.g., Fig 2), which should make it easier to learn graph cuts. Note that after uuvaturr-addjstment, the graph is weighted, i.c., we now optimize the sum of edge weights within and between pulative clusters. We will provide theoretical evidence for this obscrvation in the next section. Weighted graph cuts have been considered previously, e.g., in METIS (14), which underlies the GRACLUS pooling layer. We can adapt Eq. f to loss function for computing cluster assignment based on the (normalized) curvature-and eurvature-adjusted adjacency matrix, i.e.,
trST trST
min Lcp
Note that the entries of the degree matrix D are the weighted node degrees after adjustment. For T = 0, we recover the classical minimum graph cut objective. Notably, the resulting objective is differentiable and can be optimized with standard techniques (sce above), which is preferable over the parameter search for a suitable edge weight threshold performed in
Geometric pooling operator (ORC-Pool). The reduction and connection functions of ORC- PooL closely resemble canonical choices that are common among many pooling operators (e.g., MIN- CUTPOOL (8) DIFFPooL (65), NMF 4), among others): Given selection S E (0, N which assigns N nodes to K supernodes, we set RED(X) = STX =: X' and CoN(E) = S ES == E'. Naturally, if G is not autributed, then there is no need for a reduction function. In the coarsened graph, edge weights are again initialized either as or according to differences in supernode attributes (implemented in the Con function). This coarscning scheme may be applied multiple times, i.c., we can learn a sequence Si, S2... of selection matrices, which compute cluster assignments with varying degree of coarseness. With that, ORC-Pool may standalone coarsening clustering approach.

---

Published in Transactions on Machine Learning Research (12/2024)
3.2 ORC- Pool layer for Graph Neural Networks
While the geometric pooling operator described above applies to attributed graphs, its selection function does not evaluate similarity structure encoded in node attributes, aside from specific edge weight initialization discussed earlier. In this section, we will describe an alternative coarsening approach, where ORC-POOL is integrated into GNNs a pooling layer grounded in graph geometry. Stacked on top of MP base layers, which learn node representation that reflect similarity structure in both graph topology and node attributes, ORC-POOL integrates both types of structural information to coarsen the graph.
Pooling operators can be integrated into GNNs by stacking pooling layers, which implement the operator, on top of blocks of base layers (usually MP layers). The base layer learns node embeddings X GNN(X, A),,, where 9 are trainable parameters. If the imput graph is attributed, its node attributes are used to initialize the base layer. We implement the ORC-PoOL, layer as an MLP,, where we learn the cluster assignment matrix S' with an MLP with softmax activation (enforcing the constraint S1.- = 1x) and trainable parameters tb (S = MLP(X, v)), which are trained using Eq. particular, we optimize the training objective
7
as an auxiliary loss. The GNN is trained end-to-end. In
min (cp(0, tr ST CTS
STS STSS|F
VKlPp. IA
IK
+
(8)
ur ST DS
where the first term of the objective corresponds to the auxiliary loss and the second encourages mutually orthogonal clusters of similar size, preventing degenerate clusters (II-Iır denotes the Frobenius norm). Notice that the training objective closcly resembles that of MINCUTPOOL (8), but utilizes the curvaturo-adjusted adjacency matrix Cr. In particular, recover MINCUTPOOL
special case (T ().
We obtain the adjacency matrix of the coarsened graph by setting A" = ST AS and recomputing curvature- adjusted weights via Ricci flow on the superedges. We initialize superodge weights as a measure of similarity of the supernode attribules, which are obtained by setting X = 5"" Stacking blocks of base layers and pooling layers on top of each other yields a successive coarsening of the graph, which reflects its multi-scale structure. Notice that curvature-based importance scores influence the message functions on cach scale. The trainable parameters 1b, 9 are trained end-to-end as part of the GNN architecture and the respective loss for the downstream task, allowing for seamless integration of geometric pooling into GNN architectures. As a final remark, notice that the curvature-adjustmentt employed in ORC-POOL is flexible and could be incorporated into other (dense) pooling operators too (espocially those that build on graph cuts).
4 Properties of ORC-Pool
4.1 Basic properties
Permutation--nvarianne. Graphs and sets are permutation-invarinn,, as there is no natural order relation on the nodes. Standard MP layers are permutation-equivariin and pooling layers permutation-invariat to encode this structure inductive bias. The curvature-adjustment added in ORC-POOL preserves this property (for details see Apx.
Impact on Expressivity. A classical measure of the representational power of GNNs is expressivity, which evaluates a GNNs ability to distinguish pairs of non-isomorphic graphs. It is well-known that standard lest, classical heuristic for graph isomorphism testing [63] 38). To reap the benefits of pooling, pooling operators should preserve the expressivity of the MP base layer. Recently, (10) established conditions for this property. A simple corollary shows that ORC-POOL fulfills these conditions (see Apx. G for details):
MPGNNs are as powerful as the Weiserfeiler-LLhnmaee
Corollary 1. Consider a simple architecture with a block of MP base layers, followed by a ORC-PooL layer. Let G G2 denote two L-WL-distinguishable graphs with node attributes X, Further let X + X; denote the node representations learned by the block of MP layers. Then the coursened graphs Gp. Gp learned by the ORC-PooL layer are 1- 1-WL distinguishable.
7

---

Published in Transactions on Machine Learning Research (12/2024)
Topological Effects. MPGNNs known to suffer from oversmoothing (35), which describes the con- vergence of the representations of dissimilar nodes in densely connected subgraphs as the number of layers increases. This effect is particularly prevalent in node-level tasks, e.g., negatively impacting the GNN's ability to perform node clustering. ORC-POOL mitigates this effect in two ways: By assigning higher weights to sparse connections belween dense subgraphs and smaller weights to edges within dense subgraphs, distances between similar nodes are contracted and distances between dissimilar nodes are expanded. Moreover, pooling layers generally counteract oversmoothing, as they induce scale-separation. That is, local and global features are preserved separately, as they are encoded on different coarsening scales. Since similar nodes are merged into supernodes, neighboring supernodes tend to be less similar than neighboring nodes at the previous scale, alleviating oversmoothing on coarser scales. Conversely, limiting the number of layers could avoid over-smoothing, but shallow MPGNNs are known to suffer from under-reaching (7): If the number of layers is smaller than the graph's diameter, information between distant nodes not exchanged, which can limit the utility of the Icarned representations in downstream tasks. This underscores the value effective pooling layers for the design of deeper GNNs. We comment on the impact of pooling on over-squashing, related effect, in Apx
4.2 Structural Impact of Curvature Adjustment
We further corroborate our carlier claim that ORC flow reveals coarse geometry and emphasizes the community structure, making it easier to learn suitable graph cuts. For our analysis, we employ a graph model, which exhibits community structure as found in real data:
Definition 1. Consider a class of model graphs Ga.b (a (a b2 2), which are constructed by taking a complete graph with b nodes and replacing each by a Figure 4: Ga.b (a b 3) complete graph with a + 1 nodes.
Ga,-graphs can be seen as instances of stochastic block models with b blocks of size a a 1 and exhibit a clear community structure (see Fig. for an example). Due to the regularity of the graph, we can categorize its edges into three types: (1) bridges, which connect dissimilar nodes in different clusters; (2) internal edges, which connect similar nodes in the same cluster that are at most one hop removed from a dissimilar node; and (3) all other internal edges, which connect similar nodes within the same cluster. [4 derived an evolution equation for edge weights under ORC flow for Gu,b graphs:
Lemma (informal, (1))
Let w² = [wi, why wh] denole the weights of edges of types (1)-(3). Assuming
nh = 1,, the edge weights evolue under ORC flow (Eq. 3) as w F(a, b)w', where F F(a, b) a fired (3 3)mmatriz, depending on a, b only.
Utilizing this result we want to analyze the strength of the community structure subject to Ricci flow, employing modularity (24) as measure:
dy
Q(W) = Eav Wav
Wav
O(Cn,Co)
Eno Wuv/
as
Here, W = (wy)) denotes weighted adjacency matrix, d, weighted node degrees and 8(C,, C,) if u, are in the same cluster and zero otherwise. High modularity is often used as indicator of "strong" separation of the graph into subgraphs, under which communities can be casier to detect. We show that modularity increases as edge weights evolve under ORC flow:
Theorem 1 (informal). Let W denote the original adjacency matriz and Ct the curvature-adjusted adjacency matria after & ilerations of ORC flow. Then Q(C.) > Q(W), with Q(C.) increasing in t.
We defer the full technical statement of Lem. 3 and Thm. 5 as well proof details to supp. G
8

---

Published in Transactions on Machine Learning Research (12/2024)
Table 1: Node Clustering. Average accuracy (NMT) for ORC-PoOL. comparison with state of the art. pooling layers; average time per epoch (in seconds) given in brackets. The reported times (mean/ standard deviation) are computed based on 10 trials. Highest accuracy in bold.
Layer
Cora
CiteSeer
PubMed
No pool
0.33
(0.008 = 0.000)
0.22
(0.009 + 0.000)
0.14 = 0.02
(0.060 0.001)
Diff
0.20
(0.010=00000)
0.11
(0.009 0.000)
0.03
(0.082
Mincut
0.43 ::
(0.018 1.0001
0.31 0.04
(0.016 0.001)
0.23 = 0.04
(0.689 1.689::0000 0.003)
DMoN
1.01010.000 0.000)
(0.010 0.000)
0.04
0.000)
TV
0.33 4 0.05
00.000=
0.30 0.05
(0.010 0.000)
0.21 0.03
(0.069 0.001)
Graclus
0.43 1: 0.03
1
0.34
10.013:1.0000
0.27 1 0.00 10.138
ORC (us)
0.47 1: 0.04
(0.035 0.000)
0.35 0.04
(0.029 0.000)
0.24 0.04
(0.92 1 0.003)
Table Graph Classification. Average classification accuracy for ORC-POOL in comparison with state of the art pooling layers, averaged over 10 trials, using a 80/10/10 arain/va//teii split. Highest accuracy in bold.
Layer
MUTAG
ENZYMES
PROTEINS
IMDB- BIDD-BIIARY REDDIT-BINARY
COLLAB
PEPTIDES
No Pool
0.81
0.05
0.29 0.04
U.74 + 0.03
0.58 0.02
0.85 = 0.03
0.67 0.03
0.63 + 0.03
Diff
0.83 1 0.09
0.31 0.03
0.75
0.04
0.64 0.06
0,86 1 0.02
0.70 1. 0.02
0.66 1 0.02
Mineut
0.83 + 0.07
0.37 0.05
0.76 0.05
0.65 + 0.05
0.85 - 0.02
0.67 + 0.02
0.67 0.02
DMoN
0.00
0.05
0.67
TV
0.83 + 0.08
0.37 0.05
0.75 + 0.04
0.54 1:4::0 0.03
0.85 + 0.03
0.70 0.02
0.67
Graclus
= 0.06
0.02
0.61 0.04
0.84 0.04
0.08 +.00
0.06
ORC (us)
0.90 = 0.06
0.06 0.08
0.71 0.04
0.88 0.02
0.71 0.19
5
Experimental Analysis of Geometric Pooling
In this section we present experiments to demonstrate the advantage of our proposed pooling layer ORC-POoL. We test our hypothesis that encoding local and global geometric information into the pooling layers con increase the accuracy of the GNN in downstream tasks.
Experimental setup. We implement a simple GNN architecture, consisting of blocks of GCN base layers, followed by a pooling layer. When comparing ORC-PooL with other state of the art pooling layers, we keep the architecture fixed (only altering pooling layers) to enable fair comparison. Details on the architecture used in our node- and graph-lovel experiments be found in Apx. D.I We utilize the popular benchmarks PLANETOID H for node clustering, and TUDATASET KT and LRGBDATASET 17 for graph classification. Experiments are performed using PYTORCH GEOMETRIC, We compare ORC-POOL with four state of the art pooling layers, which were reported as best-performing across domains and graph lcarning tasks (28):
MINCUTPOOL ( DIFFPOOL (65), TVPooL (32). and Deep Modularity Networks (DMoN) (55) We
implement ORC-PooL using exact ORC, i.e., (1,).) is computed via the earth mover's distance.
Node Clustering. We compare the performance of ORC-POoL, and other pooling layers for node clustering, where we evaluate the Normalized Mutual Information (short NMI, defined sec. D.2] of the cluster assignments computed by the GNN, as well as the average runtime per epoch. The number of desirod clusters is known to the model beforehand. Results for PLANETOID are reported in Tab. F We see that ORC-POOL performed best overall with the highest average NMI on all three graphs. The second best model in all cases was MuNCU,,,LL, which the ORC-POOL layer is based on. In terms of runtime, we observed that for the larger PubMed graph, DIFFPOOL, DMoN, and TVPOOL are about one order of magnitude faster per epoch compared to MINCUTPoOL and ORC-PooL. The difference in time per epoch is much smaller for Cora and CiteSeor.
Graph classification. We further compare the performance of ORC-POOL with that of other pooling layers for graph classification, where we report the accuracy of label assignments. ORC-POOL. performed best. overall on TUDATASET, achieving the best accuracy on all datasets but ENZYMES (see Tab. 5 We further tested ORC-POOL on PEPTIDES, largo-scale graph classification tasks from the long-range graph benchmark LRGB (7]. Here, again, ORC-POOL showed superior performance. Run times for all experiments can be found in Tab. D.3
9

---

Published in Transactions on Machine Learning Research (12/2024)
Table Curvature computation. Comparison of accuracy (NMT) and runtime (in brackets) of computing ORC exactly (EMD), via Sinkhorn distances (Sinkhorn), and via combinatorial ORC approximation. Best runtime in bold.
Layer
Cora
CiteSeer
PubMed
EMD
0.47 +: 0.04 (19.97 1: 0.35)
0.35 4: 0.04 (16.77 à 0.21)
0.24 4: 0.04 (571.11 2.94)
Sinkhorn
0.03 (56.36 + 1.43)
0.35 + 0.03 (42.16 0.71)
0.22 0.05 (904.38 3.088
ORC-approx
0.45 = 0.03 (16.88 = 0.43)
0.35 = 0.03 (14.62 0.10)
0.21 = 0.03 (548. 13 1.79)
Curvature computation. ORC-POol defines a dense pooling layer; i.e, E [[//[NV] O(N). This property is inherited from the MINCUTPoOL. objective and preserved by the curvature adjustment. Hence, the complexity of the pooling operator is O(NK) for storage and O ((((E)++ NK)) for minimizing the relaxed min-cut loss (9). As noted in sec. 2.2.3] the computation of the curvature-adjustment can be costly on large, dense graphs. We test the performance of the two introduced curvature approximations (Sinkhorn distances, combinatorial ORC) on the accuracy and average runtime per epoch. We focus on node-level tasks, where input graphs are of larger size. Our results (Tab. 5 indicate that computing Sinkhorn distances is act slower than EMD on large, dense input graphs. This consistent with previous results for curvaturo-bass community detection on graphs with similar topology (5z). In contrast, the combinatorial ORC approximation delivers faster method, which retains accuracy similar to exact ORC. We note that implementation likely be further optimized, which could lead to additional speedups, e.g,, via more effective parallclizatio.
6 Discussion
We introduced ORC-PooL, a geometric pooling operator that leverages coarse geometry, characterized by Ollivier-Riie Mievier-Riici curvature and an associated geometric flow. ORC-PoOL extends class of Ricci flow based clustering algorithms to attributed graphs and can be incorporated into GNN architectures as pooling layer. While curvature adjustment leads to improved performance in downstream lasks, it also adds computational overhead, specifically in node-level tasks with large, dense input graphs. While our experiments indicate that the runtime impact is modest in most cases, we show that a combinatorial ORC approximation can reduce runtime impact while retaining accuracy on node-level tasks. Nevertheless, an important direction for future work is the study of alternative curvature notions or approximations (see Apx. A.2p. which could result more scalable implementations of ORC-POOL. In addition, it would be interesting to investigate the impact of ORC-POoL layers in other graph learning tasks, such as graph regression or reconstruction, well on wider range of graph domains. Lastly, it would be interesting to consider extensions directed input graphs.
Acknowledgements
The authors thank Yu Tian for sharing code. AF was supported by the Harvard College Research Program and a KURE Fellowship from the Kempner Institute. MW was partially supported by NSF award CBET-2112085.
References
[1] E. Abbe. Community detection and stochastic block models: Recent developments. Journal of Machine Learning Research, 18(177):1-86, 2018
[2] U. Alon and E. Yahav. On the bottleneck of graph neural networks and its practical implications. In International Conference on Learning Representations, 2021.
33] A. Arnaiz-Rodrigues, A. Begga, F. Escolano, and N. Oliver. Diffwire: Inductive graph rewiring via the Lovász bound. arXiv preprint 2rXii:2206.07369, 2022.
[4] D. Bacciu and L. D. Sotto. non-negative factorization approach node pooling graph convolutional neural networks. In International Conference of the Italian Association for Artificial Intelligence, pages 294 304-366.60 Springer, 2019.
10

---

Published in Transactions on Machine Learning Research (12/2024)
[5] J. Back, M. Kang, and S. J. Hwang. Accurate learning of graph representations with graph multiset pooling. aixiv preprint orXiv:2100 (22000211333 11533, 2021.
[6] Back, M. Kang, and S. Hwang. Accurate learning of graph representations with graph multiset pooling. CoRR, abs/2102.115533 2021.
[7] P. Barceló, E. V. Kostylev, M. Monet, Pérez, Reutter, and J. Silva. The logical expressiveness of graph neural networks. In International Conference on Learning Representations, 2020.
[8] F. M. Bianchi, D. Grattarola, and C. Alippi. Spectral clustering with graph neural networks for graph pooling. In International Conference on Machine Learning, pages 874-883. PMLR, 2020.
[9] F. M. Bianchi, D. Grattarola, L. Livi, and C. Alippi. Hierarchical representation learning in graph neural networks with node decimation pooling. IEEE Transactions on Newral Networks and Learning Systems, 2020.
[10] F. M. Bianchi and V. Lachi. The expressive power of pooling in graph neural networks. arXie preprint or Xin:2304.00 01575, 2023.
[11] V. D. Blondel, J.-L. Guillaume, R. Lambiotte, and E. Lefebvre. Fast unfolding of communtiies in large networks. Journal of statistical mechanics: theory and experiment, 2008(10):110008, 2008.
[12] F. R. Chung. Spectral graph theory, volume 92. American Mathematical Soc., 1997.
[13] I. S. Dhillon, Y. Guan, and B. Kulis. Kernel k-means: spectral clustering and normalized cuts. In Proceedings of the tenth ACM SIGKDD international conference on Knowledge discovery and data mining, pages 551-556, 2004.
[14] 1. S. Dhillon, Y. Guan, and B. Kulis. Weighted graph cuts without eigenvectors multilevel approach. IEEE transactions on pattern analysis and machine intelligence, 29(11):99441165, 2007.
[15] 1. S. Dhillon, Y. Guan, and B. Kulis. Weighted graph cuts without eigenvectors a multilevel approach. IEEE transactions pattern analysis and machine intelligence, 29(11):1944-1957, 2007.
[16] A. Duval and F, Malliaros. Higher-order clustering and pooling for graph neural networks. In Proceedings of the 31st ACM international conference on information & knowledge management, pages 426-435, 2022.
[17] V. P. Dwivedi, L. Rampášek, M. Galkin, A. Parviz, G. Wolf, A. T. Luu, and D. Beaini. Long range graph benchmark. In Thirty-siath Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022.
[18] M. Erbar and J. Maas. Ricci Curvature of Finite Markov Chains via Convexity of the Entropy. Archive for Rational Mechanics and Analysis, 206(3), Dec. 2012.
[19] L. Fesser, S. S. d. H. Iváñez, K. Devriendt, M. Weber, and R. Lambiotte. Augmentations of Forman's Ricci curvature and their applications in community detection. arXiv preprint 3rXiv:2306.06/ 2023.
[20] L. Fesser and M. Weber. Mitigating over-smoothing and over-squashing using augmentations of Forman- Ricci curvature. In Learning Graphs Conference, 2023.
[21] M. Fey and J. E. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Monifolds, 2019.
[22] M. Fiedler. Algebraic connectivity of graphs. Czechoslovak mathematical journal, 23(2):298 305, 1973.
[23] R. Forman. Bochner's Method for Cell Complexes and Combinatorial Ricci Curvature. volume 29, pages 323 374, 2003.
[24] S. Gómez, P. Jensen, and A. Arenas. Analysis of community structure in networks of correlated data. Phys. Rev. E, 80:016114, Jul 2009.
11

---

Published in Transactions on Machine Learning Research (12/2024)
[25] M. Gori, G. Monfardini, and F. Scarselli. A new model for learning in graph domains. In Proceedings. 2005 TEEE international joint conference on newral nelworks, volume 2, pages 729-734, 2005.
[26] A. Gosztolai and A. Arnaudon. Unfolding the multiscale structure of networks with dynamical Ollivier- Ricci curvature. Nature Communications, 12(1): 2021.
[27] D. Grattarola and C. Alippi. Graph neural networks in tensorflow and keras with spektral |application notes]. IEEE Computational Intelligence Magazine, 16(1):90-106, 2021.
[28] D. Grattarola, D. Zambon, F. M. Bianchi, and C. Alippi. Understanding Pooling in Graph Neural Networks, 2021. arXiv:2110.05292.
[2] A. Hagberg, P. Swart, and D. S Chult. Exploring network structure, dynamics, and function using networkx. Technical report, Los Alamos National Lab.(LANI) Los Alamos, NM (United States), 2008.
[30] W. Hamilton, Z. Ying, and Leskovec. Inductive representation learning on large graphs. Advances in neural information processing systems, 30, 2017.
[3]] J. B. Hansen and F. M. Bianchi. Clustering with total variation graph neural networks. areie preprint or Xwx:221:111000 1022188, 2022.
[32] J. B. Hansen and F. M. Bianchi. Total variation graph neural networks. In Proceedings of the 40th international conference on Machine learning. ACM, 2023.
[33] W. Hu, M. Fey, M. Zitnik, Y. Dong, H. Ren, B. Liu, M. Catasta, and J. Leskovec. Open graph benchmark: Datasets for machine learning on graphs, 2021.
[34] T. N. Kipf and M. Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint orXiv: 029007 2016.
(35) Q. Li, 7. Han, and X.-M. Wu. Deeper insights into graph convolutional networks for semi-supervised learning, In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018
(36] Y. Lin, L. Lu, and S.- Yau. Ricci curvature of graphs. Tohoku Matheratioca Journal, 63(4):605 - 627, 2011.
[37] C. Morris, N. M. Kriege, F. Bause, K. Kersting, P. Mutzel, and M. Neumann. TUDatasot: A collection of benchmark datasets for learning with graphs. In ICML 2020 Workshop on Graph Representation Learning and Beyond (GRL+ 2020), 2020.
[38] C. Morris, M. Ritzert, M. Fey, W. L. Hamilton, J. E. Lenssen, G. Rattan, and M. Grohe. Weisfeiler and Leman go neural: Higher-order graph neural networks. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages
2019.
[39] E. Müller. Graph clustering with graph neural networks. Journal of Machine Learning Reseorch, 24:1-21, 2023.
[40] K. Nguyen, T. Nguyen, N. Ho, K. Nguyen, H. Nong, and V. Nguyen. Revisiting over-smoothing and over-squashing using Ollivier's Ricci curvature, arXiv preprint ar 2r11:22211.5779,, 2022.
[41] C.-C. Ni, Y.Y. Lin, F. Luo, and J. Gao. Community detection on networks with Ricci flow. Scientific reports, 9(1):1 12, 2019.
[42] Y. Ollivier. Ricci curvature of Markov chains on metric spaces. Journal of Functional Analysis, 256(3):810 864, 2009.
[43] Y. Ollivier. A survey of Ricci curvature for metric spaces and Markov chains. In Probabilistic Approach to Geomelry, pages 343-30 381. Mathematical Society of Japan, 2010.
[44] A. Paszke, S. Gross, S. Chintala, G. Chanan, E. Yang, Z. DeVito, Z. Lin, A. Desmaison, L. Antiga, and A. Lerer. Automatic differentiation in pytorch. 2017.
12

---

Published in Transactions on Machine Learning Research (12/2024)
[45] O. Platonov, D. Kuznedelev, M. Diskin, A. Babenko, and L. Prokhorenkova. A critical look at the evaluation of gnns under heterophily: are we really making progress?, 2023.
[46] Y. Rong, W. Huang. T. Xu, and J. Huang. Dropedge: Towards deep graph convolutional networks node classification. In International Conference on Learning Represenlations, 2020.
[7]] C. Sanders, A. Roth, and T. Liebig. Curvature-based pooling within graph neural networks. arXiv preprint or 23v: 16516, 2023.
[48] R. Sandhu, T. Georgion, and A. Tannenbaum. Ricci curvature: An economic indicator for market fragility and systemic risk. Science advances, ((1):1500499,, 2016.
[9]] J. Shi and J. Malik. Normalized cuts and image segmentation. IEEE Transactions on pattern analysis and machine intelligence, 22(8))888 905, 2000.
50] J. Sia, E. Jonckheere, and P. Bogdan. Ollivier-Rieci curvature-based method to community detection in complex net works. Scientific reports, 9(1):1-12, 2019.
[5]] D. A. Spielman and S.-H. Teng. Spectral partitioning works: Planar graphs and finite element meshes. In Proceedings of 37th conference on foundations of computer science, pages 96 105. TEEE, 1996.
[2] X. Y. Stella and J. Shi. Multiclass spectral clustering. In Computer vision, IEEE international conference
on, volume 2, pages 313 313. TEEE Computer Society, 2003.
[33 A. Tannenbaum, C. Sander, L. Zhu, R. Sandhu, I. Kolesov, E. Reznik, Y. Senbabaoglu, and T. Georgiou. Ricci curvature and robust of cancer networks. arXiv preprinl 1rxiv:1502.00 04512, 2015.
[4] Y. Tian, Z. Lubberts, and M. Weber. Curvature-based clustering on graphs arXiv preprint or 0ix:223077.000 10155, 2023.
[55] A. Tsitsulin, J. Palowitch, B. Perozzi, and E. Müller. Graph clustering with graph neural networks, 2023.
[6] P. Veličković, G. Cucurull, A. Casanova, A. Romero, P. Lio, and Y. Bengio. Graph attention networks ar Xéu preprint S.::11100110000 2017.
[77]
U. Von Luxburg. A tutorial spectral clustering. Statistics and computing, 17:395-416, 2007.
[58] M. Weber, Jost, and E. Saucan. Forman-Rice flow for change detection in large dynamics data sets. Azioms, 5(4):doi-10, 2016.
[59] M. Weber, J. Jost, and E. Saucan. Detecting the coarse geometry of networks. In NewrIPS Relational Representation Learning, 2018.
[60] M. Weber, Saucan, and J. Jost. Characterizing Complex Networks with Forman- Riccii Curvature and Associated Geometric Flows. Journal of Complez Networks, 5(4):527-550, 2017.
[61] M. Weber, E. Saucan, and Jost. Coarse Geometry of Evolving Networks. Journal of Complex Networks, 6(5):706 7027706-73,,.00 2017.
[62] M. Weber, J. Stelzer, E. Saucan, A. Naitsat, G. Lohmann, and J. Jost. Curvature-based methods for brain network analysis. arXie preprint 1rXiu:1707.001 2017.
[63] K. Xu, W. Hu, J. Leskovee, and S. Jegelka. How powerful are graph neural networks? anXiv preprint or Xiv: 10::1810..m 1810.00080000 2018.
[64] Z. Yang, W. Cohen, and R. Salakhudinov. Revisiting semi-supervised learning with graph embeddings. In International conference on machine learning, pages 40-48. PMT.R, 2016.
[65] R. Ying, J. You, C. Morris, X. Ren, W. Hamilton, and Leskovec. Hierarchical graph representation learning with differentiable pooling. In arxiv preprint arXiv: 180.00000,, 2018.
[66] Z. Zhang, J. Bu, M. Ester, J. Zhang, C. Yao, Z. Yu, and C. Wang. Hicrarchical graph pooling with structure learning. Xrxiv preprint Xiv: 91 1.05954, 2019.
13

---

Published in Transactions on Machine Learning Research (12/2024)
Table of Contents
A Curvature Approximation
14
A.1 Combinatorial ORC Approximation
14
Forman's Curvature
B Illustration of Ricci Flow
15
C Additional Discussion of Related Work
15
D Additional Details on Experiments
17
D.1 Details on Experimental Setup
17
D.2 Node Clustering
17
D.3 Graph Classification
19
E
Additional Ablation Studies
20
Number of Ricci Flow iterations
20
E.2 Choice of Base Layer
21
E.3
Comparison with CURyPOOL (47
21
E.4
Comparison with additional Pooling Layers
F
Results Open Graph Benchmark
23
G Details on Theoretical Analysis
24
Properties of ORC-PooL
24
G.2 Structural Impact of Curvature Adjustment
25
Licenses
27
A
Curvature Approximation
A.1 Combinatorial ORC Approximation
For completeness, we restate the upper and lower bounds on ORC derived by Tian et al. (44 in our notation.
We compute the combinatorial ORC approximation defined in sec. 2.2.3 using these bounds.
Recall the definition of ORC given in Eq. 2 and note that we fix = 1 and a = 0. Let N(x) be the 1-hop neighborhood of node z. For vertices u, we have = N(u) \ N(u), r N(w) N(w), and e c= (w) n (v).. Mass distributions on node neighborhoods are denoted as Pu (z). We further define the shorthands Tag EePo(!) and Lap = (r())) a4 denotes the quantity m)
14

---

Published in Transactions on Machine Learning Research (12/2024)
Theorem 2 (Lower bound ( (54) Thm. 4.6)).
K u0 2 1
Whu
(c)
Σ
Wrv (r(
Wae
Way
Σ
Weo (pa(c) - -po(c))+ +
Wen
(po(c) Pu(c)))
Wuv
La + Pa(u) - -p((u) -(p((c)-Pu(())
To state the upper bound, we further define P = for E N(u) (2) : p.(()) (2) > 0} and N = {o e
N(u) U Ny)) Pu(2) - (x) 0}. For V and C V, let (z S) min(da(x, y 8}.
Theorem (Upper bound ((54), Thm. 4.6)).
do(x,N) (p))))-p(==)++
da( (po(z) - (p.(z)-Pu(z)) Pp(c)))
Aa <
max
Wuv
Wav
A.2 Forman's Curvature
Forman's curvature 233 gives combinatorial notion of Ricei curvature, which aggregates local information in the 2-hop neighborhood of an edge. That is, consider the endpoints of a certain edge and all edges that share an endpoint with the edge. One formalization of this curvature, which considers the contributions of connectivity and triangles for the curvature computation, resembles Ollivier's curvature qualitatively and has found applications in community detection. Its computation is casily parallelizable, allowing for fast and scalable implementation. This has given rise to a range of applications of Forman's curvature in Network Analysis (60 53
1S
62). While in general less suitable for identifying community structure (54), could 58
present a viable alternative to ORC in large-scale graphs, where computing ORC is infeasible. We can adapt the ORC-PooL operator proposed in this paper the case of Forman's curvature by simply exchanging the curvature notion
B Illustration of Ricci Flow
In the main text, we rated the evolution of edge weights under Rieci flow a dumbbell graph. In Fig. plot networks sampled from the stochastic block model, with two, three, four and five communities (left to right). We show the network with original weights (all equal to one) on the top and curvature-adjusted edge weights (after 5 iterations of Ricci flow, i.e., T 5) at the bottom, where again darker colors correspond to smaller curvature. We see that curvature reliable uncovers bridges between communities due to their low curvature, whereas edges within communities have a higher weight.
We further illustrate the effect of Ricci flow on molecular graphs, which are representative of graph topologies often encountered in graph-level tasks. Fig. G shows the evolution of edge weights under Ricci flow (T 4) for two graphs from the MUTAG data set, again initializing edge weights uniformely with one. We see that the curvature-adjjstment highlights "bridges" (dark color) between functional substructures, e.g,, rings. This allows the selection function to perform pooling that aligns with multi-scale structure in the graph.
C Additional Discussion of Related Work
We provide a more detailed discussion of related work by (47), which also proposes a curvature-based graph pooling approach, albeit with several key differences to our approach. We provide a detailed conceptual comparison of the two approaches below; an experimental comparison can be found in sec. F.3
Setting. Classically, graph pooling layers are designed jointly evaluate similarity structure encoded in the node features and graph connectivity, instead of connectivity only (see, e.g., DiffPool, MincutPool. well
15

---

Published in Transactions on Machine Learning Research (12/2024)
Figure Evolution of edge weights under Ricci flow the Stochastic Block Model.
Figure 6: Evolution of edge weights under Ricci flow in MUTAG.
discussions in references such as (2 Fig. illustrates that in attributed graphs both types of information need to be evaluated to identify meaningful sets of nodes to be pooled. However, the pooling layer introduced in (47) only evaluates similarity structure encoded in the connectivity. This is also corroborated by the experimental comparison of both approaches below, which ORC-PooL, achieves higher performance for all tasks on attributed graphs.
Motivation. The motivation for our proposed pooling layer (and that of related pooling layers) is to capture salient. multi-scale structure. In sec. 2.2 we give a geometric motivation for how curvature captures such structure utilizing a connection discrete Ricci curvature and random walks (see also Fig. 2 On the other hand, the approach Sanders et al. is motivated by addressing over-smoothing and over-squashing. While we agree that pooling (in general) mitigates over-smoothing due to the induced scale separation (as we over-squashing effects may actually be amplified as discuss in the appendix.
discuss in sec. 4.1
Different use of curvature and different notion. We note that while both the approach in Sanders et al. and our proposed approach leverages different notions of curvature, there are fundamental differences: (1) Our approach relies on curvaturo-adjnsted adjacency matrix, which is computed using Ricci flow, a geometric flow associated with discrete Ricci flow. The approach in Sanders al. uses discrete notion of
16

---

Published in Transactions on Machine Learning Research (12/2024)
curvature, it does not use Ricci flow. (2) Our approach utilizes Ollivier's Ricci curvature and approximations thereof. Sanders et al. use "balanced Forman curvature" instead. While related, the two curvature notions differ substantially. The geometric motivation given in our paper does not directly translate to balanced Forman curvature.
Complexity and hyperparameters. The balancod Forman curvature used in Sanders et al. has complexity O(E))emaz) On the other hand, variant of our proposed approach, which utilizes a combinatorial ORC curvature, has complexity of only al. state, their approach has several hyperparameters that have to be carefully chosen so as not to simplify the graph too much. Determining this threshold via grid search could add significant computational overhead. On the other hand, the main lyperaarameter in our method is the number Ricci flow iterations and the performance is relatively insensitive to the choice of this parameter.
dmaz), making it more scalable on large-scale graphs. As Sanders et
Overview related pooling methods. Tab. 1 gives an overview of our main pooling baselines in the SRC
framework.
Method
Select
Reduce
Connect A' = S AS
DifiPool (55
S = GNN, (A, (w/ auxiliary loss)
X' = S GNNN
MinCut
S = MLP(X) (w/ auxiliary loss)
X' = S
A' = S AS
Graclus
S = {ar, arg max, Dr, + Asi
}
X' = S
METIS
D3s S GNN(A, X) (w/ auxiliary loss)
DMoN 55
X = STX
A'= = ST AS
TV 32
S MLP(X) (w/ auxiliary loss)
X = S
A' 5 AS
Table 4: Overview of Pooling layers
D Additional Details on Experiments
D.1 Details on Experimental Setup
We implement a GNN architecture consisting of GCN base layers and a single pooling layer, which are jointly trained. For node-level tasks, one GCN layer is used to compute a graph embedding, and another GCN layer is used to compute node clusters from the embedding. Then the pooling layer computes the loss of the current node assignments, which used update the GCN parameters. For the graph-level tasks, our architecture alternates between blocks of GCN base layers and pooling layers, which are jointly trained. We train global pooling layer top, which computes the readout by aggregating node representations across the graph. To reduce the computational overhead of the ORC computation, we use the curvature-adjustmeen only in the first pooling layer, i.e, the second pooling layer is a standard MINCUTPOOL layer. We note that the choice of using MINCUTPOOL in subsequent layers was driven by a desire to maximize scalability, as the curval urenad would require recomputing curvature in the coarsened graph. However, we note that computing the addture-djjstmeett in subsequent layers has reduced computational cost due to the smaller size of the coarsened graph, which suggests that this is viable extension of the present approach that may be considered in future work.
To compare against GRACLUS, we incorporated the GRACLUS pooling layer, based on METIS, to PYTORCH GEOMETRIC. We note that our implementation is not optimized for runtime. Hence, we expect that it is possible to obtain further speedups in GRACLUS.
All experiments are run on NVIDIA A100 GPU with one CPU.
D.2 Node Clustering
Architecture. The first GCN layer embeds the nodes of the graph into an m-dimensional feature space. That is, if the graph initially had a node feature tensor of dimension n k, where is the number of nodes, in the output we have node feature tensor of dimension The second GCN layer takes this embedded
17

---

Published in Transactions on Machine Learning Research (12/2024)
X
GCN
Pool
S
Figure 7: Architecture for node clustering.
graph and computes an assignment tensor dimension n x e, where is the number of clusters. By taking the softmax of this c tensor, we obtain the final assignment of cach node a single cluster.
We note that in the original implementation of TVPooL, custom MP layers are used before pooling. Since the goal of our experiments is fair comparison of the effectiveness of the pooling layers, we use standard GCN layers as base layers for all pooling approaches.
We also performed the experiments with a GCN without pooling using the same architecture, but omitting the pooling layer. The output is computed by a softmax function, where the predicted class corresponds to the element with the highest value.
Data. The PLANETOID graphs Cora, CiteSeer, and PubMed are citation networks where documents are nodes and citations are edges.
Table 5: Planetoid
Graph
Nodes
Edges
Features
Classes
Cora
2708
10566
1433
7
CiteSeer
3327
9104
3703
6
PubMed
19717
88648
500
3
Hyperparameters. Parameters for the experiments are as follows: For PLANETOID, one GCN layer with output dimension & and an ELU activation function is used to embed the graph. The optimizer is an Adam optimizer with learning rate of 1e-2. The models trained for at most 10000 epochs, until the best NMI is found under patience constraint. That is, if we achieve the current best NMI on epoch and patience is p, stop training if the model does not have beller NMI at any epoch up to p. For the PLANETOID graphs, patience is set to 100, and for Amazon-ratings (from HETEROPHILOUS) patience is set to 250. When applying ORC-PooL to the PLANETOID graphs, four Ricci flow iterations are used, and for Amazon-ratings, two Ricci flow iterations are used.
NML. We measure accuracy as the Normalized Mutual Information (short: NMI) if the cluster assignments computed by the GNN. NMI a classical evaluation metric for community detection. In our experiments employ NMI to measure the accuracy in the node clustering task. We give a brief definition of NMI for completeness.
Let S Nx denote a vector that encodes the label assignment to & clusters, i.c., we set = 1 if node belongs to cluster Ct and 8u = 0 otherwise. The entries of S can be viewed as random varibles drawn from the distribution P(s = 1) = N and P(s= 0) = P(s 1), where N; == [Ci]. Using the marginal probability distribution Psj and the joint probability distribution P(st, sr), we further define the entropies ((s;) and ((s,, sy), as well the conditional entropy of st given sy as H(s: ((s(ss) H(s sp) NMT for two cluster assignments S' is given by
H (sv). The
NMI(S)S") = =1 (H(S)S)) H(')s)
18

---

Published in Transactions on Machine Learning Research (12/2024)
GCN
X1
Pool
[xp )
GCN
X2
Pool
[X2 A
GCN
Global Pool classifier
(X, A]
Figure 8: Architecture for graph classification.
D.3 Graph Classification
The main part of the model consists of two blocks of GCN layers followed by pooling layer. The first set of GCN layer(s) embeds the nodes of the graph into an m-dimensional feature space. That is, if the graph initially had a node feature tensor dimension n x k; where n is the number nodes, m. The second GCN layer takes this embedded the softmax of this assignment tensor of dimension X e, where is the number of clusters. By taking tensor, obtain the final assignment of each node to a single cluster. Here, the clustering simply gives us a new graph in which groups of nodes have been aggregated, which we pass into the second block. After the second block of GCN layer(s) and pooling, we pass the resulting graph into a third block of GCN layors to get the final node embeddings. These are passed into a global pooling layer, which aggregates node embeddings. In our case, we use global mean pooling, which simply computes the average of all node embeddings. The resulting vector is then passed into linear classifier.
urehitecture.
in the output we have a node feature tensor of dimension graph and computes
As an additional baseline, also GCN without pooling all data sets. The output is obtained using global mean pooling followed by linear classifier.
Data. Statistics for TUDATASET graphs are shown below where the task is to predict the mutagenic effect. For ENZYMES, we classify enzymes into one of six classes depending on which type of chemical reaction they catalyze. PROTEINS has a binary classification task for whether a protein an enzyme. REDDIT-BINARY COLLAB, and IMDB-BINARY are social networks where we predict the subreddit, research field, and genre of people in the network, respectively. For the graphs that did not have any node features (COLLAB, dummy node feature that is 1.0 for all nodes.
37 MUTAG is a common dataset of chemicals
REDDIT-BINARY), add
Peptides-func is the only dataset from LRGB asks to classify peptides based on their function.
Table 6: TUDatasel and Peptides-func
Dataset
# of Graphs
Features
Classes
MUTAG
188
7
ENZYMES
600
3
6
PROTEINS
1113
IMDB- IMDB-BINARY
1000
0
2
REDDIT-BINARY
1000
COLLAB
5000
0
3
Peptides-fune ffndde
15535
9
10
Hyperparameters. Parameters for the experiments are as follows: For all graphs, the GCN blocks depicted above consist of a GCN layer with output dimension 8 and an ELU activation function. The optimizer an Adam optimizer with learning rate of 5c-4 and a weight docay of 1c-4. The models are trained for most 10000 epochs, until the best accuracy on validation found under patience constraint. That is, if achieve the current best validation accuracy on epoch and patience is p, we stop training if the model
19

---

Published in Transactions on Machine Learning Research (12/2024)
does not have better validation accuracy any point up to epoch to 50. When applying ORC-PooL to MUTAG, ENZYMES, and PROTEINS, one iteration of Ricci flow curvature is used. For IMDB-BINARY, REDDIT-BINARY, and COLLAB, two, three, and two iterations are used, respectively. Two iterations are used for Peptides-func. We let the pooling layer approximately halve the number of nodes each time. That is, after the first. pooling layer, the number of nodes Xf is the half the average number nodes for graphs in the dataset. The number of nodes X is half the number of nodes in X
For all datasets, patience is set
Runtime comparison.
We report runtime comparison of the graph-lovel tasks.
Table 7: Average seconds/epoch on TUDataset and Peptides-func.
MUTAG
ENZYMES
PROTEINS
MNDB-BINARY
No Pool
0.023 1 0.001
0.13 + 0.02
0.22 1 0.00
0.16 1 0.001
Diff
0.057 0.001
0.17 1 0.03
0.34 0.01
0.29 1 0.00
Mincut
0.068 + 0.001
0.21 + 0.03
0.41 0.00
0.35 1.3500 0.01
DMoN
0.069 0.001
0.21 0.02
0.43 0.02
0.36 0.00
TV
0.068 + 0.000
0.20 = 0.02
0.40 + 0.01
0.34
Graclus
0.061 0.002
0.17
0.038 0.01
0.35 0,02
ORC (us)
0.190 + 0.003
0.48 + 0.03
1.32 + 0.00
1.07 + 0.02
REDDIT-BINARY
COLLAB
SEPTIDES-FUNG
No Pool
1.68 0.06
1.02 0.02
2.85 + 0.05
Diff
5.28 0.18
+
3.79 0.02
Mincut
5.39
2.05
1.44
DMoN
2.01
4.68 +
TV
0.09
2.02
4.35 0.05
Graclus
4.83 4.13
+
4.01 + 0.04
ORC (us)
12.90
6.05 + 0.01
36.27 0.10
Results for Heterophilous Graphs. The HETEROPHILOUS graph Aaazion-ratings is a network of Amazon products where edges connect products that are frequently bought together (45). The five classes correspond to product ratings. Unlike the PLANETOID graphs, we do not expect nodes the same class to consistently cluster together. The "natural" clusters of the graph generally correspond to categories items, rather than item ratings. Pooling layers combine clusters of nodes in order coarsen a graph, procedure whose utility
Table 8: Heterophilous data set.
Graph
Nodes Edges
Features
Classes
Amazon-ratings
24492
93050
300
depends on an implicit homophily assumption. Hence, we do not expect GNNs with pooling layers to perform well on data sets such the Amazon graph. This is confirmed by our experimental results, which show poor performance (in terms of NMT) for all models. In contrast, in the original HETEROPHILOUS study, GCNs without pooling layers were able to altain good performance 45
E
Additional Ablation Studies
E.1 Number of Ricci Flow iterations
We investigate the impact of using a varying number of Ricci flow iterations on the performance of ORC-POOL. In the case of node clustering, we re-run ORC-POOL on the Cora and CiteSeer graphs, using 2, 4, and iterations of Ricci flow. In the results presented in the main text, 4 iterations used.
20

---

Published in Transactions on Machine Learning Research (12/2024)
Table 9: Amazon
Model
NMI
Time/epoch (s)
Total time (s)
Diff
0.0014 + 0.00036
0.1218
53.13 + 18.47
Mincut
0.00004 = 0.00013
1.4494
364.66 + 1.89
DMoN
0.0012 = 0.00009
0.0898
35.48 0.47
TV
0.0010 + 0.00011
0.0980
52.99 + 16.34
ORC (us)
0.00067 0.00077
1.5369
405.69 25.41
Table 10: We report average accuracy (NMI) on Cora and CiteSeer when using varying Ricci flow iterations.
The average time per epoch (in seconds) is given brackets. The reported times (mean/ standard deviation) are computed based on 10 trials.
Iterations
Cora
CiteSeer
0.45 0.04 (0.033 + 0.001)
0.34 0.03 (0.028 + 0.000)
0.47 +
(0.035 + 0.000)
0.35 0.04 (0.029 + 0.001)
0.44 0.05 (0.034 + 0.001)
0.33 0.02 (0.029 = 0.001)
We find that the NMI is comparable for different numbers Ricci flow iterations. This indicates that the performance of ORC-POOL is not very sensitive the number of iterations used, and that the curvature-adjustment captures crucial multi-scale structure with only few iterations.
We also investigate the impact of different numbers of iterations of Ricci flow a graph classification task. We compared different numbers of iterations for the REDDIT-BINARY data set, testing 2, 3, and 4 iterations. In the main results, iterations were used. Our results show comparable results across all experiments, implying that the observations node-level tasks above extend to this setting. Based on our results, we expect that the performance of ORC-PoOT is not very scnsitive to the number of itcrations of Ricci flow and computing the curvalure-aljusttenn based small number of ilerations already leads to significant. performance increases. This is consistent with previous findings for curvature- curvature-based methods, see, e.g., 154
Table 11: We report accuracy (NMI) and time per epoch on REDDIT-BINARY when using varying Ricci flow iterations. The reported times (mean/ standard deviation) are computed based on 10 trials.
Iterations
Accuracy
Seconds/epoch
2
0.86 1 0.02
11.28 1 0.19
0.88 0.02
12.90 + 0.17
0.87 0.03
13.72 + 0.18
E.2 Choice of Base Layer
For all of the experiments above, used standard GCN layers as base layers for each experiment. Here, we investigate the performance of our model architecture using GAT and GLN instead of GCN.
For node clustcring, we re-run our CiteSeer experiment, again comparing ORC-POOL against four baselines. The experimental setup follows the description in Apx. DD.2 apart from GIN and GAT replacing GCN. Our results show a high NMI and overall similar performance gains using ORC-PoOL. For graph classification, make analogous observations when re-running our PROTEINS experiment with GAT and GIN replacing GCN base layers in the model architecture described in Apx. D.3
E.3
Comparison with CurvPool 447
We compare the results of ORC-POOL and CURYPOOL on PROTEINS and IMDB-BINARY, using the architecture, hyperparameters, and experimental setup of the CURvPOOL paper. (4 presents three versions
21

---

Published in Transactions on Machine Learning Research (12/2024)
Table 12: We report accuracy (NMT) on CiteSeer when using GAT and GIN in place of GCN base layers. The average time per epoch (in soconds) is given brackets. The reported times (mean/ standard deviation) are computed based 10 trials.
GAT
GIN
Diff'
0.28 0.031
(0.062 4
0.29 0.025
(0.097 + 0.003)
Mincut
0.32 0.018
(0.083 0.002)
0.33 0.023
(0.118 0.004)
DMoN
0.063
0.008)
0.30
(0.094
(0.086 0.007)
0.30 0.026
(0.095 0.008)
ORC (us)
0.34 + 0.040
(0.156 + 0.016)
0.35 0.040
0.011
Table 13: We report average classification accuracy for PROTEINS with GAT and GIN replacing GCN base layers. We use 80/10/ 10/010 rrin/val/ttett split. Averages are determined based on 10 trials. Highest accuracy in bold.
CAT
GIN
Diff
0.017
(3.82
0.73
(4.41 0.56)
Mineut
0.019
(4.16
(4.91 0.084
0.64)
DMoN
0.74 0.030
04.33
(5.17 0.75
TV
0.028
(4.01
(4.96 0.048) 0.76
ORC (us)
0.79 0.022
+.60330.3)
0.77 + 0.039
(9.21 0.73)
of CURVPOOL that differ in their pooling scheme. Below, we compare against the highest performing one, HIGHCURYPook
Table 14: We report average accuracy (NMI) for PLANETOID for both ORC-Pool and CURVPOOL
Layer
Cora
CiteSeer
PubMed
Curv
0.35
0.32
0.15
ORC (us) 0.47 1 0.04
0.35 1 0.04
0.24 1 0.04
For node clustering, we compare the two models on three PLANETOID data sets. The CURYPOOL clustering is deterministic, as it is based only the edge curvatures of the original graph (hence variance is reported). We used HIGHCURYPooL. with threshold of -0.2. We see that ORC-Pool. outperforms CURYPOOI on all three PLANETOID graphs, by an especially large margin for Cora and PubMed. This is expected, since ORC-POol utilizes information encoded in both node altribule and connectivity, whereas CURVPOOL only utilizes connectivity.
For graph classification, we compared ORC-POOL and CURVPOOL on one attributod (PROTEINS) and one unattributed data set. (IMDB-BINARY). We observe that ORC-PooL. achieved higher performance than CURvPOOL PROTEINS and equal performance IMDB-BINARY. This illustrates that ORC-PooL is able to utilize crucial information encoded in the node attributes and that this is crucial for the performance gains observed across data sets.
We note that achieving high performance with CURYPOOL requires careful tuning of the curvature threshold for pooling nodes. shown (47), the choice of this hyperparameter can cause accuracy change by as much as 5-10 percentage points, depending on the dataset. Therefore, hyperparameter tuning, e.g., grid search needs to be performed to determine the optimal threshold. In contrast, the key hyperparameter in ORC-POOLis the choice of the number of Ricci flow iterations to which downstream performance is fairly insensitive, as the experiments above indicate.
E.4 Comparison with additional Pooling Layers
We provide an additional comparison with recently introduced pooling layers HoscPool (66) MVPooL (66), GMTPool G. We follow the previously described experimental setup (see Apx. D.2 and D.3 for details). Our
22

---

Published in Transactions on Machine Learning Research (12/2024)
Table 15: We report average classification accuracy for PROTEINS (attributed) and IMDR-BINARY (unattributed) for ORC-POOL and CURVPOOL. We use 10-fold cross validation. Averages are computed over the test sets for each of the folds. Highest accuracy in bold.
PROTEINS
IMDE- BINARY
Curv
0.77 0.77 0.061
0.71 + 0.051
ORC (us)
0.79 0.053
0.71 + 0.049
experiments show that ORC-Pool achieves the highest accuracy two graph-level tasks (Mutag, Proteins) and the second highest node-level task (Cora). For the latter, only HoscPool achieves higher accuracy.
Cora
MUTAG
Proteius
GMTPool
0.001)
(0.066 0.001)
(0.40 0.00)
MVPool
0.32 = 0.03 (0.015 + 0.001
0.83
0.08 (0.071 + 0.001)
0.74 + 0.05
(0.42 + 0.01)
HoacPool
0.51 0.05 (0.020 # 0.001)
0.87
(0.089 0.002)
0.76
0.04
(0.58 + 0.01)
ORC-Pool (us)
0.47 0.04 (0.035 + 0.000)
0.90 0.08 (0.190 0.003)
0.78 0.04
(1.32 0.00)
Table 16: Comparison of ORC-Pool with additional baselines.
F Results on Open Graph Benchmark
We include additional experiments on large-scale data sets from the Open Graph Benchmark B3 For node clustering, we use OGBN-ARXIV, and for graph classification, we use OGBG-MOLHIV.
Like Cora and CiteSeer, OGBN-ARXIV is citation network. It has 169,343 nodes and 1,166,243 edges; the goal is the classification of papers by subject area. OGBG-MOLHIV is molecular property prediction dataset. It has 41,127 graphs with an average of 25.5 nodes. We follow the preferred evaluation metrics for OGB, reporting accuracy for OGBN-ARXIV and ROC-AUC for OGBG-MOLHIV. When running ORC-Pool on OGBN-ARXIV, iterations Ricci flow were used. For OGBG-MOLHIV, iterations of Ricci flow were used.
Table 17: We report accuracy and time per epoch on OGBN-ARXIV for all models.
Accuracy
Seconds/epoch
No Pool
0.59 + 0.037
0.33 + 0.002
Diff
0.65 1 0.031
0.59 1: 0.002
Mincut
0.70 1 0.027
3.26 1 0.009
DMoN
0.68 1 0.028
0.42 1 0.001
TV
0.68 1 0.040
0.41 +: 0.001
ORC (us)
0.72 = 0.036
9.29 + 0.011
Table 18: We report ROC-AUC and time per epoch on OGBG-MOLHIV for all models.
ROC-AUC
Seconds/epoch
No pool
0.64 + 0.038
12.37 + 0.78
Diff
0.68 + 0.052
16.87 + 1.51
Minent
0.70 + 0.040
18.38 t 2.43
DMoN
0.73 - 0.037
17.44 0.89
TV
0.69 + 0.034
17.58
0.79
ORC (us)
0.73 0.048
34.90 2.16
We observe that ORC-POOL performs best on OGBN-ARXIV; on OGBN-MOLHIV DMoNPoOL performs best. As in the experiments above, these results based off 10 runs.
23

---

Published in Transactions on Machine Learning Research (12/2024)
G Details on Theoretical Analysis
G.1 Properties of ORC-Pool
Exprcssivity. In the main text we stated the following result on the impact of ORC-POOL layers on the expressivily of the GNN:
Corollary 2 (Expressivity of ORC-PoO1.). Consider simple archilecture with block of MP base layers, following by a ORC-POOL layer. Let G1, G2 denote two lWLluiittngu graphs with node attributes X2. Further let X + X2 denote the node representations learned by the block of MP layers. Then the coarsened graphs Gf, G2 learned by the ORC-Pool layer are 1-WL distinguishable.
Proof. This result is simple adaptation of a recent result by 10 Thm.1), which establishes conditions under which pooling layers preserve expressivity:
Theorem 4 ( (OO, Thm. 1). Let G -G(XXE)) (X E RNx) denote the input graph and G' the graph oblained after applying a block of MP base layers to G; X' denoting the new multiset of node features. Let POOL: : G' G denote on SRC pooling layer after the MP layers, which produces pooled graph GP with mulli-sels XP. POOL preserves the expressivily of the MP layers, provided that the following conditions hold:
Let G1, G2 denote two WL-distinguishable grophs and X2 the node representations learnt by the MP layers. Then Σxex; x # Σxex; T.
The selection function assigns nodes a unqiue supernode; i.e., E/= sf = = all [N].
3. The reduction function assigns supernode representations as x.. = EN 2/8/.
The authors show that conditions (3). hold for dense pooling layers, such as MINCUTPoOL, and, hence, also ORC-POOL. In particular, by construction, graph cuts or partition-based community detection algorithms produce non-overlapping communities; hence the selection function assigns nodes to unique cluster, which becomes supernode. The reduction function computes supernode attributes as X = S X', which aligns with condition (3). Condition (1) is independent of the choice of pooling operator and is fulfilled for any MP layer that is as powerful as the 1-WL lest, e.g,, GIN 163
We note that the inherent lémitations in representational power in the MP layers may affect the quality of the coarsening learned by the pooling layer for certain classes of graphs. However, the empirical results presented in this work and the related literature indicate that, in practice, pooling layers coarsen graphs effectively.
Remark 1.
Permutation-invarianne. In the main paper, we stated the following property:
Lemma 2. ORC-POOL is permautaiioniimaarinn.
This property is inherited from MINCUTPOOL. For completeness, we provide the argument below.
Proof. Again, we utilize the SRC framework to describe the ORC-PooL operator. The seletion function SEL computes a cluster assignment matrix S. It is easy to see that the graph cut objective (Eq. (3.5) is permutation-equuiariaa... The cermutatioo--quivaraannn it simply performs a re-woighting of the odges. It further easy to see that the roduction and connection lunctions (RED and CON) are permutation-invarian With that, the composition (SEL RED) CoN is ionuutation-invariinn as desired.
is not impacted by the curvatue-adjustment as
Pooling and Over-squashing. As the number of layers in an MPGNN increases, one observes the formation of bottlenecks (2): Information from far distant nodes is encoded into fixed-length node representations during message passing, leading to information loss. This effect is particularly strong for bridges between clusters, which connect dissimilar nodes whose neighborhoods have a small intersection. Previous literature
24

---

Published in Transactions on Machine Learning Research (12/2024)
has characterized oversquashing via discrete Ricci curvature (40) 20): Edges that connect nodes in different clusters (bridges), have low Ricci curvature. This effect is emphasized under ORC flow, as the weight of the bridges increases. With that, our proposed curvature adjustment may amplify oversquashing effects. Recent literature has proposed graph rewiring as a mitigation for oversquashing effects. During rewiring, edges are re-samlled proportional to a relevance score (edge weight), which may be assigned utilizing the Lovasz bound (3, random edge dropping (46) or discrete curvature (40) 20), among others. Since we already compute discrete curvature during pooling, we could add a rewiring step with little overhead. We leave an investigation of this approach for future work.
G.2 Structural Impact of Curvature Adjustment
In the main text, we have given brief argument for how coarsening approaches based on ORC-Poo1, preserve the structural integrity of the graph. In this section, we expand on these results. To corroborate our claim that ORC flow reveals coarse geomelry and emphasizes the community structure, we employ graph model, which exhibits community structure as found in real data:
Definition 2. Consider class of model graphs Ga, (a b 2), which are constructed by taking complete graph with b nodes and replacing each by complete graph with + 1 nodes.
Figure 9: Ga,b (a b b=3) 3)
Graphs of the form Ga,b are special cases of stochastic block models (I with blocks (communilliss) of size a ++1 (see Fig. for an example). Due to the regularity of the graph, we can categorize its edges into three types: (1) bridges, which connect dissimilar nodes in different clusters; (2) internal edges, which connect similar nodes in the same cluster that are at most one hop removed from dissimilar node; and (3) all other internal edges, which connect similar nodes within the same cluster. [T derived an evolution equation for edge weights under ORC flow for Ga.b graphs:
Lemma 3 (AII)) Let w' = [w1, wh wh]l denote the weights of edges of types (1)-(3). Assuming initialization
w = 11, 1}, the edge weights evolue under ORC flow w (1 - Kwoda wu as
9-1
20
a+bb
w
ab-o-
6
w.
(9)
a(a+b)
a++
0
=:.(ab))
Building on this characterization, (41) show the follow result, which expresses edge weights w² with respect to the eigenvalues of the matrix Fa.::
Lemma 4 (AIT,, adapted to assumptions). The matric Fa,b has three real eigenvalues, Az, As, for which \ > X2 1 and d3 < 0, and corresponding eigenvectors Wi, W2, W3. Then one show
/111 + +((f)) w = a, Niw + a2Xiw2 agNgu3 = kai + +((1i)
(2)²
where & E (0, 1) and lo
o(Ai) 0.
Utilizing this result we want to analyze the strength of the community structure subject to curvature- adjustment, employing modularity (44 as a measure:
Q(W)
1
Σ
Wao Ene Wnv D(CCn, Co)
dydo )
(10)
Eav Wav as
= (wy) denotes a weighted adjacency matrix, d, weighted node degrees and 8(C,, C.) = 1, if u, v are in the same cluster and zero otherwisc. It can be shown that higher modularity corresponds to a "stronger" separation of the graph into subgraphs, which is easier to detect; e.g., the corresponding min-cut problem is easier to solve. We show that modularity increases, edge weights evolve under QRC flow:
Here, W =
25

---

Published in Transactions on Machine Learning Research (12/2024)
Theorem 5. Let W denote the adjacency matriz of the input graph and C₁ the prrvature-adjusten adjacency matrix after I iterations of ORC flow; i.e,, matriz whose entries corvespond to w/ for edges of type j. Then Q(C,) > Q(W,, with ((1,) increasing in (for not too small).
Proof. By construction, a Gu,b graph has 6(b-1) Lemma we have itcration (up o(Ai) 2 type (1), ab type (2) and ula-1) type (3) edges. Using io
error):
Σ w b(b
b(b-1 - 1) wt + abub + a(a 1)
b(b-1) 1)
2 a1 + abbka 21 a(a- 1) ((1)) =: Σ
bug
2
well degrees
dg (b 1) + aw3 (b 1)a1 aka, A1
c-1 (2)-²
dj w awi = kazdi +
for internal nodes (subscript i) and bridge nodes (subscript 6), the latter being adjacent to edges of type 1. We can write the modularity of a Ga, graph with weights evolved after iterations of Ricci flow as
Q(Ci)
1 b(b - 1) 2
w
(db)
8(Cus Co) +ab +ab(us, wh
d.d D
8(C,, Co)
(d)) 2
a(a- 1).
+
b
w/
8(Cuu Ci)
2
Σ
1
d'd
a(a - 1)
ww's (d))
ab
(1b
E
Dr
2
E
Notice that as t increases, we have did < 2 and (dt)² < Et, which implies that the terms
d'd'
and
decrease fast with t. Moreover, w/ = (1) is decreasing fast in t. Hence, the asymptotics of the sum in brackets are dominated by the term abwa. We introduce the shorthands Σ/ Eg, Σ/ for the contributions of edges of types (1)-(3) to Σt. With that the asymptotics of Q(C,) are dominated by
M
+
Σ3
+ 1+
Σ
We see that
Σ + Eg
b(b-1)
(2)
Σ
abk
decreases in & with the assumptions in Lemma asymptotically under the Ricci flow.
Putting everything together implies that Q(C,) increases
26

---

Published in Transactions on Machine Learning Research (12/2024)
G.3 Licenses
We provide below the licenses of all packages and data sets used in this work.
Model/Datasee
License
Notes
LRGB 1
Custom
See here for license
TUDatasel 37
Open
Open sourced here
Planetoid 64
MIT
See here for license
Pytorch Geometric (21
MIT
See here for license
Pytorch 4
3-clause BSD
See here for license
NetworkX 20
3-clause BSD
See here for license
GraphRiieeCurvaatuee
Apache-2.0 2.0 license
See here for license
27


Reviews:
[{"review_id": "635SDMfGru", "reviewer": "Reviewer_SpMC", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\nCurvature based methods for clustering are explored in non attribute graphs and demonstrate improved performance by capturing the geometrical aspect. It is an interesting idea to integrate this to attributed graphs which are commonly observed. Moreover, the proposed technique is very simple and easy to apply with graph neural networks. \n\n### Weaknesses\n* Although the technique is simple, it is computationally expensive for large dense graphs, limiting its practicality. \n* The writing is unclear in some places and the presentation can be improved by covering some preliminaries making the draft self contained (see requested changes)", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a new graph pooling operator, ORC Pool, which is based on Olliver's Ricci Curvature. The main idea of ORC Pool is to integrate the features associated with the nodes in graphs. This allows for an effective application of ORC Pool in node clustering and graph classification tasks using graph neural networks with ORC Pool similar to min-cut based pooling layers. The experimental results show improved performance with the ORC Pool layer.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. The node clustering experiments were on homophilic graphs where the graph pooling has shown good performance. How does ORC Pool compare to other pooling operators on heterophilic graphs? Exploring this could provide valuable insights, and there are heterophilic node classification datasets like Cornell, Texas, Wisconsin, Chameleon, etc. \n\n2. Graclus pooling is closely related to ORC Pool. Including a comparison to graclus pooling would further strengthen the paper.\n\n3. Suggestion on writing improvements to the draft:\n* Abstract: \"In both settings\" - it is not clear what this means as there have been no two settings described before. \n* Page 2, Related Work: \"The SRC framework\" - please provide the expansion before using abbreviation. \n* What is \"putative clusters\"?\n* Please add $1_{i\\sim j}$ used in equation 4 to notation \n* Add definition of $\\hat{D}$ in equation 7\n* Please add a paragraph on the methods compared in the experiments such as DIFFPool, DMoN, TVPool", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "WmNSl4g03k", "reviewer": "Reviewer_RZnK", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n* **Novelty and Relevance**: The proposed method is new and highly relevant to the graph neural network and geometric deep learning communities.\n* **Technical Contribution**: The work is technically sound, with some theoretical properties of the proposed method discussed and proved.\n* **Performance**: The proposed method outperforms existing pooling methods on standard graph classification tasks.\n\n## Weaknesses\n\n* **Presentation**: The presentation and organization of the paper could be improved to better communicate the contribution of the work. Improving the flow and clarity of the text would make the work more accessible, especially for readers unfamiliar with the topic. (There are also several typographical errors, such as proper name capitalization (e.g., \"graclus\" on page 6) and inconsistent section and paragraph title capitalization.) The description of ORC-Pool is spread across Sections 2 and 3, making it difficult for readers to follow without reading these sections multiple times.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces ORC-Pool, a novel graph pooling method for graph neural networks (GNNs) that uses Ricci flow for coarsening graphs. The method demonstrates strong performance on standard benchmarks, outperforming other popular pooling techniques.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. While this work will be a valuable addition to the literature, the current presentation detracts from clearly communicating its content. Below are some suggestions/recommendations for improving readability:\n    * Proofread the manuscript for typographical errors.\n    * Be mindful of readers who are not familiar with the topic by, e.g., replacing \"SRC framework\" on page 2 with \"the select-reduce-connect framework (SRC)\".\n    * End Section 1 with an overview of the paper's organization (e.g., \"This paper is organized as follows: ...\").\n    * Add algorithm boxes or similar to Section 3 to summarize the key components of ORC-Pool.\n    * Perhaps move the discussion on graph cuts from Section 3 to Section 2.\n    * Add a table to the appendix that summarizes the select, reduce, and connect operations of DiffPool, MinCutPool, DMoN, TVPool, and ORC-Pool.\n2. Could you clarify: Are the trainable parameters in Eq. 8 optimized as part of the GNN architecture's end-to-end training by adding the loss function from Eq. 8 to the downstream task loss, or is Eq. 8 minimized separately with gradient propagation through the solution?\n3. Could you add the baseline \"None,\" representing training without pooling, to Table 2 for comparison?\n4. Could you add a link to an anonymized repository with your code?", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "KXC2l1C67A", "reviewer": "Reviewer_YHFn", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n* The paper presents novel ideas on how to incorporate topological information into graph pooling operations.\n* Extensive theoretical background and analysis of Ricci-Flow-based node clusterings are provided.\n* The paper is well-written and easy to follow.\n\n\n**Weaknesses**:\n\nThe most significant weakness is a lack of discussion of potential limitations when predicting cluster assignments end-to-end with message-passing GNNs. When integrating ORC-Pool into a GNN, the cluster assignment $S$ is computed by an MLP that is applied to a vertex embedding $\\tilde{X}$ which is computed by a message-passing GNN:\n\n$S = MLP(\\tilde{X}, \\psi)$\n\nIf I understand this correctly, this is a row-wise operation that applies the same MLP independently to the embedding $\\tilde{X}(v)$ of each vertex $v$. Since $\\tilde{X}$ is computed through message passing, there are inherent restrictions that apply to the embeddings $\\tilde{X}$ and assignment $S$. In particular, if the GNN has $\\ell$ message passing layers, then $\\tilde{X}(v)$ is fully determined by the structure of the depth-$\\ell$ subtree rooted in $v$. If two vertices $u,v \\in V$ have structurally identical $\\ell$-hop neighborhoods, then $\\tilde{X}(u)=\\tilde{X}(v)$ and therefore the corresponding cluster assignments $S(u)=S(v)$ are also identical by construction. This will always hold true regardless of the model parameters and training objective. The assignment $S$ computed by a GNN in this manner may, therefore, have very different properties from assignments that are optimal with respect to equation 7. In particular:\n1. The graph from Figure 4 can not be properly clustered by a message-passing GNN since for each vertex $v$, we can find a vertex $u$ from another cluster with an identical local structure (i.e., the same 1-WL color).\n2. In a $k$-regular graph we would always have $S(v)=S(u)$ for all $v,u \\in V$ regardless of structure and curvature.\n\nTherefore, it is important to differentiate between the properties of a theoretical clustering that is optimal with respect to Ricci Flow and a clustering that is actually computed by a GNN. Currently, the paper lacks a clear discussion of this difference and its implications.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces ORC-Pool, a graph-pooling operator that leverages discrete Ricci curvature to cluster vertices and capture multi-scale graph structures while accounting for both the topology and the vertex attributes of a graph. It further integrates ORC-Pool as a trainable module within Message-Passing Graph Neural Networks. The theoretical properties of the proposed Ricci-Flow-based clustering are formally examined, and experiments demonstrate ORC-Pool’s competitive empirical performance when compared to alternative graph-pooling methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "One critical change is a clear discussion on the point raised in the weakness section about what restrictions apply to the GNN-based cluster assignments S.\n\nBeyond this, some additions to the experiments would further strengthen this work:\n* It would be helpful to visualize cluster assignments that are actually predicted by a trained GNN model.\n* It would also be useful to add results for GNNs without any pooling as a reference baseline in Table 2.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper introduces a trainable pooling operator that extending Ricci flow based clustering to attributed graphs. Building on this, the paper presents a pooling layer for incorporating the proposed operator within a message-passing graph neural network. The paper also discusses the structural properties of the pooling operator.

Claims: The key claim is that the provided numerical experiments demonstrate that the introduced pooling layers lead to improvements in several node-based and graph-level inference tasks. 

The reviewers agree that the presented experiments are sufficient to support the claim, and this is in line with my assessment. In most of the experiments, a simple GCN is used as the base GNN, although some experiments using GAT and GIN are included in the appendix. All three of these base GNNs are far from the state-of-the-art in terms of performance, so although the proposed pooling layer leads to a significant performance improvement, it is not so clear that this would carry over to scenarios where a more recent GNN is employed.

Audience: All reviewers concurred that the presented work is technically correct and is likely to be of interest to some members of the TMLR audience.
Decision: Accept (poster)