🤖 machine learning

Pair-Centric Graph Rewiring for Over-Squashing via Optimal Transport-Guided Communication Alignment

This paper introduces PairAlign, a pair-centric graph rewiring framework that addresses over-squashing in message-passing neural networks by identifying poorly supported pairwise communications via a demand-support shortage metric and optimizing edge additions using Optimal Transport to align structural capacity with communication needs.

Original authors: Yan Wang, Chuan-Xian Ren

Published 2026-08-12

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Yan Wang, Chuan-Xian Ren

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to send a secret message across a crowded city using a chain of friends. You whisper to your neighbor, who whispers to theirs, and so on, until the message reaches the other side of town. This is how computers called "Graph Neural Networks" (or MPNNs) learn from data that looks like a map of connections, such as social networks or chemical molecules. They pass information from one node (a person or atom) to its immediate neighbors, layer by layer.

But here's the problem: if the message has to travel a long way, it gets squeezed. Imagine trying to fit a whole novel into a single postcard. As the message hops from friend to friend, the computer has to compress all that distant information into a tiny, fixed-size summary. If the path is too long or the "roads" between friends are too narrow, the message gets distorted or lost entirely. Scientists call this "over-squashing." It's like trying to pour a firehose of water into a thimble; the most important details get crushed out. To fix this, researchers have tried to "rewire" the city—adding new shortcuts or removing bad roads—to help the message flow better. But the big question remains: which new roads should we build when we only have a limited budget?

This paper introduces a clever new way to answer that question, called PairAlign. Instead of just looking at the whole city or fixing one specific traffic jam at a time, PairAlign zooms in on specific pairs of friends who are struggling to talk to each other. It calculates a "shortage score" for every pair: how much they need to communicate versus how much the current map actually allows them to. If two friends are far apart and the path between them is a bottleneck, they get a high shortage score.

The authors suggest that the best way to fix over-squashing is to treat it like a delivery service with a limited number of trucks. You don't just drop packages randomly; you match your trucks to the most desperate deliveries. PairAlign does this using a mathematical tool called Optimal Transport. Think of it as a super-smart dispatcher that looks at all the "shortage" pairs (the ones crying out for help) and all the possible new roads you could build. It then figures out the perfect way to assign your limited road-building budget so that the most critical conversations get the most support.

The paper shows that this method works. When they tested PairAlign on standard graph benchmarks (like citation networks and chemical datasets), it consistently improved the computer's ability to learn compared to other methods. The results suggest that by focusing on the specific pairs that are suffering from over-squashing and using a global strategy to allocate new connections, we can help these networks "hear" distant signals much more clearly. It's not just about adding more roads; it's about building the right roads for the people who need them most.

Technical Summary: Pair-Centric Graph Rewiring for Over-Squashing via Optimal Transport-Guided Communication Alignment

1. Problem Statement

Message-passing neural networks (MPNNs) struggle when task-relevant information is distributed across distant regions of a graph. This phenomenon, known as over-squashing, occurs because local propagation must compress remote signals through limited structural interfaces, causing information to be attenuated or lost. While graph rewiring has emerged as a structural solution to mitigate over-squashing, existing methods primarily rely on:

Edge-level metrics: Local geometric cues (e.g., curvature) or bottleneck scores.
Graph-level surrogates: Global connectivity measures (e.g., spectral expansion, effective resistance).

The paper argues that these approaches leave a critical gap: under a limited rewiring budget, it is unclear which specific pairwise communications most urgently require structural support. Graph-level improvements may enhance aggregate connectivity without addressing the most constrained node-pair interactions, while edge-level repairs may miss the broader distribution of communication needs.

2. Methodology: PairAlign

The authors propose PairAlign, a pair-centric graph rewiring framework that explicitly targets the "demand-support shortage" between node pairs. The methodology consists of three core components:

A. Pairwise Communication Shortage Formulation

PairAlign defines over-squashing risk as a ratio between structural demand and propagation support:

Demand ( $\omega$ ): Defined by the shortest-path distance on the original graph ( $d_G(u, v)$ ), raised to a power $p$ . This captures the inherent difficulty of communicating between distant nodes.
Support ( $s$ ): Defined by the finite-hop propagation mass on the current computational graph, calculated as a weighted sum of powers of the row-normalized propagation matrix $P$ .
Shortage Score ( $S$ ): The ratio $S(u, v) = \omega(u, v) / (s(u, v) + \epsilon)$ $S (u, v) = ω (u, v) / (s (u, v) + ϵ)$ .
- Theoretical Basis: The paper proves that this computable shortage score serves as a proxy for the Jacobian-based shortage (the normalized influence of node $u$ on node $v$ ). Specifically, $S$ provides a lower bound and constant-factor equivalence to the Jacobian-based metric under Lipschitz assumptions, making the intractable Jacobian analysis feasible through graph structure.

