Source-linked AI summary

DRew: Dynamically Rewired Message Passing with Delay

Benjamin Gutteridge, Xiaowen Dong, Michael Bronstein, Francesco Di Giovanni

arXiv:2305.08018v2cs.LGcs.AIstat.ML

TL;DR

Long-range tasks challenge MPNNs because local message passing requires depth, while distant nodes communicating instantly can discard graph-distance information. DRew rewires message passing by layer and adds distance-dependent delays. Across four benchmark tasks, DRew variants compare favorably with classical, static-rewiring, and multi-hop baselines, while GraphGPS leads the best DRew model on two tasks.

  • Problem

    Long-range interactions require MPNNs to stack layers of local message passing, while deep networks are prone to over-squashing; global rewiring can discard locality induced by graph distance.

    Source: §2.2 · p. 2, §3.1 · p. 4
  • Method

    DRew dynamically rewires message passing by layer, and νDRew adds a tunable delay for interactions between nodes at different distances.

    Source: §1 · p. 1, §3.1 · p. 4
  • Results

    Across all four long-range benchmark tasks, DRew outperforms static-rewiring and multi-hop benchmarks, and at least one DRew-MPNN matches or beats the best classical graph Transformer baseline; GraphGPS outperforms the best DRew model on PascalVOC-SP and Peptides-struct.

    Source: §5.1 · p. 7
  • Takeaways & Limitations

    The shortest-path-based multi-hop aggregations in νDRew distinguish any pair of graphs distinguished by the shortest path kernel, making (ν)DRew at least more expressive than 1-WL and classical MPNNs.

    Source: §4 · p. 6

    The authors do not expect DRew to be advantageous on homophilic node-classification tasks where shallow GNNs suffice, and the extent of long-range interactions in PascalVOC-SP is uncertain.

    Source: §7 · p. 9

Abstract

from arXiv · show

Message passing neural networks (MPNNs) have been shown to suffer from the phenomenon of over-squashing that causes poor performance for tasks relying on long-range interactions. This can be largely attributed to message passing only occurring locally, over a node's immediate neighbours. Rewiring approaches attempting to make graphs 'more connected', and supposedly better suited to long-range tasks, often lose the inductive bias provided by distance on the graph since they make distant nodes communicate instantly at every layer. In this paper we propose a framework, applicable to any MPNN architecture, that performs a layer-dependent rewiring to ensure gradual densification of the graph. We also propose a delay mechanism that permits skip connections between nodes depending on the layer and their mutual distance. We validate our approach on several long-range tasks and show that it outperforms graph Transformers and multi-hop MPNNs.

1. Introduction

Classical MPNNs pass information only between immediate neighbours, so long-range interactions require deeper networks that are prone to over-squashing and over-smoothing. DRew introduces layer-dependent message passing that preserves the graph’s distance bias, alongside a delay mechanism for adapting when nodes interact.

1. Introduction

  • Classical MPNNs pass information only between immediate neighbours, requiring deeper networks for distant nodes to share information.Deep MPNNs are particularly prone to over-squashing and over-smoothing.Source: §1 · p. 1
  • DRew is a layer-dependent framework for multi-hop message passing that can be applied to any MPNN architecture.It aims to support long-range information flow while respecting the graph’s distance-based inductive bias.Source: §1 · p. 1
  • The paper introduces νDRew, which adds a tunable delay to message passing to address issues such as oversmoothing in deep MPNNs.The delay mechanism generalizes DRew and is intended to adapt propagation to the task.Source: §1 · p. 1
  • The authors report theoretical analysis that their frameworks can mitigate over-squashing, and experiments on synthetic and real-world datasets showing robustness and effectiveness for deep MPNNs or long-range tasks.Source: §1 · p. 1

2. Message Passing Neural Networks

The section introduces graph and MPNN concepts before explaining why long-range tasks challenge local message passing. It motivates rewiring and delay as ways to adapt information flow to graph distance and network depth.

