PaperlayerView exact revision on arXiv

Source-linked AI summary

Equivariant Subgraph Aggregation Networks

Beatrice Bevilacqua, Fabrizio Frasca, Derek Lim, Balasubramaniam Srinivasan, Chen Cai, Gopinath Balamurugan, Michael M. Bronstein, Haggai Maron

arXiv:2110.02910v3cs.LGstat.ML

TL;DR

MPNNs are scalable and simple but limited in expressive power, motivating methods that can distinguish graphs they cannot separate. ESAN represents graphs as bags of policy-selected subgraphs and processes them equivariantly, with theory and experiments showing stronger expressivity and improved benchmark performance.

  • Problem

    MPNNs are at most as expressive as the Weisfeiler-Leman test and cannot distinguish some simple graphs, motivating improved expressiveness.

  • Method

    ESAN represents each graph as a bag of subgraphs selected by predefined policies and processes the bag with an equivariant neural architecture.

  • Results

    ESAN increases the expressive power of existing GNN architectures, achieves perfect accuracy on synthetic expressivity datasets, and consistently outperforms base GNNs on popular graph benchmarks.

  • Takeaways & Limitations

    Subgraph-based equivariant processing provides a simple framework that improves the expressive power and performance of popular GNN architectures.

  • Takeaways & Limitations

    The framework assumes that nodes can be ordered consistently across all subgraphs originating from the same graph.

Abstract

from arXiv · show

Message-passing neural networks (MPNNs) are the leading architecture for deep learning on graph-structured data, in large part due to their simplicity and scalability. Unfortunately, it was shown that these architectures are limited in their expressive power. This paper proposes a novel framework called Equivariant Subgraph Aggregation Networks (ESAN) to address this issue. Our main observation is that while two graphs may not be distinguishable by an MPNN, they often contain distinguishable subgraphs. Thus, we propose to represent each graph as a set of subgraphs derived by some predefined policy, and to process it using a suitable equivariant architecture. We develop novel variants of the 1-dimensional Weisfeiler-Leman (1-WL) test for graph isomorphism, and prove lower bounds on the expressiveness of ESAN in terms of these new WL variants. We further prove that our approach increases the expressive power of both MPNNs and more expressive architectures. Moreover, we provide theoretical results that describe how design choices such as the subgraph selection policy and equivariant neural architecture affect our architecture's expressive power. To deal with the increased computational cost, we propose a subgraph sampling scheme, which can be viewed as a stochastic version of our framework. A comprehensive set of experiments on real and synthetic datasets demonstrates that our framework improves the expressive power and overall performance of popular GNN architectures.

1 INTRODUCTION

ESAN addresses the limited expressiveness of MPNNs by representing graphs as bags of subgraphs and processing them with equivariant architectures. Theory and experiments show stronger expressiveness and improved performance, while sampling mitigates added computational cost.

  • Motivation: MPNNs are scalable and simple but are limited to the expressiveness of the 1-WL graph isomorphism test.Consequently, they cannot distinguish some simple graphs.
  • Approach: ESAN represents each graph as a bag of subgraphs selected by a predefined policy and processes that bag with an equivariant architecture.The architecture combines graph neural networks with set networks.
  • Efficiency: Subgraph sampling provides a stochastic version intended to reduce the computational burden of the full framework.The approach is presented as a flexible framework for improving existing GNN architectures.
  • Theory: The framework analyzes how subgraph selection policies and equivariant architectures affect expressiveness, and proves separation of some 3-WL-indistinguishable graphs using a 1-WL encoder.It also enhances stronger architectures such as PPGN.
  • Experiments: ESAN achieves perfect accuracy on the synthetic EXP, CEXP, and CSL expressive-power tasks.These tasks are designed so that 1-WL GNNs cannot outperform random guessing.
  • Experiments: On popular graph benchmarks, ESAN variants consistently outperform their base GNNs and perform better than or on par with state-of-the-art methods.The reported evaluation covers both synthetic and real datasets.

2 EQUIVARIANT SUBGRAPH AGGREGATION NETWORKS (ESAN)

