Understanding over-squashing and bottlenecks on graphs via curvature

Most graph neural networks (GNNs) suffer from information distortion when propagating messages across graphs that possess topological bottlenecks, leading to a phenomenon termed 'over-squashing'. This paper introduces a geometric framework using combinatorial edge-based curvature to precisely describe how negatively curved edges cause these bottlenecks and proposes a novel curvature-based graph rewiring method, Stochastic Discrete Ricci Flow (SDRF), to alleviate the over-squashing issue.

The gist

Negatively curved edges are responsible for the formation of bottlenecks in graph neural networks, which leads to the over-squashing of information.

Analysis of Over-Squashing and Bottlenecks

The paper first formalizes over-squashing by analyzing the Jacobian of node representations, proposing it as an explicit and formal way of assessing the over-squashing effect. Lemma 1 establishes that if message functions have bounded derivatives, the propagation is controlled by a power of the augmented normalized adjacency matrix, leading to exponential decay in dependence on input features at distance r when certain topological conditions are met. The sensitivity analysis connects this Jacobian to graph topology via powers of the augmented normalized adjacency matrix.

Graph Curvature and Bottleneck Induction

To understand how topology induces the bottleneck, the authors introduce a new combinatorial edge-based curvature called Balanced Forman curvature, which constitutes a sharp lower bound to the standard Ollivier curvature on graphs. The intuition is that positive curvature corresponds to complete graphs (cliques), zero curvature to grids (Euclidean geometry), and negative curvature to trees (hyperbolic geometry). Theorem 2 proves that for any edge i ∼ j, the Ollivier Ricci curvature is greater than or equal to the Balanced Forman curvature: κ(i, j) ≥ Ric(i, j).

Curvature-Based Rewiring Method

The core contribution is a new curvature-based method for graph rewiring called Stochastic Discrete Ricci Flow (SDRF). This method is designed to guide the rewiring steps in a surgical way by modifying the negatively-curved edges. Algorithm 1 describes SDRF: it iteratively adds an edge that improves the minimal Ricci curvature of an existing edge, and removes an edge with maximal positive curvature if its value exceeds a threshold C+. This process aims to homogenize edge curvatures, which is inspired by continuous Ricci flow.

Experimental Validation

The theoretical results are validated through experimental comparisons on nine standard graph learning datasets (Cornell, Texas, Wisconsin, Chameleon, Squirrel, Actor, Cora, Citeseer and Pubmed). The experiments demonstrate that SDRF consistently improves upon baseline methods like DIGL and +FA in terms of node classification accuracy. Crucially, the paper shows that SDRF preserves the graph topology to a far greater extent than DIGL due to its surgical nature, resulting in better preservation of the degree distribution compared to diffusion-based rewiring schemes. The results suggest that curvature-based rewiring is a viable candidate for improving GNN performance, especially on low-homophily datasets where long-range dependencies are more relevant.

Conclusion and Implications

The paper concludes that negatively curved edges are the ones causing bottlenecks and thus leading to over-squashing, providing a mechanism to address this issue. By linking local curvature properties to the Jacobian of node representations, the authors establish that negatively-curved edges are responsible for over-squashing. The proposed SDRF method is shown to be superior because it surgically targets these problematic edges while maintaining structural integrity, unlike diffusion-based rewiring methods which might fail to correct the bottleneck by only acting on short diffusion distance nodes. This opens the door for curvature-based rewiring methods in GNNs.

Limitations and Future Directions

The current theoretical results do not extend to multigraphs, and the methodology remains agnostic to information beyond graph topology, such as node features. Future work will focus on developing a notion of curvature that can incorporate these feature-dependent aspects. The paper also relates betweenness centrality to the Jacobian, proposing it as a topological characterization of bottleneckedness in graphs.

Key Findings Summary

1. Over-squashing is formally measured by the Jacobian of node representations, which is shown to be controlled by powers of the augmented normalized adjacency matrix.

2. Negative curvature (as defined by Balanced Forman curvature) is directly responsible for inducing graph bottlenecks and over-squashing in GNNs.

3. Stochastic Discrete Ricci Flow (SDRF) effectively alleviates the bottleneck by surgically modifying negatively-curved edges, outperforming random-walk based rewiring methods on low-homophily datasets.

4. SDRF preserves the graph topology better than diffusion methods, leading to a more stable and potentially more efficient downstream GNN performance.

5. A positive lower bound on curvature guarantees an upper bound on the diameter of the graph, linking curvature control to network structure (Proposition 5).

As a fastidious researcher, I have analyzed this paper, which introduces a geometric perspective (using edge-based curvature) to diagnose and mitigate over-squashing in Graph Neural Networks (GNNs).

Here are the specific improvements and capabilities derived from this research that can be implemented in AI systems:

The core improvement lies in shifting from generic graph rewiring to a mathematically grounded, curvature-guided structural modification process. This allows for targeted alleviation of information distortion.

I propose implementing a two-stage system:

1. A diagnostic module using the proposed curvature metrics to identify bottlenecks.

2. A corrective module utilizing the Stochastic Discrete Ricci Flow (SDRF) algorithm to surgically modify the graph topology based on those diagnoses.

This leads to an improved AI system capable of:

1. Developing GNNs that maintain high performance when processing graphs with complex, long-range dependencies (e.g., social networks, knowledge graphs).

2. Achieving better generalization across different graph structures by explicitly addressing topological limitations rather than relying solely on feature learning.

3. Improving computational efficiency by preserving desirable structural properties (like degree distribution) during topology modifications, ensuring the resulting GNN remains tractable for downstream tasks.

Here are the specific technical improvements and capabilities:

1. The system will incorporate a diagnostic tool based on the proposed metric:

2. Calculate the local Balanced Forman curvature for every edge in a graph.

3. Identify edges with high negative curvature, as these are mathematically proven to be responsible for information bottlenecks (over-squashing).

4. Implement the corrective mechanism using the SDRF algorithm (Algorithm 1):

5. Surgically rewire the graph by iteratively adding edges that support negatively curved regions and removing edges that contribute excessive positive curvature, guided by a temperature parameter to control stochasticity.

This process allows for:

1. Mitigating over-squashing in deep GNNs, ensuring that representations of distant nodes remain robust against distortion during message passing (as shown by the reduction in the Jacobian of node representations).

2. Improving performance on low-homophily datasets where long-range dependencies are crucial, as the curvature-based approach is empirically shown to outperform diffusion-based rewiring methods (DIGL).

3. Preserving critical structural properties: The SDRF method is designed to be surgical, meaning it modifies edges locally around the bottleneck without drastically changing the overall degree distribution of important nodes, leading to better computational stability for the final GNN layer.

4. Achieving theoretically sound bounds on graph bottleneckedness (Cheeger constant, spectral gap) by relating local curvature to global topological properties, allowing for a more principled approach to graph preprocessing than simple random-walk methods.