Understanding over-squashing and bottlenecks on graphs via curvature

Most graph neural networks ( GNNs ) use the message passing paradigm , in which node features are propagated on the input graph . Recent works pointed to the distortion of information flowing from distant nodes as a factor limiting the efficiency of message passing for tasks relying on long-distance interactions . This phenomenon , referred to as ‘ over-squashing ’ , has been heuristically attributed to graph bottlenecks where the number of k-hop neighbors grows rapidly with k. We provide a precise description of the over-squashing phenomenon in GNNs and analyze how it arises from bottlenecks in the graph . For this purpose , we introduce a new edge-based combinatorial curvature and prove that negatively curved edges are responsible for the over-squashing issue . We also propose and experimentally test a curvature-based graph rewiring method to alleviate the over-squashing . 1 INTRODUCTION In the past few years , deep learning on graphs and in particular graph neural networks ( GNNs ) ( Sperduti , 1994 ; Goller & Kuchler , 1996 ; Sperduti & Starita , 1997 ; Frasconi et al. , 1998 ; Gori et al. , 2005 ; Scarselli et al. , 2008 ; Bruna et al. , 2014 ; ? ; Kipf & Welling , 2017 ; Gilmer et al. , 2017 ) have become very popular in the machine learning community due to their ability to deal with broad classes of systems of relations and interactions , ranging from social networks to particle physics ( Shlomi et al. , 2021 ) . The vast majority of GNNs follow the message passing paradigm ( Gilmer et al. , 2017 ) , using learnable non-linear func- tions to diffuse information on the graph . Multiple popular GNN architectures such as GCN ( Kipf & Welling , 2017 ) and GAT ( Veličković et al. , 2018 ) can be posed as particular flavors of this scheme and considered instances of a more general framework of geometric deep learning ( Bronstein et al. , 2021 ) . Some of the drawbacks of the message passing paradigm have now been identified and formalized , including the limits of expressive power ( Xu et al. , 2019 ; Morris et al. , 2019 ; Maron et al. , 2019 ) and the problem of over-smoothing ( NT & Maehara , 2019 ; Oono & Suzuki , 2020 ) . On the other hand , much less is known about the phenomenon of over-squashing , consisting in the distortion of messages being propagated from distant nodes . Alon & Yahav ( 2021 ) proposed rewiring the graph as a way of reducing the bottleneck , defined as those topological properties in the graph leading to over-squashing . This approach is in line with multiple other results e.g . using connectivity diffusion ( Klicpera et al. , 2019 ) as a preprocessing step to facilitate graph learning . Yet , the exact understanding of the over-squashing and how it originates from the bottlenecks in the topology of the underlying graph are still elusive . Consequently , there is currently no consensus on the right method ( either based on graph rewiring or not ) to address the bottleneck and hence alleviate the over-squashing . In this paper , we address these questions using tools from differential geometry , which traditionally is concerned with the study of manifolds . It offers an appealing framework to study the properties of graphs , in particular arguing that graphs , like manifolds , exhibit curvature that makes them more suitable to be realized in spaces with hyperbolic geometry ( Liu et al. , 2019 ; Chami et al. , 2019 ; Boguna et al. , 2021 ) . One notion of curvature that has received attention for graph learning is Ricci curvature ( Hamilton , 1988 ) , also known in geometry for its use in Ricci flow and the subsequent proof of the Poincaré conjecture ( Perelman , 2003 ) . Certain graph analogues of the Ricci curvature ( Forman , 2003 ; Ollivier , 2009 ; Sreejith et al. , 2016 ) were used in Ni et al . ( 2018 ) for a discrete version of Ricci flow to construct a metric between graphs . Graph Ricci flow was also used in Ni et al . ( 2019 ) for community detection . Both of these methods use the edge weights as a substitute for the metric of a manifold , and do not change the topological structure of the graph . Contributions and Outline . This paper , to our knowledge , is the first theoretical study of the bottleneck and over-squashing phenomena in message passing neural networks from a geometric perspective . In Section 2 , we propose the Jacobian of node representations as a formal way of measuring the over-squashing and we show that the graph topology may compromise message propagation in graph neural networks by creating a bottleneck . In Section 3 , we investigate how such a bottleneck is induced which leads to the over-squashing of information . To this aim , we introduce a new combinatorial edge-based curvature called Balanced Forman curvature that constitutes a sharp lower bound to the standard Ollivier curvature on graphs , and prove that negatively curved edges are responsible for the formation of bottlenecks ( and hence for over-squashing ) . In Section 4 , we present a new curvature-based method for graph rewiring called Stochastic Discrete Ricci Flow . According to the theoretical results in Section 3 , this rewiring method is suited to address the graph bottleneck and hence alleviate the over-squashing by surgically targeting the edges responsible for the issue . By contrast , we rigorously show that a recently introduced diffusion-based rewiring scheme might generally fail to reduce the bottleneck . Finally , in Section 5 , we compare different rewiring strategies experimentally on several standard graph learning datasets . 2 ANALYSIS OF THE OVER-SQUASHING PHENOMENON . 2.1 PRELIMINARIES . Let G = ( V , E ) be a simple , undirected , and connected graph , where ( i , j ) ∈ E iff i ∼ j . We focus on the unweighted case , although the theory extends to the weighted setting as well . We denote the adjacency matrix by A and let Ã = A+ I be the adjacency matrix augmented with self-loops . Similarly we let D̃ = D + I , with D the diagonal degree matrix , and let Â = D̃− 1 2 ÃD̃− 1 2 be the normalized augmented adjacency matrix ( self-loops are commonly included in GNN architecture , and in Section 2 we formally explain why GNNs are expected to propagate information more reliably when self-loops are taken into account ) . Given i ∈ V , we denote its degree by di and let Sr ( i ) : = { j ∈ V : dG ( i , j ) = r } , Br ( i ) : = { j ∈ V : dG ( i , j ) ≤ r } , where dG is the standard shortest-path distance on the graph and r ∈ N. The set Br ( i ) represents the receptive field of an r-layer message passing neural network at node i . Message passing neural networks ( MPNNs ) . Assume that the graph G is equipped with node features X ∈ Rn×p0 where xi ∈ Rp0 is the feature vector at node i = 1 , . . . , n = |V | . We denote by h ( ` ) i ∈ Rp ` the representation of node i at layer ` ≥ 0 , with h ( 0 ) i = xi . Given a family of message functions ψ ` : Rp ` × Rp ` → Rp ′ ` and update functions φ ` : Rp ` × Rp ′ ` → Rp ` +1 , we can write the ( ` + 1 ) -st layer output of a generic MPNN as follows ( Gilmer et al. , 2017 ) : h ( ` +1 ) i = φ ` h ( ` ) i , n∑ j=1 Âijψ ` ( h ( ` ) i , h ( ` ) j ) . ( 1 ) Here we have used the augmented normalized adjacency matrix to propagate messages from each node to its neighbors , which simply leads to a degree normalization of the message functions ψ ` . To avoid heavy notations the node features and representations are assumed to be scalar from now on ; these assumptions simplify the discussion and the vector case leads to analogous results . 2.2 THE OVER-SQUASHING PROBLEM . Multiple recent papers observed that MPNNs tend to perform poorly in situations when the learned task requires long-range dependencies and at the same time the structure of the graph results in exponentially many long-range neighboring nodes . We say that a graph learning problem has long-range dependencies when the output of a MPNN depends on representations of distant nodes interacting with each other . If long-range dependencies are present , messages coming from nonadjacent nodes need to be propagated across the network without being too distorted . In many cases however ( e.g . in ‘ small-world ’ graphs such as social networks ) , the size of the receptive field Br ( i ) grows exponentially with r. If this occurs , representations of exponentially many neighboring nodes need to be compressed into fixed-size vectors to propagate messages to node i , causing a phenomenon referred to as over-squashing of information ( Alon & Yahav , 2021 ) . In line with Alon & Yahav ( 2021 ) , we refer to those structural properties of the graph that lead to over-squashing as a bottleneck1 . Sensitivity analysis . The hidden feature h ( ` ) i = h ( ` ) i ( x1 , . . . , xn ) computed by an MPNN with ` layers as in equation 1 is a differentiable function of the input node features { x1 , . . . , xn } as long as the update and message functions φ ` and ψ ` are differentiable . The over-squashing of information can then be understood in terms of one node representation h ( ` ) i failing to be affected by some input feature xs of node s at distance r from node i . Hence , we propose the Jacobian ∂h ( r ) i /∂xs as an explicit and formal way of assessing the over-squashing effect2 . Lemma 1 . Assume an MPNN as in equation 1 . Let i , s ∈ V with s ∈ Sr+1 ( i ) . If |∇φ ` | ≤ α and |∇ψ ` | ≤ β for 0 ≤ ` ≤ r , then ∣∣∣∣∣∂h ( r+1 ) i∂xs ∣∣∣∣∣ ≤ ( αβ ) r+1 ( Âr+1 ) is . ( 2 ) Lemma 1 states that if φ ` and ψ ` have bounded derivatives , then the propagation of messages is controlled by a suitable power of Â . For example , if dG ( i , s ) = r + 1 and the sub-graph induced on Br+1 ( i ) is a binary tree , then ( Âr+1 ) is = 2−13−r , which gives an exponential decay of the node dependence on input features at distance r , as also heuristically argued by Alon & Yahav ( 2021 ) . The sensitivity analysis in Lemma 1 relates the over-squashing – as measured by the Jacobian of the node representations – to the graph topology via powers of the augmented normalized adjacency matrix . In the next section we explore this connection further by analyzing which local properties of the graph structure affect the right hand side in equation 2 , hence causing the bottleneck . We will address this problem by introducing a new combinatorial notion of edge-based curvature and showing that negatively curved edges are those responsible for the over-squashing phenomenon .

Summary: In this paper, the authors work on the over-squashing effect that has been recently observed in the GCN literature, namely the effect that, in a message-passing paradigm and for a learning problem that necessitates long-range interaction between distant nodes, the existence of bottlenecks in the graph (for instance an edge e with very large betweenness centrality) will intuitively distort the information that needs to travel from distant nodes and thus hinder the GCN's performance.  One natural direction of research is to add edges (and sometimes remove others in order to keep the complexity under control) to alleviate the bottlenecks. This rewiring process may be done in different ways and the authors suggest one particular algorithm that they compare with state-of-the-art.  In particular, the contributions are: - a precise (even if somewhat arbitrary) definition of an over-squashing measure between two nodes $i$ and $s$ that is simply the Jacobian of the node representation at $i$ with respect to the entry $x_s$. The smaller this measure, the less $i$ "feels" $s$, the larger the over-squashing effect. The goal of the paper is to find a way to understand how to alleviate the bottlenecks responsible for this effect  - to this end, the authors suggest a new definition of Ricci-like curvature defined over the edges of the graphs. This new definition has one main property which is theorem 3: it lower bounds the Ollivier curvature. This then enables to obtain Cor. 4, and finally the main result Thm 5, that states, in a nutshell, that negatively curved edges are the ones causing bottlenecks.  - a concrete rewiring algorithm is suggested, in which two steps are repeated until convergence or max iteration is reached: a first stochastic step is performed where an edge is added to alleviate the edge that has minimal curvature, before a deterministic step where the edge with maximal curvature (typically those in cliques) is removed if it is larger than a threshold  - a tentative analysis, in the form of Thm 8, as well as experiments show the authors' method performs as well or outperforms other rewiring methods (they mainly compare with DIGL) while keeping the number of total edges after rewiring under control