ESAN represents a graph as a tensorized bag of subgraphs and applies architectures equivariant to both subgraph and node permutations. DSS-GNN combines independent and information-sharing graph processing with subgraph readout and set encoding, while selection policies determine the resulting expressive power.

  • Overview: ESAN defines graph prediction by representing a graph as a bag of subgraphs and learning a function over that multiset.The framework must specify both the bag-processing architecture and the subgraph selection policy.
  • DSS-GNN: DSS-GNN layers combine a Siamese component that processes subgraphs independently with an information-sharing component that aggregates across the bag.The two components are parameterized by equivariant GNN layers.
  • Bag-of-Graphs Encoder Architecture: The bag is tensorized with m adjacency matrices and node-feature matrices, whose arbitrary subgraph and node orderings require equivariance.The architecture uses the product symmetry group H = S_m × S_n.
  • Bag-of-Graphs Encoder Architecture: The tensor action permutes subgraphs with τ and consistently permutes their nodes with σ.Consistent node ordering is assumed because all subgraphs originate from the same graph.
  • DSS-GNN: DSS-GNN encodes subgraphs, reads out an invariant vector for each, and applies a universal set encoder to obtain a graph representation.Its three blocks are an equivariant feature encoder, subgraph readout layer, and set encoder.
  • DS-GNN: DS-GNN disables information sharing, encoding subgraphs independently before applying a set encoder.The resulting architecture is effectively Siamese and can have different symmetry properties from DSS-GNN.
  • Subgraph Selection Policies: Subgraph selection policies map each graph to a permutation-invariant collection of subgraphs and directly affect complexity and expressive power.The paper studies node-deleted, edge-deleted, and two ego-network policies, plus random sampling for efficiency.

3 THEORETICAL ANALYSIS

Theoretical analysis shows that ESAN's subgraph-based WL analogues and neural architectures can be strictly more expressive than 1-WL, MPNNs, and, under some choices, 3-WL. Expressiveness depends on the architecture, base encoder, and subgraph selection policy.

  • 3.1 A WL Analogue for ESAN: DSS-WL and DS-WL encode bags of subgraphs using refinement procedures inspired by the WL test.DSS-WL shares information across subgraphs, whereas DS-WL refines each subgraph independently.
  • 3.2 WL Analogues and Expressive Power: There exist subgraph selection policies for which DS(S)-WL is strictly more powerful than 1-WL.The variants refine WL and can separate some WL-indistinguishable CSL graph pairs, including CSL(n, 2) versus CSL(n, k).
  • 3.2 WL Analogues and Expressive Power: DS(S)-GNNs can attain the distinctions made by DS(S)-WL, and suitable policies make both DSS-GNN and DS-GNN strictly more powerful than standard MPNNs.The result applies when each original-graph edge appears at least once in the subgraph bag.
  • 3.3 Theoretical Expressiveness of Design Choices: DSS-GNN is at least as powerful as DS-GNN for every policy and strictly more powerful for some policies.This establishes that information sharing across subgraphs can increase expressive power.
  • 3.3 Theoretical Expressiveness of Design Choices: Using a 3-WL base encoder with depth-1 EGO or EGO+ makes DS-GNN strictly more powerful than both 3-WL and its 1-WL-base version.Thus, ESAN's expressiveness gains are not limited to MPNN or 1-WL base encoders.
  • 3.3 Theoretical Expressiveness of Design Choices: With depth-n EGO+ policies, ND and EGO+ match 3-WL on SR graphs, while ED is strictly stronger than 3-WL.ND and EGO+ distinguish SR graphs with different parameters but not all graphs with identical parameters; ED separates some such pairs.

4 EXPERIMENTS

