PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding and Extending Subgraph GNNs by Rethinking Their Symmetries

Fabrizio Frasca, Beatrice Bevilacqua, Michael M. Bronstein, Haggai Maron

arXiv:2206.11140v3cs.LG

TL;DR

The paper addresses the incomplete understanding of node-based Subgraph GNNs, especially their expressive-power bound and permissible equivariant layers. It analyzes their aligned symmetries through invariant graph networks, proves a 3-WL upper bound, and introduces a general layer family culminating in SUN. The resulting framework unifies prior node-based architectures and supports improved empirical generalisation.

  • Problem

    The expressive-power upper bound and general aggregation and information-sharing rules for node-based Subgraph GNNs were not fully understood.

  • Method

    The paper uses node-subgraph alignment to derive a single symmetry group, connects Subgraph GNNs to invariant graph networks, and develops general equivariant layers with SUN.

  • Results

    Node-based Subgraph GNNs are bounded by 3-WL, while SUN theoretically generalises previous node-based architectures and can outperform them empirically.

  • Takeaways & Limitations

    The framework unifies and extends node-based Subgraph GNNs while supporting systematic study of models whose expressivity lies between 1-WL and 3-WL.

  • Takeaways & Limitations

    The proposed equivariant-layer design space is not exhausted, and extending the analysis to higher-order node-based policies remains future work.

Abstract

from arXiv · show

Subgraph GNNs are a recent class of expressive Graph Neural Networks (GNNs) which model graphs as collections of subgraphs. So far, the design space of possible Subgraph GNN architectures as well as their basic theoretical properties are still largely unexplored. In this paper, we study the most prominent form of subgraph methods, which employs node-based subgraph selection policies such as ego-networks or node marking and deletion. We address two central questions: (1) What is the upper-bound of the expressive power of these methods? and (2) What is the family of equivariant message passing layers on these sets of subgraphs?. Our first step in answering these questions is a novel symmetry analysis which shows that modelling the symmetries of node-based subgraph collections requires a significantly smaller symmetry group than the one adopted in previous works. This analysis is then used to establish a link between Subgraph GNNs and Invariant Graph Networks (IGNs). We answer the questions above by first bounding the expressive power of subgraph methods by 3-WL, and then proposing a general family of message-passing layers for subgraph methods that generalises all previous node-based Subgraph GNNs. Finally, we design a novel Subgraph GNN dubbed SUN, which theoretically unifies previous architectures while providing better empirical performance on multiple benchmarks.

1 Introduction

Subgraph GNNs apply GNNs to bags of extracted subgraphs, but their expressive-power ceiling and permissible information-sharing rules remain insufficiently understood. This paper analyzes node-based methods through a smaller symmetry group, establishes a 3-WL upper bound, and proposes a general layer family culminating in SUN.

  • Subgraph GNNs apply MPNNs to collections of subgraphs and aggregate their representations, extending beyond standard message-passing GNNs.
  • Their expressive-power upper bound and the space of valid aggregation and information-sharing rules were previously unknown.
  • Node-based policies align each subgraph with a unique node, allowing one joint permutation group to replace independently permuted node and subgraph groups.
  • The paper connects node-based Subgraph GNNs to invariant graph networks and proves that their expressive power is bounded by 3-WL.
  • A general layer formulation captures local and global operations, recovers previous methods, and yields SUN, which theoretically generalises earlier node-based Subgraph GNNs and can outperform them empirically.

2 Previous and related work

Prior work established Subgraph GNNs as a way to increase GNN expressivity by processing subgraphs, while invariant graph networks provided a related tensor-based framework for equivariant graph computation. The paper builds on both lines of work to study expressive power and layer design.

  • Expressive power of GNNs: Expressive GNN research includes subgraph methods, positional and structural encodings, higher-order message passing, and equivariant models addressing limitations of standard MPNNs.
  • Invariant graph networks: Invariant graph networks represent graphs or hypergraphs with higher-order tensors and apply permutation-equivariant linear layers with pointwise nonlinearities.
  • Invariant graph networks: IGN equivariant layers can be expressed through operators associated with equality patterns over tensor indices.
  • Invariant graph networks: These layers also admit pooling-broadcasting interpretations between index orbits, such as broadcasting diagonal sums to off-diagonal entries.
  • Subgraph GNNs: Subgraph GNNs apply GNNs to bags of subgraphs, with prior models differing in whether subgraphs are processed independently or share information across aligned nodes.

3 Node-based Subgraph GNNs

