Source-linked AI summary
From Stars to Subgraphs: Uplifting Any GNN with Local Structure Awareness
Lingxiao Zhao, Wei Jin, Leman Akoglu, Neil Shah
TL;DR
MPNNs offer scalable local aggregation but are limited by 1-WL expressiveness, motivating methods that improve structural discrimination without the costs of higher-order designs. GNN-AK wraps a base GNN as a kernel over locally induced subgraphs, with GNN-AK+ and SubgraphDrop extending the framework. The method achieves stronger theoretical expressiveness, practical scalability, and new state-of-the-art results on several benchmarks.
Problem
MPNNs are scalable but 1-WL-bounded, whereas more expressive k-order GNNs can be impractical for large graphs and may face generalization concerns.
Method
GNN-AK uplifts any GNN by applying it as a kernel to locally induced rooted subgraphs, with GNN-AK+ adding richer subgraph and context encodings and SubgraphDrop reducing overhead.
Results
GNN-AK is strictly more powerful than 1&2-WL and not less powerful than 3-WL, while GNN-AK+ achieves new state-of-the-art performance on several graph-level benchmarks.
Takeaways & Limitations
The framework provides a middle ground between scalable but less-expressive MPNNs and highly expressive yet practically infeasible higher-order designs.
Takeaways & Limitations
For limited layer depth, rooted local subgraphs may fail to capture global differences between some non-isomorphic graphs, including members of the CFI(k) family.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
MPNNs are efficient and widely used, but their expressiveness is bounded by 1-WL and cannot capture some basic structural concepts. GNN-AK addresses this gap by replacing star-shaped aggregation with GNN encoding of local induced subgraphs while retaining practical scalability.
- MPNNs recursively aggregate immediate-neighbor information, making them intuitive, effective, and efficient but limiting their expressiveness to 1-WL.
- Prior attempts to reach higher Weisfeiler-Leman expressiveness often require O(k)-order tensors, creating scalability and feasibility problems for large graphs.
- GNN-AK+ extracts rooted subgraphs, convolves them with a base GNN, and combines centroid, subgraph, and context encodings into node representations.
- GNN-AK uplifts any GNN by applying it as a kernel to locally induced subgraphs instead of the whole input graph, generalizing star-pattern aggregation.
- The framework reports theoretical expressiveness beyond 1-WL and practical gains across graph-learning benchmarks, with implementations accepting GNNs from PyG.
2 RELATED WORK
Prior work incorporates subgraph or higher-order structure through Weisfeiler-Leman tests, motifs, handcrafted features, adjacency powers, and meta-learning approaches.
- Existing methods use k-WL, motif information, handcrafted subgraph features, or adjacency-matrix powers to incorporate structural information into GNNs.
3 GENERAL FRAMEWORK AND THEORY
The framework replaces MPNNs’ star-shaped local aggregation with induced-subgraph encoding, using a GNN as a wrapper around a base GNN. Its theory establishes stronger expressiveness than 1&2-WL and connections to 3-WL, while highlighting conditions and limits of local subgraphs.
- From stars to subgraphs: GNN-AK generalizes the star around each node to an induced k-hop egonet or other local subgraph, capturing connections among neighbors.This motivates Subgraph-1-WL, which replaces 1-WL color refinement with an injective hash of the rooted induced subgraph.
- MPNN limitations: MPNNs repeatedly aggregate self and neighbor representations, producing graph embeddings through global pooling, but their expressiveness is tied to 1-WL.Their update and aggregation functions, together with pooling, determine practical expressiveness; injective choices achieve the 1-WL maximum.
- Neural framework: GNN-AK transfers the subgraph framework to neural networks by using any GNN, mainly an MPNN, as the subgraph encoder and wrapper around a base GNN.The framework computes representations from attributed induced subgraphs across layers and pools node embeddings into a graph representation.
- Expressiveness correspondence: Under sufficient layers and injective functions, MPNN-AK matches Subgraph-1-WL*, while the general GNN-AK matches Subgraph-1-WL when its base GNN matches the hash function.These results characterize neural expressiveness through the corresponding subgraph refinement procedures.
- Theoretical guarantees: Subgraph-1-WL* is strictly more powerful than 1&2-WL, and when the hash is 3-WL expressive, Subgraph-1-WL is no less powerful than 3-WL.Consequently, 1-WL-expressive MPNN-AK is strictly stronger than 1&2-WL, while PPGN-AK can distinguish some 3-WL-failed non-isomorphic graphs.
- Conditions and limits: Successful discrimination requires sufficiently informative rooted subgraphs and a discriminative hash, but limited-radius egonets can miss global differences between non-isomorphic graphs.Theorem 4 gives such a failure for t ≤4 on CFI(k) graph pairs; larger subgraphs may also require a more expressive base model.
4 CONCRETE REALIZATION
GNN-AK realizes local subgraph-aware aggregation by applying a base GNN to rooted induced subgraphs, while GNN-AK+ adds context and distance-to-centroid encodings. The framework also supports random-walk subgraphs and incurs overhead tied to rooted-subgraph size.
- GNN-AK: GNN-AK applies a base GNN to rooted induced subgraphs, producing subgraph and centroid encodings that are fused into node representations.The rooted subgraph is pooled with SUM or MEAN, optionally concatenating the root representation to distinguish its centroid.
- GNN-AK+: GNN-AK+ adds context encoding to capture each node from different subgraph contexts or points-of-view.This supplements information from intermediate embeddings that the basic realization does not fully exploit.
- GNN-AK+: GNN-AK+ uses distance-to-centroid features to augment hidden representations and gate subgraph and context encodings before pooling.The gating design weights nodes according to their distance from the rooted node, with element-wise multiplication used in the formal mechanism.
- Expressiveness: GNN-AK+ is at least as powerful as GNN-AK under the proposed realization.The paper states this relationship as Proposition 2.
- Subgraph extraction: Random-walk extraction limits subgraph size when k-hop egonets become too large in dense or small-diameter graphs.Multiple truncated walks are combined, and online re-extraction during training further reduces the effect of sampling randomness through regularization.
- Complexity: O(s) is the runtime and memory overhead factor induced by rooted subgraphs of size s over the base GNN.Subgraph extraction itself takes O(k|E|) time, while the union graph determines the wrapped model’s computational cost.
5 IMPROVING SCALABILITY: SUBGRAPHDROP
SubgraphDrop reduces GNN-AK(+) overhead by selecting a small, covering set of rooted subgraphs during training while retaining all subgraphs for evaluation. Its alignment procedures address incomplete or differently scaled encodings caused by sampling.
- SubgraphDrop: SubgraphDrop drops only a subset of rooted subgraphs during training, while evaluation uses all subgraphs to preserve information.Randomized sampling varies selected subgraphs across epochs, providing an amortized approximation to full-mode training.
- Trade-off: SubgraphDrop significantly reduces memory overhead while keeping performance nearly the same as training with all subgraphs.The method is designed to align training and evaluation encodings despite subsampling.
- Subgraph sampling strategies: R times the overhead of the base model is targeted by sampling subgraphs that cover each node approximately R times, with R≤3 in practice.The strategies aim to select minimally redundant subgraphs while evenly covering the graph.
- Encoding alignment: Subgraph and centroid encodings for unselected roots are estimated by propagating encodings from selected roots through distance-based node partitions.Nodes are grouped by their shortest-path distance to the selected-root set and processed iteratively.
- Encoding alignment: Context encodings are rescaled for SUM pooling to match full-mode evaluation, while MEAN pooling requires no modification.Every node remains covered at least R times during training, enabling context encoding throughout the graph.
6 EXPERIMENTS
Experiments evaluate expressiveness on simulation tasks, practical performance across real-world datasets, sampling efficiency, and architectural ablations. GNN-AK(+) improves multiple base GNNs and reaches new state-of-the-art results on several benchmarks, while sampling addresses resource pressure.
- Experimental design: GNN-AK(+) is evaluated on four simulation datasets to test expressiveness and on five real-world datasets to test practical performance.The experiments also assess SubgraphDrop and ablate structural components and model depth.
- Empirical verification of expressiveness: GNN-AK(+) boosts base GNN performance across simulation tasks, with Table 1 reporting accuracy, mean absolute error, and out-of-memory outcomes.The table is presented as empirical verification of the framework’s expressiveness lift.
- Comparing with SOTA and generality: GNN-AK+ achieves new state-of-the-art performance on ZINC-12K, CIFAR10, and PATTERN.PPGN is omitted from the real-world comparison because of scalability, while lifts are observed across datasets and base GNNs.
- Comparing with SOTA and generality: 60.3%, 50.5%, and 39.4% are the relative error reductions on ZINC, CIFAR10, and PATTERN for GCN, GIN, and PNA∗ base models, respectively.These reductions correspond to the benchmark comparisons reported for GNN-AK+.
- Experimental setup: PNA∗ changes degree-based scaling to degree encoding concatenated with node embeddings, eliminating the need to compute dataset average degree.This variant is used among the experimental base models.
- Scaling up by subsampling: GNN-AK(+)-S can avoid out-of-memory failures and can match or improve the performance of its non-sampled counterpart.The resource analysis studies GIN-AK+-S with varying R on ZINC-12K and CIFAR10.
7 CONCLUSION
GNN-AK is a general-purpose wrapper that applies a base GNN to induced subgraphs, improving expressiveness and benchmark performance while retaining practical scalability.
- GNN-AK uses a base GNN as a kernel on induced subgraphs, generalizing classical MPNN star-pattern aggregation.The framework is designed to uplift the expressiveness of different base GNN choices.
- GNN-AK(+) improves expressiveness across multiple base GNN choices, including beyond 1&2-WL for MPNNs and beyond 3-WL for PPGN.
- The expressiveness gains translate into performance improvements and state-of-the-art results on graph-level benchmarks.
- SubgraphDrop reduces runtime and memory footprint while retaining scalability to practical graphs.The method introduces SubgraphDrop as a subsampling-based realization of the framework.
A.1 ADDITIONAL RELATED WORK
Prior expressive GNN approaches improve graph discrimination through permutation-sensitive operations, node indicators, or convolution-like constructions, but introduce feasibility, invariance, or complexity trade-offs.
- Permutation-sensitive functions can achieve universality, but summing across a combinatorial number of permutations limits feasibility.
- Node indicators can improve distinguishability, while preserving invariance may require quadratic time complexity.
- MoNet and PatchySAN pursue CNN-like graph convolutions using pseudo-coordinates or learned neighbor ordering and truncation.
A.2 SAMPLING BY SUBGRAPHDROP
SubgraphDrop accelerates GNN-AK-S training by selecting a subset of rooted subgraphs that repeatedly covers every node, then propagating estimated encodings to uncovered nodes.
- GNN-AK-S extracts subgraphs and subsamples m≪|V| subgraphs so every node is covered at least R times.Multiple sampling strategies are supported.
- During training, the base GNN computes intermediate node embeddings within the selected subgraphs.
- Subgraph and centroid encodings initially exist only for selected roots and are propagated to estimate encodings for other nodes.Context encodings are scaled to match evaluation.
- Random, farthest, and min-set-cover sampling provide alternative procedures for selecting subgraphs until the coverage requirement is met.Farthest sampling uses shortest-path distance, while min-set-cover sampling follows a greedy minimum set-cover procedure.
A.3 PROOF OF THEOREM 1
The proof establishes that, with injective components and sufficient depth, MPNN-AK distinguishes graphs according to Subgraph-1-WL*.
- With sufficient layers and injective functions, an MPNN is as powerful as 1-WL.
- MPNN-AK outputs different vectors exactly when Subgraph-1-WL* produces different fingerprints, assuming an injective pooling operation.
- Therefore, MPNN-AK is as powerful as Subgraph-1-WL*.
A.4 PROOF OF THEOREM 2
The appendix proves that subgraph-based refinement can distinguish graph pairs that 1-WL, 2-WL, or 3-WL cannot, while also documenting construction details, limitations, and ablation findings.
- Proof of Theorem 2: Subgraph-1-WL* is strictly more powerful than 1-WL, and therefore also exceeds 2-WL in expressiveness.The proof combines an upper-bound argument with 4-regular graphs distinguishable by 1-hop egonets despite being indistinguishable by 1-WL.
- Proof of Theorem 3: 3-WL-failed strongly regular graphs can be distinguished by Subgraph-1-WL when its hash function is 3-WL-discriminative.Their rooted 1-hop egonets differ in triangle counts: one contains 8 triangles and the other 6.
- Proof of Theorem 4: The CFI construction replaces each degree-d base-graph vertex with X_d, whose endpoints encode incident edges and whose middle section has size 2^d−1.X_d has exactly 2^d−1 automorphisms, determined by swapping endpoint pairs for even-cardinality subsets.
- Proof of Theorem 4: The theorem does not imply that Subgraph-1-WL with t ≥5-egonets always distinguishes the constructed graph pair; this depends on the base graph.A stated sufficient scenario involves the left and right parts of the rooted subgraph being connected only through a specific separator structure.
- Ablation Study: D2C makes GNN-AK+ less sensitive to base-model depth, while models without D2C face a trade-off between base expressiveness and outer depth.With D2C, one-layer base models suffice for strong substructure-counting performance and achieve the best graph-property regression performance; without D2C, base depth remains important.
- Additional Results on TU Datasets: GIN-AK+ does not improve IMDB-B because every graph is itself an egonet, while MUTAG and PTC results are unstable on very small datasets.The datasets contain 188 and 344 graphs, respectively, and evaluation averages validation curves across folds.