Experiments on synthetic and real benchmarks show that ESAN improves the expressive power and performance of base GNN encoders across tasks. The gains vary with the dataset, encoder, and subgraph policy, while stochastic sampling provides an efficient alternative.

  • Expressive power: EXP, CEXP, CSL: On EXP, CEXP, and CSL, ESAN variants substantially exceed the random-guess limitations of 1-WL GNNs and often achieve perfect accuracy.DS- and DSS-GNN perfectly solve EXP and CEXP with GIN or GraphConv, while both achieve perfect CSL performance across tested 1-WL encoders and policies.
  • TUDatasets: On TUDatasets, ESAN variants rank first, second, or third on seven datasets and improve over their base encoders in many cases.DSS-GNN improves over the base encoder in 51/56 cases and DS-GNN in 42/56 cases, with gains up to almost 5%.
  • TUDatasets: The stochastic ESAN variant generally outperforms the base encoder and occasionally surpasses the full approach.DSS-GNN tends to outperform DS-GNN and is less sensitive to subgraph sampling, though this comes with slightly greater computational complexity.
  • OGB: On OGBG-MOLHIV and OGBG-MOLTOX21, GIN-based ESAN variants improve base-encoder accuracy by up to 2.4% and 3%, respectively.DSS-GNN improves the base encoder in 14/16 cases, compared with 10/16 for DS-GNN.
  • ZINC12K: All ESAN configurations significantly outperform their GIN base encoder on ZINC12K under the 100k parameter budget.ESAN is reported as the best-performing provably expressive, domain-agnostic GNN and remains competitive with domain-specific expressive models.

5 CONCLUSION

ESAN enhances graph learning by decomposing graphs into subgraphs and processing their symmetry structure, improving expressiveness and benchmark performance while incurring higher computational complexity mitigated by stochastic sampling.

  • ESAN decomposes each graph into a set of subgraphs and processes that set while respecting its symmetry structure.
  • The framework increases the expressive power of existing graph-learning architectures and performs well on multiple graph-classification benchmarks.
  • The main limitation is increased computational complexity relative to standard MPNNs, which stochastic subgraph sampling is designed to mitigate.
  • ESAN builds on equivariant learning for structured sets and relates to methods including DropEdge, Ego-GNNs, and Reconstruction GNNs.
  • ESAN uses local, equivariant, and provably expressive subgraph policies that are computationally tractable independently of the graph distribution.

B.1 COMPARISON TO RECONSTRUCTION GNNS

Compared with Reconstruction GNNs, ESAN provides a more general equivariant framework with broader policies, higher-order encoders, and new WL-based expressiveness analyses.

  • Comparison to Reconstruction GNNs: Reconstruction GNN is exactly DS-GNN with a k-node-deletion subgraph selection policy.
  • Architecture for processing bags of subgraphs: DSS-GNN is theoretically stronger than DS-GNN for certain policies and achieves better empirical results in most tested settings.
  • General subgraph policies: ESAN studies node-deletion, edge-deletion, ego-net, and ego-net+ policies, while emphasizing that many other policies are possible.
  • General subgraph policies: On ZINC12K, EGO and EGO+ perform significantly better than node deletion, showing that node deletion is not always the optimal policy.
  • More general theoretical analysis: ESAN integrates higher-order graph encoders into its general framework, beyond the MPNNs and universal feed-forward networks considered by Reconstruction GNN.
  • Formulation of new WL variants: DS-WL and DSS-WL characterize DS-GNN and DSS-GNN expressiveness, whereas Reconstruction GNN develops no corresponding WL variant.

C.1 DATASETS AND MODELS IN COMPARISON

The experiments span thirteen diverse graph-classification datasets and compare DS-GNN and DSS-GNN, built on common MPNN encoders and four subgraph-selection policies.

  • Datasets: Experiments cover thirteen datasets from RNI, CSL, TUD, OGB, and ZINC12K, varying in graph count, features, and sparsity.
  • Expressiveness benchmarks: On RNI and CSL expressiveness benchmarks, ESAN variants are compared against base encoders and methods including ID-GNN, RNI, and DropEdge.
  • Models and policies: DS-GNN and DSS-GNN are evaluated with GIN, GraphConv, and GCN base encoders using node-deletion, edge-deletion, EGO, and EGO+ policies.
  • Evaluation setup: The evaluation follows prescribed benchmark procedures, with dataset-specific optimization, architecture, and hyperparameter settings.
  • OGB: The OGB experiments use scaffold splits for OGBG-MOLHIV and OGBG-MOLTOX21, with GIN and GCN baselines.
  • Expressiveness benchmarks: For CSL(41, 9) versus CSL(41, 12), DS-GNN with node deletion needs at least five layers, while EGO and EGO+ need depth 4.

