Efficient Subgraph GNNs by Learning Effective Selection Policies

Subgraph GNNs are provably expressive neural architectures that learn graph representations from sets of subgraphs. Unfortunately, their applicability is hampered by the computational complexity associated with performing message passing on many subgraphs. In this paper, we consider the problem of learning to select a small subset of the large set of possible subgraphs in a data-driven fashion. We first motivate the problem by proving that there are families of WL-indistinguishable graphs for which there exist efficient subgraph selection policies: small subsets of subgraphs that can already identify all the graphs within the family. We prove that, unlike popular random policies and prior work addressing the same problem, our architecture is able to learn the efficient policies mentioned above. In essence, a Subgraph GNN first transforms an input graph into a bag of subgraphs, obtained according to a predefined generation policy. For instance, each subgraph might be generated by deleting exactly one node in the original graph, or, more generally, by marking exactly one node in the original graph, while leaving the connectivity unaltered (Papp & Wattenhofer, 2022). Then, it applies an equivariant architecture to process the bag of subgraphs, and aggregates the representations to obtain graphor node-level predictions. The popularity of Subgraph GNNs can be attributed not only to their increased expressive power compared to MPNNs but also to their remarkable empirical performance, exemplified by their success on the ZINC molecular dataset (Frasca et al., 2022; Zhang et al., 2023). Unfortunately, Subgraph GNNs are hampered by their computational cost, since they perform message-passing operations on all subgraphs within the bag.