B. Edge Insertion Dynamics

The framework acknowledges that adding an edge is not monotonically beneficial. A new edge creates useful walks but simultaneously dilutes the transition mass of existing neighbors due to row normalization.

First-Order Analysis: The authors derive the first-order change in support ( $g_e$ ) for a target pair when an edge is inserted.
Selection Criterion: An edge is only useful if the "added-path contribution" exceeds the "normalization loss." PairAlign dynamically recomputes shortage during the rewiring process to ensure selected edges actually reduce the shortage score.

C. Optimal Transport (OT)-Guided Allocation

To manage a finite edge budget, PairAlign avoids greedy local selection, which tends to concentrate edges on easily fixable pairs while leaving severe shortages uncovered. Instead, it formulates the allocation as an Optimal Transport problem:

Source: The distribution of candidate added edges ( $q$ ).
Target: The distribution of high-shortage node pairs ( $p$ ), derived by centering the original shortage scores to emphasize severe deficits.
Transport Cost ( $C$ ): Encodes structural compatibility between a candidate edge and a target pair, considering:
1. Endpoint Alignment: Distance between edge endpoints and target pair nodes.
2. Bridge Suitability: Whether the edge's span is commensurate with the target pair's distance.
Objective: The total loss combines the shortage reduction on the rewired graph and the OT alignment cost. This ensures the budget is distributed globally to cover the most critical shortage targets rather than clustering locally.

D. Optimization

The method uses a continuous relaxation with logits for candidate edges, employing a straight-through estimator for gradient propagation. It iteratively updates edge scores, re-evaluates the shortage on the current rewired graph, and projects the final scores to a discrete set of added edges.

3. Key Contributions

Pair-Centric Formulation: Introduces a demand-support shortage metric that ranks node pairs by the lack of current-graph support relative to original structural demand. Theoretically, this provides a tractable proxy for Jacobian-based over-squashing.
OT-Guided Rewiring: Proposes a framework that coordinates the limited edge budget across shortage targets using Optimal Transport. This prevents budget concentration on a few targets and ensures broader coverage of critical communication bottlenecks.
Empirical Validation: Demonstrates that PairAlign improves performance across standard graph benchmarks (node and graph classification) and various backbones (GCN, GIN, GAT), validating pair-level repair as an effective route for alleviating over-squashing.

4. Experimental Results

The paper evaluates PairAlign (PAR) against state-of-the-art rewiring methods (SDRF, BORF, FoSR, GTR, LASER, etc.) on:

Node Classification: Standard datasets (Cora, Citeseer) and heterophilic benchmarks (Roman-Empire, Amazon-Ratings, etc.).
Graph Classification: TUDatasets (ENZYMES, MUTAG, etc.).

Key Findings:

Performance: PairAlign achieves the best average ranking across all datasets and backbones. For instance, on the Texas dataset with a GIN backbone, accuracy improved from 53.5% to 68.8%.
Structural Diagnostics: PAR shows superior $\Delta$ Shortage (reduction in communication shortage) and Coverage@10 (fraction of top shortage pairs receiving support) compared to greedy baselines.
Correlation: The proposed shortage score correlates strongly (Spearman $\rho \approx 0.93$ ) with pairwise effective resistance, confirming its validity as a bottleneck diagnostic.
Heterophily: On heterophilic graphs, PAR provides stable improvements without the indiscriminate mixing that harms class-relevant structure, outperforming methods that rely solely on feature similarity or community structure.

5. Significance and Claims

The paper claims that PairAlign offers a theoretically grounded and practically effective approach to graph rewiring by shifting the focus from local edges or global connectivity to pairwise communication compatibility.

Theoretical Insight: It bridges the gap between the abstract Jacobian view of over-squashing and practical structural modification, proving that optimizing the computable shortage score effectively reduces the Jacobian-based bottleneck.
Mechanism: It demonstrates that Optimal Transport is necessary to coordinate limited structural resources, ensuring that rewiring efforts are distributed to the most critical node pairs rather than being captured by local greedy heuristics.
Generalizability: The method is presented as a structural preprocessing step that is compatible with various message-passing backbones, suggesting that improving the underlying graph topology's ability to support specific pairwise interactions is a universal strategy for mitigating over-squashing.

The authors conclude that while the method is effective, future work should explore task-aware demand modeling and explicit locality constraints to improve scalability on very large graphs.