Flood and Echo Net: Algorithmically Aligned GNNs that Generalize

TL;DR

Flood and Echo Net introduces a novel GNN architecture inspired by distributed algorithms, utilizing flooding and echo phases to improve expressiveness and generalization across different graph sizes.

Contribution

It proposes a new GNN framework aligned with distributed algorithms, enhancing expressiveness and scalability beyond traditional message-passing neural networks.

Findings

More expressive than traditional MPNNs

Proven to be more efficient in message complexity

Improves generalization to larger 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.