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

This paper proposes PairAlign (PAR), a pair-centric graph rewiring framework that addresses the over-squashing problem in message-passing neural networks (MPNNs) by treating it as a pair-level communication shortage. The authors, Yan Wang and Chuan-Xian Ren from Sun Yat-sen University, frame the core question as: With a limited rewiring budget, the key question is which pairwise communications most need structural support.

The paper notes that "Message-passing neural networks (MPNNs) often struggle when task-relevant information is distributed across distant regions of a graph, since local propagation must compress remote signals through limited structural interfaces. Existing rewiring methods fall into two main categories: graph-level methods that improve aggregate connectivity and edge-level methods that repair local bottleneck structures. However, these approaches leave a practical allocation problem under limited budgets: structural support must be assigned to the node-pair communications that need it most. The authors argue that Without an explicit pair-level notion of communication shortage, rewiring can improve aggregate connectivity while leaving the most constrained pairwise communications only indirectly supported."

PairAlign defines three key concepts:

Pairwise Communication Demand (Definition 1): "The pairwise communication demand is the function ω: P → R≥0 defined by ω(u, v) = dG(u, v) p, where dG(u, v) is the shortest-path distance between u and v on G, and p > 0 controls the growth rate of communication demand with respect to distance."

Pairwise Propagation Support (Definition 2): The finite-hop pairwise propagation support on W is the function s(·; W): P → R≥0 defined by s(u, v; W) = Σ l=1 K α l [P l] uv, where α 1,..., α K ∈ R≥0 are nonnegative hop weights satisfying Σ l=1 K α l = 1.

Pairwise Communication Shortage (Definition 3): "The pairwise communication shortage on W is defined as S(u, v; W) = ω(u, v) / (s(u, v; W) + ε), where ε > 0 is a small positive constant."

The paper provides several theoretical results:

Lemma 1 shows that the support s(u, v; W) upper bounds the normalized finite-depth Jacobian influence I(u, v) for K-layer MPNNs with Lipschitz message and update functions.

Proposition 1 establishes that the computable shortage S(W) is a proxy for the Jacobian-based shortage S Jac(W), showing that S(W) ≤ S Jac(W) and, in the non-degenerate regime where I(u, v) ≥ c·s(W), S(W) ≤ S Jac(W) ≤ c-1 S(W).

Proposition 2 reveals a two-sided effect of edge insertion: a new edge can create useful walks and simultaneously dilute existing normalized transition mass. The first-order support change is decomposed as g e(u, v) = B e(u, v) − D e(u, v), where B e is the added-path contribution and D e is the normalization loss over the original neighbors of a.

Proposition 3 formalizes the coverage advantage of OT-guided allocation over greedy-local assignment, showing that with probability at least 1 − δ, Cov p(Γ OT) − Cov p(Γ g) ≥ TV(π C B, p) − sqrt(T log(2T/δ) / (2n)).

PairAlign formulates the communication alignment between the candidate edge budget and the shortage targets as an Optimal Transport problem. The OT formulation is:

Γ* = arg min Γ∈Π(q(B),p) ⟨Γ, C⟩ − εH(Γ)

where q(B) is the added-budget distribution, p is the shortage-target distribution, and C is the transport cost matrix. The cost matrix combines two terms:

• Endpoint alignment: min(d(a, u) + d(b, v), d(a, v) + d(b, u))

• Bridge suitability: −λψ(d(a, b) / (d(u, v) + ε))

The total objective combines shortage reduction with OT-guided alignment: min B L shortage(A + B) + λ OT L OT(B).

The optimization uses continuous logits on candidate non-edges with a top-k hard mask applied in the forward pass and gradients propagated through the corresponding soft variables by a straight-through style approximation. The algorithm has three phases: shortage-aware pairwise target construction, budgeted near-discrete rewiring, and final projection to the discrete added-edge set.

The paper evaluates PairAlign on three benchmark groups:

Node Classification: PAR achieves the best average rank under GCN and GIN, reaching 1.50 and 1.00. Under GIN, Texas rises from 53.5 to 68.8 and Cornell from 36.5 to 51.0.

Graph Classification: PAR attains the top result on five of the six datasets under GCN, with clear gains on ENZYMES and MUTAG.

Heterophilous Node Classification: PAR obtains the best average rank under both GCN and GAT, improving over the closest competitor ComFy in each case.