2.1. Preliminaries

  • The preliminaries assume an undirected, connected graph and describe locality using shortest-path distance, which partitions nodes into exact k-hop sets.The k-hop set contains nodes at distance exactly k, not all nodes reachable by a walk of length k.Source: §2.1 · pp. 1–2
  • Classical MPNNs exchange information between immediate neighbours at each layer, so nodes at distance k need at least k layers to interact.An MPNN layer applies learnable aggregation and update maps to node features and graph structure.Source: §2.1 · p. 2

2.2. Long-range dependencies and network depth

  • A task has long-range interactions when a node’s representation must account for information from another node at graph distance much greater than one.Source: §2.2 · p. 2
  • Over-squashing can arise when an expanding k-hop neighborhood sends exponentially many messages through fixed-size node representations, potentially losing information.Graph rewiring approaches may aggregate across multiple hops locally or connect all node pairs as in graph Transformers, which can discard distance-based locality.Source: §2.2 · pp. 2–3
  • Because depth-related performance degradation may reflect over-smoothing, over-squashing, or vanishing gradients, the section motivates message passing that adapts between faster and slower information flow.The proposed principles are dynamically rewired message passing and a delay mechanism.Source: §2.2 · p. 3

3. Dynamically Rewired MPNNs

Classical MPNNs make distant nodes communicate through intermediate neighbors, with interaction delayed by graph distance. DRew gradually adds distance-based connections, while νDRew also controls how delayed messages are; both changes retain the graph’s distance structure.

3. Dynamically Rewired MPNNs

  • Classical MPNNs make nodes at distance r communicate through intermediate neighbors, with interaction delayed by exactly r layers.This fixed, local flow can filter communication through intermediate nodes and contribute to over-squashing.Source: §3 · p. 3

3.1. A new framework: (ν)DRew message passing

  • DRew adds each k-hop neighborhood from layer k onward, so the receptive field grows by one hop per layer and distant nodes can interact directly.Unlike static multi-hop rewiring, it fills the computational graph gradually according to distance.Source: §3 · p. 3, §3.1 · p. 4
  • νDRew controls message delay with τν(k) = max(0, k − ν), letting nodes access past representations according to distance and the tunable parameter ν.With ν = ∞ there is no delay; the experiments treat ν as a tunable hyperparameter.Source: §3.1 · p. 4

3.2. Instances of our framework

  • The framework supplies instances for GCN, GIN, and GatedGCN, with neighborhood weight sharing used for GatedGCN but not GCN or GIN.For GatedGCN, shared weights achieve comparably strong results to non-shared weights at lower parameter count; weight sharing lowers performance for GCN and GIN.Source: §3.2 · pp. 4–5

3.3. The graph-rewiring perspective: νDRew as distance-aware skip-connections

  • Viewed as rewiring, DRew adds edges to the next distance shell at each layer, while delayed νDRew also moves messages across layers.This makes νDRew equivalent to geometric skip connections between different nodes, using graph distance to determine connections.Source: §3.3 · p. 5

4. Why νDRew Improves Information Processing

νDRew adds distance-aware skip connections that let distant nodes interact through fewer message-passing steps, mitigating over-squashing. Its delay controls how much surrounding graph structure is mixed into messages, and its shortest-path-based aggregations are at least more expressive than classical MPNNs.

Mitigating over-squashing

  • Distance-aware skips let nodes r hops apart interact through products with fewer than r message-passing factors, including a direct Γr interaction that mitigates over-squashing.Unlike powers of the same adjacency matrix, Γr matrices are described as avoiding the exponential decay associated with distance.Source: §4 · p. 5

Interpreting delay as local smoothing

  • Larger delay lets nodes aggregate features before substantial smoothing, while smaller delay incorporates structure encoded around the source in earlier layers.This gives delay a task-adaptive role in controlling how quickly topological information is mixed across graph regions.Source: §4 · pp. 5–6

Expressivity

  • Shortest-path multi-hop aggregation makes (ν)DRew at least more expressive than 1-WL and classical MPNNs, though finer analysis remains future work.The stated lower bound follows because the aggregations distinguish graph pairs distinguished by the shortest path kernel.Source: §4 · p. 6

5. Empirical Analysis

The experiments compare DRew variants with benchmark models on long-range tasks, then test robustness as network depth increases on synthetic and molecular tasks.