Node-based Subgraph GNNs generate one subgraph per graph node and process the resulting bag with equivariant layers and invariant pooling. Common policies and architectures differ in how they modify subgraphs and share information, while their upper bound was previously unknown.

  • Notation: Graphs use adjacency and feature matrices, while a bag of subgraphs is represented by corresponding tensors indexed by subgraph, node, and feature dimensions.
  • Formalising Subgraph GNNs: Subgraph GNNs formalise selection, stacked node- and subgraph-permutation-equivariant layers, invariant pooling, and an MLP around a 1-WL-maximally expressive MPNN base encoder.
  • Node-based selection policies: A node-based selection policy produces the bag {f(G, v)} for graph nodes v, associating every subgraph with a unique node.
  • Node-based selection policies: Common node-based policies include node deletion, node marking, and ego-networks, each producing a subgraph associated with a root node.
  • Known node-based Subgraph GNNs: Known architectures include reconstruction, GNN-AK, NGNN, ID-GNN, DS-GNN, and DSS-GNN variants, which differ in how information is shared between subgraphs and nodes.

4 Symmetries of node-based subgraph selection policies

Node-based subgraph bags possess a node-subgraph alignment that invalidates independent node and subgraph permutations. Their correct symmetry is a single joint permutation group, matching the tensor symmetries of 3-IGNs and 2-IGNs.

  • Previous formulations used separate symmetric groups for node and subgraph permutations, acting independently on bag tensors.
  • Because node-based policies define a bijection between nodes and subgraphs, the bag can be ordered consistently and acquires a node-subgraph alignment.
  • Only permutations acting identically on nodes and subgraphs preserve this aligned structure, so the appropriate symmetry is one joint Sn action.
  • The joint group Sn is smaller than independent product groups, imposing fewer equivariance constraints and potentially permitting more expressive models.
  • The resulting tensor symmetries match third-order tensors in 3-IGNs for adjacency and second-order tensors in 2-IGNs for features, enabling transfers from IGN theory.

5 A representational bound for Subgraph GNNs

The paper bounds known node-based Subgraph GNNs by showing that 3-IGNs can implement their selection policies and layers. Since 3-IGNs match 3-WL in expressive power, 3-WL upper-bounds these Subgraph GNNs.

  • Layer implementation: 3-IGNs can implement the generalized message-passing and pooling layers of known node-based Subgraph GNNs.This establishes an IGN realization of each intermediate Subgraph GNN layer.
  • Selection-policy implementation: 3-IGNs implement every known node-based selection policy in class Π through suitable stacked layers.The result covers the policies used by previously proposed node-based Subgraph GNNs.
  • Expressive-power bound: For any distinguished non-isomorphic graph pair, a corresponding 3-IGN instance also distinguishes the pair.The theorem assumes Subgraph GNNs use Morris et al. message-passing base encoders.
  • Expressive-power bound: 3-WL upper-bounds node-based Subgraph GNNs: whenever such a model distinguishes two graphs, 3-WL distinguishes them as well.This follows because 3-IGNs have the same expressive power as 3-WL.

6 A design space for Subgraph GNNs

The section characterizes a tractable equivariant layer space for node-based Subgraph GNNs by extending 2-IGN operations with local and global aggregation. This framework recovers known methods, supports novel operations, and is bounded in expressive power by 3-WL under stated conditions.

  • Layer-space motivation: The layer-space problem asks whether existing Subgraph GNN update rules exhaust all equivariant options.A full 3-IGN basis is impractical because it is vast, can require O(n^3) space, and does not directly support local message passing.
  • Extended 2-IGNs: ReIGN(2) extends 2-IGNs for n subgraphs with n nodes by allowing both global and connectivity-based local pooling.Local aggregation can follow adjacency in either the selected subgraph or the original graph, while global pooling aggregates across relevant nodes or subgraphs.
  • Extended 2-IGNs: The extended operations recover sparse message passing, global pooling across subgraphs, and additional operations such as transposition.Root and non-root entries can use distinct update rules, reflected by diagonal and off-diagonal positions in the n × n representation.
  • ReIGN(2) framework: ReIGN(2) provides a tractable pool of atomic operations that can be selected and combined into equivariant layers, recovering previously proposed node-based Subgraph GNNs.Theorem 8 states that every model in the considered family with a message-passing base encoder has an equivalent ReIGN(2) instance.
  • Expressive-power bounds: Under policies in Π and 3-IGN-computable invariant pooling, ReIGN(2) is upper-bounded by 3-WL.This follows from the result that 3-IGNs implement ReIGN(2) models under those conditions.
  • Scope and open space: ReIGN(2) may not capture every S_n-equivariant layer over R^n2, and SUN omits some ReIGN(2) operations for future exploration.The section specifically notes that SUN does not use all possible operations in the ReIGN(2) framework.
  • SUN: SUN selects a tractable subset of ReIGN(2) operations, uses separate parameters for root and non-root nodes, and theoretically implements known Subgraph GNN layers.Proposition 12 states that a stacking of SUN layers can implement any layer of a known node-based Subgraph Network.

