From Stars to Subgraphs: Uplifting Any GNN with Local Structure Awareness
arXiv:2110.03753
Abstract
Message Passing Neural Networks (MPNNs) are a common type of Graph Neural Network (GNN), in which each node's representation is computed recursively by aggregating representations (messages) from its immediate neighbors akin to a star-shaped pattern. MPNNs are appealing for being efficient and scalable, how-ever their expressiveness is upper-bounded by the 1st-order Weisfeiler-Lehman isomorphism test (1-WL). In response, prior works propose highly expressive models at the cost of scalability and sometimes generalization performance. Our work stands between these two regimes: we introduce a general framework to uplift any MPNN to be more expressive, with limited scalability overhead and greatly improved practical performance. We achieve this by extending local aggregation in MPNNs from star patterns to general subgraph patterns (e.g.,k-egonets):in our framework, each node representation is computed as the encoding of a surrounding induced subgraph rather than encoding of immediate neighbors only (i.e. a star). We choose the subgraph encoder to be a GNN (mainly MPNNs, considering scalability) to design a general framework that serves as a wrapper to up-lift any GNN. We call our proposed method GNN-AK(GNN As Kernel), as the framework resembles a convolutional neural network by replacing the kernel with GNNs. Theoretically, we show that our framework is strictly more powerful than 1&2-WL, and is not less powerful than 3-WL. We also design subgraph sampling strategies which greatly reduce memory footprint and improve speed while maintaining performance. Our method sets new state-of-the-art performance by large margins for several well-known graph ML tasks; specifically, 0.08 MAE on ZINC,74.79% and 86.887% accuracy on CIFAR10 and PATTERN respectively.
Expressive GNN framework
References in corpus (8)
- Fast Graph Representation Learning with PyTorch Geometric
- Simplifying Graph Convolutional Networks
- Deep learning for molecular design - a review of the state of the art
- The Weisfeiler-Lehman Method and Graph Isomorphism Testing
- Hierarchical Inter-Message Passing for Learning on Molecular Graphs
- Breaking the Limits of Message Passing Graph Neural Networks
- Ego-GNNs: Exploiting Ego Structures in Graph Neural Networks
- Counting Substructures with Higher-Order Graph Neural Networks: Possibility and Impossibility Results