C.3.1 STOCHASTIC SAMPLING RESULTS

Stochastic subgraph sampling reduces ESAN’s computational burden while retaining strong performance, and the broader comparisons indicate gains are not explained simply by larger models.

  • Stochastic sampling: Stochastic training samples 5%, 20%, or 50% of subgraphs per graph and uses majority voting over five sampled sets during evaluation.
  • Model size: The best ESAN configuration uses fewer learnable parameters than the best baseline on three of seven TUDatasets.
  • Model size: Increasing baseline GIN parameters slightly worsens average results and does not change the method ranking.
  • Model size: All ESAN configurations significantly outperform GIN on ZINC12K while remaining within the 100k parameter budget.
  • Timing: On NCI1, sampling reduces ESAN epoch times by 10–15%, while full DS-GNN and DSS-GNN take around 2x and 3x the GIN baseline time.
  • Timing: On ZINC12K, DS-GNN requires 3–3.5 seconds per epoch and DSS-GNN 3.7–4.25 seconds, compared with slightly less than 1.5 seconds for the base encoder.

C.4 DISCUSSION

ESAN improves expressiveness and benchmark performance over base encoders, with benefits depending on architecture, policy, and computational budget. Its strongest practical results occur when tasks require power beyond 1-WL, while OGB performance and some sampling-policy combinations remain less reliable.

  • Empirical expressiveness: ESAN perfectly solves the RNI and CSL expressiveness benchmarks, whereas GIN performs as a random guesser.
  • Real-world benchmarks: 91% and 75% of cases favor ESAN over its base encoder for DSS-GNN and DS-GNN, respectively, on real-world benchmarks.
  • Real-world benchmarks: ESAN performance is comparable to state-of-the-art methods on TUDatasets and significantly exceeds provably expressive, domain-agnostic GNNs on ZINC12K.
  • Policy comparison: Subgraph policies show no definite domain-performance correlation, although EGO(+) policies are generally more consistent across datasets.
  • Stochastic variant: Subgraph sampling usually has marginal to no impact, but it can negatively affect DS-GNN with EGO(+) policies while DSS-GNN is more robust.
  • Practical implications: ESAN is recommended when tasks may require expressive power beyond 1-WL or when relevant graph substructures lack domain-specific engineering.
  • Limitations: ESAN has not shown a consistently unsuccessful application, but its combination with GCN encoders on OGB datasets performs poorly enough to require further investigation.

D.2.2 EXPRESSIVENESS OF WL VARIANTS

The proposed DSS-WL and DS-WL variants retain at least the distinguishing power of standard WL under specified policies and can separate graph pairs that WL cannot. Their relative strength and guarantees depend on whether policies preserve vertices or include the original graph.

  • Relation to WL: DSS-WL and DS-WL are at least as powerful as standard WL under non-trivial subgraph selection policies.
  • DSS-WL versus WL: For vertex set preserving policies, DSS-WL refines WL and distinguishes every non-isomorphic pair that WL distinguishes.
  • DS-WL guarantees: Including the original graph in an augmented policy gives DS-WL a lower bound matching standard WL.
  • Separating WL-indistinguishable graphs: DS-WL separates CSL(n, 2) from CSL(n, k) for k in [3, n/2 −1] under ND, EGO, and EGO+ policies, although 1-WL cannot.
  • Separating WL-indistinguishable graphs: The pair C6 and 2 × C3 is disambiguated by both DS-WL and DSS-WL under edge-deleted policy, despite 1-WL assigning every node degree-2 color.

D.2.3 NEURAL COUNTERPARTS