7 Experiments

Experiments evaluate SUN on synthetic expressiveness tasks, molecular benchmarks, and limited-data generalisation. SUN achieves strong benchmark performance and generally better generalisation than competing Subgraph GNNs.

  • Synthetic: SUN attains state-of-the-art results on 3 of 4 substructure-counting tasks.
  • Real-world: SUN performs best among domain-agnostic GNNs under the 500k parameter budget on ZINC-12k.
  • Real-world: On Molhiv, SUN either outperforms or approaches GNNs that explicitly model rings, including HIMP, GSN, and CIN.
  • Generalisation from limited data: SUN generalises better than other Subgraph GNNs on cycle-counting except during a short initial phase for the EGO policy.
  • Generalisation from limited data: On ZINC-12k, SUN surpasses DSS-GNN, DS-GNN, and GNN-AK variants from 20%, 30%, and 40% of samples, respectively.

8 Conclusions

The paper unifies, extends, and analyses node-based Subgraph GNNs through a 3-WL expressivity bound and a general layer family. SUN also shows improved generalisation, while the work notes possible pharmaceutical relevance and future extensions.

  • Node-based Subgraph GNNs have expressive power bounded by the 3-WL test.
  • The paper proposes a new family of layers for studying models with expressivity between 1-WL and 3-WL.
  • SUN shows considerable improvement in the generalisation abilities of Subgraph GNNs.
  • The authors identify possible positive applications in pharmaceutical research, including drug discovery, based on molecular property results.

Checklist

The checklist records affirmative answers for theoretical assumptions and proofs, experimental reproducibility, reporting, asset licensing, and discussion of limitations and societal impact.

  • The paper states that it accurately reflects its contributions and scope and discusses its limitations and potential negative societal impacts.
  • The authors state that they provide the full assumptions and complete proofs for theoretical results.
  • The paper states that code, data, experiment instructions, and training details are provided for reproducing the main experiments.
  • Standard deviations are reported across multiple seeds for ZINC12k, ogbg-molhiv, generalisation experiments, and counting-substructure experiments.
  • Existing assets are cited and licensed, while new assets and human-subject procedures are marked not applicable.

A.1 Review of Subgraph GNN architectures

The reviewed architectures apply GNNs to selected subgraph bags and differ in selection policies, within-subgraph processing, cross-subgraph aggregation, and global pooling. Their complexity depends on bag size, graph degree, and policy-specific subgraph size.

  • Reconstruction GNN / NGNN / DS-GNN: Reconstruction GNN applies a Siamese MPNN to node-deleted subgraphs and aggregates their representations with a set network.
  • ESAN: ESAN generalises Reconstruction GNNs with broader selection policies and equivariant layers that share information across subgraphs.
  • NGNN: NGNN applies an independent MPNN to each ego-network and then pools local representations before running an additional MPNN on the original graph.
  • ID-GNN: ID-GNN distinguishes messages propagated by ego-network roots through heterogeneous message passing on each subgraph.
  • Computational complexity: For node-based policies, generic subgraph methods have time complexity O(n^2 · d) and memory complexity O(n · (n + nd)).
  • Computational complexity: SUN uses local message passing and global pooling operations with the same asymptotic bounds as the corresponding analysis.
  • Computational complexity: For ego-network policies with maximum subgraph size c, time complexity becomes O(n · c · d), which can be smaller when c ≪ n.

B Proofs for Section 5 – Subgraph GNNs and 3-IGNs

The proofs represent node-based subgraph bags with tensor structures governed by shared permutation symmetries, connecting them to 3-IGNs. This framework establishes that such Subgraph GNNs are bounded by 3-WL.

  • 3-IGN data structure: 3-IGNs operate on cubed tensors whose first three dimensions are acted on by node permutations, while channels remain unpermuted.Orbits partition index tuples into equality patterns that define face-vectors and orbit representations.
  • Subgraph interpretation: Node-based subgraph bags fit this tensor representation by indexing subgraphs on one axis and nodes on the other two.Diagonal and off-diagonal entries encode node representations and connectivity, respectively.
  • Subgraph interpretation: The orbit representations distinguish root nodes, non-root nodes, and incoming, outgoing, and non-root connectivity patterns.Xiii and Xijj store root and non-root representations, while Xiij, Xiji, and Xijk store connectivity types.
  • Computational primitives: Pooling and broadcasting operations provide the primitive transformations used to express aggregation, readout, and message-passing updates.Pooling can implement global summation, subgraph readout, or identity-like operations depending on which indices are removed.
  • Expressivity bound: 3-WL bounds the expressive power of node-based Subgraph GNNs, including ego-network and marked ego-network policies.The bound holds for every ego-network depth h > 0, so deeper ego-networks cannot exceed 3-WL.