5.1. Long-range graph benchmark

  • Across all four LRGB tasks, νDRew-MPNNs outperform classical MPNNs and static-rewiring and multi-hop benchmarks; at least one matches or beats the best classical graph Transformer on each task.GraphGPS outperforms the best DRew model on PascalVOC-SP and Peptides-struct, while DRew surpasses it on Peptides-func and PCQM-Contact.Source: Table 2 · p. 7, §5.1 · p. 7

5.2. QM9

  • 12 out of 13 regression targets have DRew matching or beating SPN, and DRew without delay achieves the best average performance across targets.The authors note that slow-flow models cannot guarantee direct interaction between every node pair when graph diameter exceeds network depth.Source: §5.2 · pp. 7–8, Table 3 · p. 8

5.3. Validating robustness

  • νDRew robustness at increasing depth is assessed on a synthetic task and a real-world molecular dataset, rather than through raw performance alone.Source: §5.3 · p. 8

5.3.1. RINGTRANSFER

  • Full-delay ν1DRew-GCN maintains perfect accuracy at 30 or more layers, while classical GCN fails before 10 layers and SP-GCN fails at shallower depth than νDRew.The authors attribute the full-delay result to direct interaction between the target and delayed source node.Source: §5.3.1 · p. 8, Figure 2 · p. 8

5.3.2. LAYERWISE PERFORMANCE ON PEPTIDES-FUNC

  • On Peptides-func, greater delay supports better performance at large depth, while non-delay DRew often performs better at low depth, especially with positional encoding.ν1DRew consistently improves with more layers on this long-range task.Source: §5.3.2 · pp. 8–9, Figure 3 · p. 9

6. Related Work

The section groups graph-rewiring approaches by how they alter connectivity, distinguishing local multi-hop aggregation from global connectivity.

6. Related Work

  • Graph rewiring methods may remove elements judged irrelevant or add elements to increase connectivity.Source: §6 · p. 9
  • Added connectivity can take forms such as a global node, a global layer, positional or structural encodings, or a fully adjacent layer after message passing.Source: §6 · p. 9
  • Local rewiring includes multi-hop MPNNs, whereas global rewiring includes graph Transformers that fully connect the input graph.Source: §6 · p. 9

7. Conclusion and Future Work

The conclusion presents layer-dependent rewiring and distance-aware delayed skip connections as the framework's core contributions. It also limits the expected advantage to long-range or deep-GNN settings and outlines computational considerations.

7. Conclusion and Future Work

  • DRew adapts layer-dependent dynamic rewiring to MPNNs, and its delay mechanism allows local skip connections between nodes based on their mutual distance.Comparing expressive power across static and dynamic rewiring, and studying distance-aware skip connections, is identified as future work.Source: §7 · p. 9
  • The authors expect DRew to help with long-range interactions or very deep GNNs, but not necessarily shallow homophilic node classification; PascalVOC-SP's long-range dependencies are uncertain.They suggest this may partly explain why GraphGPS outperforms DRew on PascalVOC-SP.Source: §7 · p. 9

A. Time and Space Complexity

  • Shortest-path information up to k hops can be precomputed in worst-case O(|V||E|) time and reused across runs.Source: A. Time and Space Complexity · p. 13
  • When layer depth exceeds graph diameter, DRew can aggregate over O(V^2) node pairs, but parallel neighborhood aggregation and gradual densification are described as making it faster than comparable multi-hop MPNNs or Transformers.Source: A. Time and Space Complexity · p. 13
  • Delayed representations require storing ℓ node-feature representations at layer ℓ, scaling linearly with network depth; the authors say this has not been a practical bottleneck.Batch size could be reduced if needed.Source: A. Time and Space Complexity · p. 13

B. Further Experimental Details

The appendix provides further experimental details and Peptides-func results, including the hardware used.

B. Further Experimental Details

B.1. Hardware

  • Experiments ran on server nodes using a single GPU.Source: §B.1 · p. 13
  • The hardware included P100, V100, A100, Titan V, and RTX GPUs, alongside Broadwell, Haswell, and Cascade Lake CPUs.Source: §B.1 · p. 13

B.2. Long range graph benchmark

The Long Range Graph Benchmark experiments use matched parameter budgets and data splits, with additional settings specified for individual methods and datasets.