The ablation study shows that the full model gives the clearest evidence for this mechanism, with PAR obtaining the best task performance on Texas, Roman-Empire, and MUTAG, while also producing the largest ∆Shortage and Coverage@10. The paper demonstrates that Greedy-Local exhibits a more concentrated allocation pattern while OT-guided allocation spreads the budget over a wider target set and produces stronger pair-centric structural repair.

The paper concludes: "This work proposes PairAlign (PAR), an OT-guided graph rewiring framework that treats over-squashing as a pair-level communication shortage. PAR centers rewiring on node pairs whose structural demand is high but whose propagation support remains limited. Future work includes task-aware demand modeling and explicit locality constraints for improving scalability on large graphs."

Improved AI System: Pair-Centric Communication-Aware Graph Neural Networks

Based on PairAlign, I can implement the following specific improvements:

1. Pairwise Shortage-Aware Message Routing: Replace uniform message aggregation in MPNNs with a shortage-weighted routing mechanism. The system computes S(u,v;W) = d(u,v) p / (Σ α l[P l] uv + ε) for all node pairs and dynamically adjusts attention weights to prioritize high-shortage pairs, enabling the model to allocate computational resources to structurally constrained communications.

2. Budgeted Optimal Transport Edge Selection: Implement an OT-based edge rewiring module that, given a rewiring budget B, solves min Γ ⟨Γ,C⟩ − εH(Γ) where C combines endpoint distance and bridge suitability. This allows the system to select edges that maximally reduce communication shortage across the graph, rather than relying on local greedy heuristics that miss global bottlenecks.

3. Two-Sided Edge Insertion Evaluation: Incorporate Proposition 2's decomposition g e(u,v) = B e(u,v) − D e(u,v) into the edge selection process. The system can now explicitly evaluate both the added-path benefit and the normalization loss of each candidate edge, rejecting edges that create new shortcuts but dilute existing transition mass disproportionately.

4. Jacobian-Influence Proxy for Training: Use Lemma 1's bound (s(u,v;W) ≥ I(u,v)) as a differentiable training signal. The system can optimize a surrogate loss that directly maximizes propagation support for high-demand pairs, improving gradient flow to distant nodes without requiring expensive Jacobian computations.

5. Adaptive Demand Growth Control: Implement the tunable parameter p in ω(u,v) = d(u,v) p as a learned, task-adaptive hyperparameter. The system can adjust p during training to control how aggressively it prioritizes long-range communications, automatically balancing between local feature learning and global information integration.

What the improved system can do:

• Achieve superior performance on heterophilous graphs (e.g., Texas, Cornell) where task-relevant information is distributed across distant, structurally constrained regions

• Scale to large graphs by using the OT-guided allocation to focus rewiring budget on the most impactful pairs, avoiding exhaustive pairwise computations

• Provide interpretable shortage maps (S(u,v;W)) that explain which communications are most constrained, enabling debugging of model failures on long-range dependencies

• Dynamically rewire graphs during training to adapt to evolving task demands, rather than using static graph structures

• Generalize across node classification, graph classification, and heterophilous benchmarks with consistent top-rank performance (achieving rank 1.0 under GIN, 1.5 under GCN)

Message-passing neural networks (MPNNs) often struggle when task-relevant information is distributed across distant regions of a graph, since local propagation must compress remote signals through limited structural interfaces. Graph rewiring provides a structural response to over-squashing. Most existing methods rely on edge-level bottleneck scores or graph-level connectivity surrogates. With a limited rewiring budget, the key question is which pairwise communications most need structural support. This paper proposes PairAlign, a pair-centric graph rewiring framework that makes this question explicit through demand-support shortage. Specifically, PairAlign combines original-graph structural demand with current-graph finite-hop propagation support; their ratio highlights interactions whose communication demand is poorly supported by topology, and our theory shows that this score provides a computable proxy for the corresponding Jacobian-based shortage with a pair-level interpretation of over-squashing. Our theory reveals a two-sided effect of edge insertion: a new edge can create useful walks and simultaneously dilute existing normalized transition mass. Guided by this observation, PairAlign optimizes shortage to favor edge additions that alleviate over-squashing. Beyond selecting useful additions, PairAlign further introduces an Optimal Transport-guided rewiring mechanism to coordinate the finite edge budget for pair-level structural compatibility and shortage-target coverage. It formulates communication alignment between the candidate edge budget and the shortage targets, and the theory shows that this allocation covers shortage targets more broadly and effectively than a greedy-local assignment. Experiments on standard graph benchmarks show PairAlign's improvement across message-passing backbones, validating pair-level repair as an effective route for alleviating over-squashing.