The neural counterparts simulate DSS-WL through equivariant layers operating on edge-covering subgraph bags. Under bounded graph size and degree, a finite-depth DSS-GNN can reproduce any finite number of DSS-WL refinement steps.

  • Neural simulation: DSS-GNN uses GraphConv encoders and aggregated subgraph adjacency information to construct representations for nodes within subgraphs.
  • Expressive equivalence: For edge-covering policies and bounded graph families, a DSS-GNN with D = 10T layers uniquely represents DSS-WL colors after T iterations.
  • Assumptions: The simulation relies on finite color palettes and bounded bag cardinalities, enabling unique encodings of all relevant multisets.
  • Simulation construction: The construction initializes subgraph node colors, simulates refinement through successive layers, and uses one-hot outputs to represent HASH results uniquely.
  • Neural simulation: One-hot encodings allow injective multiset aggregation, while additional layers uniquely encode finite collections of intermediate representations.

E.1 PROOFS OF DSS VERSUS DS

These proofs show that DSS-GNN can distinguish graph pairs that DS-GNN cannot, because node-level interactions across subgraphs capture structure that independent subgraph representations miss.

  • DSS-GNN is at least as powerful as DS-GNN, since setting L2 = 0 recovers any DS-GNN layer.The strict separation is established with two graph-pair examples.
  • For a path and a 3-star, DS-GNN cannot distinguish the graphs because their bags contain equally colored single-edge subgraphs.DSS-GNN instead uses graph-aware operations that recover differing node structure.
  • DSS-GNN distinguishes the path and 3-star by applying graph encoders that expose adjacency structure within each subgraph.The encoders can be implemented by MPNNs such as GIN and GraphConv.
  • The first example is limited because adding the original graph to the subgraph bag would let DS-GNN distinguish the graphs using their different degree distributions.This illustrates that separation depends on the chosen subgraph policy.
  • For two triangles versus a 6-cycle, DS-GNN still fails under dSE because the original graphs and all single-edge subgraphs are 1-WL equivalent.Both bags have size 7, and their multisets of subgraph colors are equal.
  • DSS-GNN separates the two-triangle graph from the 6-cycle after two layers by detecting whether an edge’s endpoints share a common neighbor.The relevant second-layer feature exists in the triangle graph but not in the cycle.

E.2 PROOFS OF INCREASING ENCODER EXPRESSIVENESS

The encoder-expressiveness proofs compare DS-GNN variants using 1-WL and 3-WL encoders on strongly regular graphs. A 3-WL encoder separates the Rook’s and Shrikhande graphs through their ego-nets, although 1-WL cannot.

  • The Rook’s and Shrikhande graphs are non-isomorphic strongly regular graphs with the same parameters and are indistinguishable by 3-WL.They therefore provide a test pair for comparing encoder expressiveness.
  • DS-GNN with a 3-WL encoder applies a DeepSets module to the multiset of 3-WL representations of selected subgraphs.This generalizes the DS-GNN construction by replacing the base graph encoder.
  • With depth-1 EGO or EGO+ policies, DS-GNN using 1-WL cannot distinguish the graphs because their ego-nets reduce to indistinguishable two-triangle and 6-cycle structures.The root is connected to all nodes in both ego-nets.
  • A 3-WL encoder distinguishes those ego-nets because 3-WL can distinguish two disjoint triangles from a 6-cycle by counting triangles at initialization.Thus the encoder choice increases the separation available to the subgraph architecture.

E.3 PROOFS OF POLICY CHOICE AND EXPRESSIVENESS

