Table of Contents
Fetching ...

Flood and Echo Net: Algorithmically Aligned GNNs that Generalize

Joël Mathys, Florian Grötschla, Kalyan Varma Nadimpalli, Roger Wattenhofer

TL;DR

Flood and Echo Net introduces an algorithmically aligned GNN framework that executes graph computation via a flooding phase followed by an echo phase, initiated from a single origin node. By propagating information in a wave-like manner and activating only nodes at each distance, it achieves $O(m)$ messages per phase and can generalize across graph sizes, offering expressive power beyond the 1-WL limit and potential efficiency gains. Theoretical results establish expressiveness, compare message complexity to traditional MPNNs, and demonstrate that a single phase can solve certain tasks with $O(m)$ messages. Empirically, the method improves generalization on simple algorithmic tasks and on the SALSA-CLRS benchmark, with the mode and number of phases providing flexibility for different algorithmic trajectories. This work suggests that aligning GNN execution with distributed algorithm patterns can yield practical benefits for algorithm learning and size generalization in graphs.

Abstract

Most Graph Neural Networks follow the standard message-passing framework where, in each step, all nodes simultaneously communicate with each other. We want to challenge this paradigm by aligning the computation more closely to the execution of distributed algorithms and propose the Flood and Echo Net. A single round of a Flood and Echo Net consists of an origin node and a flooding phase followed by an echo phase. First, during the flooding, messages are sent from the origin and propagated outwards throughout the entire graph. Then, during the echo, the message flow reverses and messages are sent back towards the origin. As nodes are only sparsely activated upon receiving a message, this leads to a wave-like activation pattern that traverses the graph. Through these sparse but parallel activations, the Net becomes more expressive than traditional MPNNs which are limited by the 1-WL test and also is provably more efficient in terms of message complexity. Moreover, the mechanism's inherent ability to generalize across graphs of varying sizes positions it as a practical architecture for the task of algorithmic learning. We test the Flood and Echo Net on a variety of synthetic tasks and the SALSA-CLRS benchmark and find that the algorithmic alignment of the execution improves generalization to larger graph sizes.

Flood and Echo Net: Algorithmically Aligned GNNs that Generalize

TL;DR

Flood and Echo Net introduces an algorithmically aligned GNN framework that executes graph computation via a flooding phase followed by an echo phase, initiated from a single origin node. By propagating information in a wave-like manner and activating only nodes at each distance, it achieves messages per phase and can generalize across graph sizes, offering expressive power beyond the 1-WL limit and potential efficiency gains. Theoretical results establish expressiveness, compare message complexity to traditional MPNNs, and demonstrate that a single phase can solve certain tasks with messages. Empirically, the method improves generalization on simple algorithmic tasks and on the SALSA-CLRS benchmark, with the mode and number of phases providing flexibility for different algorithmic trajectories. This work suggests that aligning GNN execution with distributed algorithm patterns can yield practical benefits for algorithm learning and size generalization in graphs.

Abstract

Most Graph Neural Networks follow the standard message-passing framework where, in each step, all nodes simultaneously communicate with each other. We want to challenge this paradigm by aligning the computation more closely to the execution of distributed algorithms and propose the Flood and Echo Net. A single round of a Flood and Echo Net consists of an origin node and a flooding phase followed by an echo phase. First, during the flooding, messages are sent from the origin and propagated outwards throughout the entire graph. Then, during the echo, the message flow reverses and messages are sent back towards the origin. As nodes are only sparsely activated upon receiving a message, this leads to a wave-like activation pattern that traverses the graph. Through these sparse but parallel activations, the Net becomes more expressive than traditional MPNNs which are limited by the 1-WL test and also is provably more efficient in terms of message complexity. Moreover, the mechanism's inherent ability to generalize across graphs of varying sizes positions it as a practical architecture for the task of algorithmic learning. We test the Flood and Echo Net on a variety of synthetic tasks and the SALSA-CLRS benchmark and find that the algorithmic alignment of the execution improves generalization to larger graph sizes.
Paper Structure (35 sections, 8 theorems, 8 equations, 15 figures, 14 tables, 1 algorithm)

This paper contains 35 sections, 8 theorems, 8 equations, 15 figures, 14 tables, 1 algorithm.

Key Result

Theorem 4.1

On connected graphs, the Flood and Echo Net is at least as expressive as any MPNN . Furthermore, it exchanges at most as many messages.

Figures (15)

  • Figure 1: The Flood and Echo Net propagates messages in a wave-like pattern throughout the entire graph. Starting from an origin (orange), messages are sent toward the origin's neighbors and then continuously sent or "flooded" farther away outwards (blue). Afterward, the flow reverses, and messages are "echoed" back (red) toward the origin. Throughout the computation, only a small subset of nodes is active at any given time, passing messages efficiently throughout the entire graph. Moreover, the mechanism naturally generalizes to graphs of larger sizes.
  • Figure 2: Flood and Echo Net
  • Figure 3: Example of two $4$-regular graphs which cannot be distinguished using standard MPNNs as they are 1-WL equivalent. However, no matter which starting point is chosen, the Flood and Echo Net can easily distinguish between them through the derived distance to the starting node.
  • Figure 4: Extrapolation on the PrefixSum task. All models are trained with graphs of size 10 and then tested on larger graphs. The Flood and Echo models are able to generalize well to graphs $100$ times the sizes encountered during training. We report both the node accuracy on the left and the graph accuracy on the right.
  • Figure 5: Graph Accuracy on the SALSA-CLRS benchmark for the BFS task on Erdős--Rényi graphs. The Flood and Echo Net can generalize almost perfectly to graphs 10 times larger while the baselines already have a significant drop in performance. Moreover, it does not rely on the number of steps given by the hints and executes a single phase.
  • ...and 10 more figures

Theorems & Definitions (17)

  • Theorem 4.1
  • Theorem 4.2
  • Lemma 4.3
  • Corollary 5.1
  • Corollary 5.2
  • Definition B.1
  • Lemma D.1
  • Corollary D.2
  • proof
  • proof
  • ...and 7 more