Asynchronous Graph Neural Networks
Extending GNN expressivity beyond the classical 1-WL limit using Asynchronous Message Passing.
During the summer of 2022, I conducted research at the Vikas Lab at Aalto University (Espoo, Finland) under supervision of Prof. Vikas K. Garg. Our work focused on exploring the theoretical foundations of Graph Neural Networks (GNNs), specifically investigating how altering the timing of message-passing mechanism could fundamentally increase a network’s expressivity.
The Problem: Limitations of Standard GNNs
Graph Neural Networks traditionally rely on synchronous Message Passing (MP), where all nodes in a graph update their hidden states simultaneously based on their neighbors’ messages. While effective, this approach has strict theoretical limitations:
- The 1-WL Limit: The expressiveness of a standard GNN—its ability to map different graph structures to different embeddings—is strictly limited to the 1-Weisfeiler-Lehman test for identifying graph isomorphisms. Even with an infinite number of layers, there are distinct, non-isomorphic graph structures that standard GNNs will map to the exact same embedding.
- Information Bottlenecks: Finding the optimal number of layers is a constant trade-off between underreaching (not gathering enough neighborhood context) and over-smoothing/over-squashing (aggregating too much information into a single node, causing distinct nodes to become indistinguishable).
The Solution: Asynchronous Message Passing (AMP)
To create more powerful GNNs, we explored Asynchronous Message Passing (AMP). Instead of updating all nodes simultaneously, AMP updates nodes sequentially based on a calculated order.
This subtle shift in dynamics yields significant theoretical advantages:
- Breaking Symmetry: AMP can successfully distinguish nodes that have similar neighborhood states by leveraging information trees based on the precise time of updating.
- Beyond 1-WL Expressivity: We gathered evidence demonstrating that AMP can embed non-isomorphic graphs into different embeddings—a feat synchronous message passing cannot achieve, effectively pushing the model’s expressivity beyond the 1-WL limit.
- Easier Information Flow: In an asynchronous setup, nodes can reach information from further away in the graph much more easily, resulting in a more relaxed state of over-smoothing and over-squashing.