The proofs show that ESAN’s expressive power depends strongly on the subgraph selection policy and ego-net depth. Different policies distinguish different classes of strongly regular graphs, while ego-net depth can help or hurt separation.

  • EGO policy depth: Depth-1 and depth-3 EGO policies fail to distinguish CSL(12,3) from CSL(12,5), whereas depth-2 EGO succeeds.The depth-2 ego-nets differ in degree distributions, including the presence of degree-3 nodes only for CSL(12,3).
  • EGO policy depth: Depth-1 EGO fails on CSL(12,3) and CSL(12,5) because all selected ego-nets are isomorphic K1,4 graphs.Consequently, 1-WL cannot distinguish any selected subgraphs.
  • EGO policy depth: Depth-3 EGO also fails because both graphs have diameter 3, making each selected ego-net the original regular graph.The original graphs have the same node count and degree, so 1-WL cannot separate them.
  • Strongly regular graphs: ND distinguishes strongly regular graphs with different parameters but cannot distinguish strongly regular graphs sharing the same parameters.The stable 1-WL coloring of node-deleted subgraphs depends only on those parameters.
  • Strongly regular graphs: Depth-n EGO+ has the same parameter-based boundary: it distinguishes strongly regular graphs with different parameters but not graphs with identical parameters.This result parallels the ND policy’s limitation.

E.4.2 ED ON STRONGLY REGULAR GRAPHS

For edge-deleted subgraphs of strongly regular graphs, three 1-WL iterations encode the strongly regular parameters, while the Rook’s and Shrikhande graphs require a fourth iteration to separate.

  • Three 1-WL iterations on edge-deleted strongly regular graphs produce colorings determined by the parameters and distinguish graphs with different parameters.The proof tracks colors of the deleted edge’s endpoints and nodes adjacent to zero, one, or two endpoints.
  • The iteration-3 color counts are λ, 2(k −1 −λ), and n + λ −2k for the three degree-k node colors.These counts depend only on the strongly regular parameters.
  • Graphs differing in n, k, or λ are separated by the third iteration, and graphs differing in µ are separated through its relation to the other parameters.Thus the edge-deleted coloring distinguishes strongly regular graphs with different parameters.
  • The 4 × 4 Rook’s and Shrikhande graphs are distinguished by DS-GNN and DSS-GNN with the ED policy.All edge-deleted subgraphs within each graph are isomorphic, so one representative from each graph suffices.
  • The Rook’s and Shrikhande edge-deleted subgraphs first receive different 1-WL colorings at iteration 4.Tables 12 and 13 document the respective color updates for deleting edge (1,2).

E.5 PROOF OF PROPOSITION 3

The proof establishes that ESAN-based policies distinguish several strongly regular graph families that standard WL tests cannot separate. In particular, edge-deleted subgraphs distinguish the Rook’s and Shrikhande graphs, exceeding 3-WL on this family.

  • Node-deleted, depth-n EGO+, and edge-deleted policies distinguish strongly regular graphs with different parameters.
  • The Rook’s graph and the Shrikhande graph are distinguished by DS-GNN with the edge-deleted policy despite sharing strongly regular parameters.
  • 3-WL cannot distinguish strongly regular graphs with the same parameters, whereas edge-deleted ESAN is strictly more powerful on this family.

F.1 FORWARD PASS

The forward-pass complexity depends on the number and density of selected subgraphs, with sparse graphs yielding substantial improvements over several expressive graph networks. ESAN can match GSN only for constant-size stochastic subgraph sets, while its preprocessing remains policy- and graph-dependent.

  • O(|S|n^2) time applies to dense graphs when each of |S| subgraphs is processed by an O(n^2) MPNN.
  • O(n^3) complexity matches PPGN and 3-IGN for node-deleted and ego-net policies with |S| = O(n).
  • O(n^4) complexity matches 3-GNN for edge-deleted policies with |S| = O(n^2).
  • O(|S|n∆max) time benefits sparse graphs, gaining O(n) over PPGN and 3-IGN and O(n^2) over 3-GNN when ∆max = O(1).
  • ESAN matches GSN’s complexity only when |S| = O(1), corresponding here to stochastic policies, while CWN’s complexity can grow substantially with general graph distributions.
  • Preprocessing is a one-off computation whose costs include O(nm) for node-deleted and O(m^2) for edge-deleted policies, with practical runtime depending on graph sparsity.
  • GSN and CWN preprocessing is generally O(n^k) for generic substructures, and worst-case exponential complexity may hinder use on graphs with uncharacterized distributions.
Loading 2110.02910v3…