A Flexible, Equivariant Framework for Subgraph GNNs via Graph Products and Graph Coarsening

TL;DR

This paper presents a flexible and scalable Subgraph GNN framework that leverages graph products and coarsening to improve expressivity and efficiency, outperforming previous methods across multiple benchmarks.

Contribution

It introduces a novel subgraph association via graph products, enabling arbitrary subgraph selection and revealing new symmetries, with theoretical analysis and practical improvements.

Findings

Outperforms baseline methods on multiple benchmarks

Handles any number of subgraphs seamlessly

Demonstrates improved expressive power and scalability

Abstract

Subgraph GNNs enhance message-passing GNNs expressivity by representing graphs as sets of subgraphs, demonstrating impressive performance across various tasks. However, their scalability is hindered by the need to process large numbers of subgraphs. While previous approaches attempted to generate smaller subsets of subgraphs through random or learnable sampling, these methods often yielded suboptimal selections or were limited to small subset sizes, ultimately compromising their effectiveness. This paper introduces a new Subgraph GNN framework to address these issues. Our approach diverges from most previous methods by associating subgraphs with node clusters rather than with individual nodes. We show that the resulting collection of subgraphs can be viewed as the product of coarsened and original graphs, unveiling a new connectivity structure on which we perform generalized message passing. Crucially, controlling the coarsening function enables meaningful selection of any number of subgraphs. In addition, we reveal novel permutation symmetries in the resulting node feature tensor, characterize associated linear equivariant layers, and integrate them into our Subgraph GNN. We also introduce novel node marking strategies and provide a theoretical analysis of their expressive power and other key aspects of our approach. Extensive experiments on multiple graph learning benchmarks demonstrate that our method is significantly more flexible than previous approaches, as it can seamlessly handle any number of subgraphs, while consistently outperforming baseline approaches. Our code is available at https://github.com/BarSGuy/Efficient-Subgraph-GNNs.