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: The authors propose a new graph rewiring approach that utilizes a discrete notion of Ricci curvature to mitigate over-squashing. This is motivated by a link between negatively curved edges and graph bottlenecks. The paper has a theoretical focus, but also provides a set of validation experiments to demonstrate the proposed approach.