B.2. Long range graph benchmark

  • Tables 1 and 2 use the same data splits as classical MPNN counterparts and a roughly 500,000-parameter budget, adjusted through hidden dimensions.Table 4 lists parameter counts and significant hyperparameters for DIGL, MixHop-GCN, and (ν)DRew-MPNN experiments.Source: §B.2 · p. 13, Table 4 · p. 14
  • Experiments train for up to 300 epochs and average three seeds, using AdamW, batch normalization, and task-specific batch sizes.The listed batch sizes are 128 for Peptides, 256 for PCQM-Contact, and 32 for PascalVOC-SP.Source: §B.2 · p. 13
  • Peptides and PCQM-Contact use the Atom node encoder, while PascalVOC-SP uses a linear node encoder with output dimension 14.Table 2 experiments use no edge encoding or edge features; PascalVOC-SP superpixels use SLIC compactness 30.Source: §B.2 · p. 13
  • Laplacian positional-encoding experiments use hidden dimension 16 and two layers, while DIGL and MixHop-GCN use method-specific settings.DIGL uses PPR diffusion with α = 0.2; MixHop-GCN hyperparameters were selected by searching over maximum power and depth.Source: §B.2 · pp. 13–14

B.3. QM9

The QM9 performance experiments use a fixed parameter budget and a setup adapted from prior work, while omitting edge features in SPN and DRew-GIN.

B.3. QM9

  • Performance experiments use a fixed 800k parameter budget, tuned hidden dimension across network depths L, learning rate 0.001, mean pooling, no dropout, batch normalization, and averages over three runs.They train for 300 epochs and use data splits from (Brockschmidt, 2020).Source: §B.3 · p. 14
  • SPN and DRew-GIN omit edge features, although many comparison benchmarks are relational MPNNs that use separate weights for each edge type.The authors expect R-GIN's improvement to be slight because over 92% of QM9 graph edges have only one type.Source: §B.3 · p. 14

B.4. RingTransfer

RingTransfer experiments compare DRew-GCN with GCN and SP-GCN under matched parameter counts across ring lengths.

B.4. RingTransfer

  • DRew-GCN uses a varying hidden dimension to match GCN and SP-GCN parameter counts for each ring length and corresponding depth.The synthetic dataset contains 2,000 examples and five classes; models train for 50 epochs and average three experiments.Source: §B.4 · p. 15
  • SP-GCN is defined as an instance of the SP-MPNN framework, with maximum hop distance kmax and learned weight vectors for hop aggregation.The passage identifies Pkmax_ij as defined in Eq. (8).Source: §B.4 · p. 15

B.5. Ablation on Peptides-func

The Peptides-func ablation varies depth and delay, with and without Laplacian positional encoding. Results show different depth preferences across delay settings, while positional encoding further accelerates information flow.

B.5. Ablation on Peptides-func

  • Low-depth, low/no-delay and high-depth, high-delay νDRew-GCN outperform GCNs and the best-performing Transformer, with or without positional encoding.Figures 4 and 5 show Peptides-func results without and with Laplacian positional encoding, respectively.Source: §B.5 · p. 15, Figure 4 · p. 16, Figure 5 · p. 16
  • At high depth, ν1DRew continues improving while GCN and ResGCN degrade; low-depth ν1DRew remains comparable to classical MPNNs on this long-range task.ν1DRew allows nodes r edges apart to interact from the rth layer, so shallow models may lack enough layers for long-range interactions.Source: §B.5 · p. 15
  • Low-delay νDRew performs strongly at low depth but worsens at great depth, where faster information flow can make the computational graph too dense.The authors describe delay as a way to tune information flow to the task.Source: §B.5 · p. 15
  • Adding positional encoding exacerbates the depth-and-delay behaviors and accelerates information flow by preceding message passing with a global layer.Source: §B.5 · p. 15
  • The authors posit that positional encoding, rather than Transformer architecture, primarily explains Transformers' strength on Peptides; vanilla GCN with LapPE outperforms SAN at five layers.νDRew outperforms MPNN and Transformer-with-PE benchmarks with or without using positional encoding itself.Source: §B.5 · p. 15
Loading 2305.08018v2…