D.2 Proofs for Section 6.1

The proofs show that ReIGN(2) layers can implement the update rules of known node-based Subgraph GNNs. The construction uses expanded local and global aggregation terms, with 3-IGN layers realizing the resulting computations.

  • General simulation: Known Subgraph GNN update rules can be implemented by appropriate stacks of ReIGN(2) layers.The proof embeds each model’s layer equations into a common equivariant layer formulation.
  • Existing architectures: DS-GNN uses one ReIGN(2) layer for independent subgraph-wise message passing, while DSS-GNN requires two layers for message passing and cross-bag aggregation.DSS-GNN first computes subgraph-local updates and cross-bag features, then propagates the aggregated representations over the original graph.
  • Existing architectures: GNN-AK and GNN-AK-ctx are recovered by selecting appropriate global or local ReIGN(2) aggregation terms, with some terms nullified for GNN-AK.The construction uses terms such as [#2.on], [#3.on], [#2.off], and [#4.off].
  • Existing architectures: ID-GNN is implemented with two ReIGN(2) layers that use distinct transformations for root and non-root nodes.The first layer separates transformed channels, and the second combines and aggregates them to complete message passing.
  • 3-IGN construction: Local aggregation is constructed through message broadcasting, sparsification, aggregation, and update steps implemented with 3-IGN layer stacks.Sparsification follows subgraph or original-graph connectivity before the messages are aggregated and incorporated into node representations.

F Extension to higher-order node policies

The paper proposes extending node-based subgraph policies from single nodes to k-tuples, yielding bags represented by higher-order tensors. It conjectures a corresponding (k+2)-WL expressivity bound but leaves the proof for future work.

  • Higher-order policies: Higher-order node policies select one subgraph for each k-tuple of nodes, producing bags with n^k subgraphs.A k-node deletion policy is a natural example, and the resulting bags can be arranged in tensors with k+2 dimensions.
  • Expressivity conjecture: The conjectured expressivity bound for k-order deletion, marking, and ego-network policies is (k + 2)-WL.This generalizes the paper’s 3-WL bound for single-node policies.
  • Open problem: The proposed proof strategy is to show that (k+2)-IGNs implement the higher-order bag construction and Subgraph GNN update steps.A corresponding ReIGN(k) could be developed by following the construction used for ReIGN(2).

G.1 Additional experiments

Additional experiments evaluate SUN across selection policies, synthetic graph-property and counting tasks, limited-data generalisation, training behaviour, and ablations. SUN performs strongly across these settings, with EGO policies leading on most TUDatasets and smaller train-test gaps than the GNN baseline.

  • TUDatasets: EGO policies obtain the best results on 6 of 7 TUDatasets, while NULL is not advantageous on this benchmark suite.SUN compares well with the best-performing approaches across domains and has smaller result variations than GNN-AK+.
  • Synthetic graph-property prediction: SUN achieves state-of-the-art performance on the Radius graph-property task and reports variability estimates across synthetic benchmarks.Radius targets are defined by the largest absolute eigenvalue of the graph adjacency matrix.
  • Limited-data generalisation: SUN generalises better than other Subgraph GNNs on 4-Cycles counting under node marking, after a short initial phase with similar performance.The comparison uses increasing fractions of available training data.
  • Training and validation analysis: SUN achieves better results than the GNN baseline on training, validation, and test sets while exhibiting a smaller train-test gap.The analysis attributes the GNN’s weaker overall learning behaviour to difficult optimisation, suboptimal solutions, and partial memorisation rather than complete underfitting.
  • Ablation study: Sequential ablations of SUN generally degrade performance, while removing P_j has little impact under EGO and can help under EGO+.Under EGO+, sharing parameters between root and non-root updates deteriorates performance.

G.2.3 OGBG-molhiv dataset

The OGBG-molhiv experiments follow a controlled evaluation procedure with fixed architectural choices and validation-based model selection. The setup also addresses validation instability through optimizer and regularization choices.

  • Evaluation protocol: OGBG-molhiv experiments use 10 seeds and report results at the epoch with the best validation metric.This evaluation procedure follows Hu et al.
  • Architecture: The models use depth-3 ego-networks, two GIN layers, residual connections, dropout 0.3, and 64-dimensional embeddings.These architectural choices follow Zhao et al.
  • Optimization: ASAM is used after Adam produced large validation fluctuations, with linear layers adopted instead of MLPs to further prevent overfitting.ASAM uses its default ρ value of 0.5 and promotes flatter minima.
  • Model selection: Architectures are selected by tuning hyperparameters on the full training and validation sets and choosing the configuration with the best validation performance.The competing models use matched tuning grids, while SUN’s grids are specified separately in the appendices.
Loading 2206.11140v3…