Equivariant Subgraph Aggregation Networks

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 . Owing to their scalability and simplicity , Message-Passing Neural Networks ( MPNNs ) are the leading Graph Neural Network ( GNN ) architecture for deep learning on graph-structured data . However , Morris et al . ( 2019 ) ; Xu et al . ( 2019 ) have shown that these architectures are at most as expressive as the Weisfeiler-Lehman ( WL ) graph isomorphism test ( Weisfeiler & Leman , 1968 ) . As a consequence , MPNNs can not distinguish between very simple graphs ( See Figure 1 ) . In light of this limitation , a question naturally arises : is it possible to improve the expressiveness of MPNNs ? Several recent works have proposed more powerful architectures . One of the main approaches involves higher-order GNNs equivalent to the hierarchy of the k-WL tests ( Morris et al. , 2019 ; 2020b ; Maron et al. , 2019b ; a ; Keriven & Peyré , 2019 ; Azizian & Lelarge , 2021 ; Geerts , 2020 ) , offering a trade-off between expressivity and space-time complexity . Unfortunately , it is already difficult to implement 3rd order networks ( which offer the expressive power of 3-WL ) . Alternative approaches use standard MPNNs enriched with structural encodings for nodes and edges ( e.g . based on cycle or clique counting ) ( Bouritsas et al. , 2020 ; Thiede et al. , 2021 ) or lift graphs into simplicial- ( Bodnar et al. , 2021b ) or cell complexes ( Bodnar et al. , 2021a ) , extending the message passing mechanism to these higher-order structures . Both types of approaches require a precomputation stage that , though reasonable in practice , might be expensive in the worst case . Our approach . In an effort to devise simple , intuitive and more flexible provably expressive graph architectures , we develop a novel framework , dubbed Equivariant Subgraph Aggregation Networks ( ESAN ) , to enhance the expressive power of existing GNNs . Our solution emerges from the observation that while two graphs may not be distinguishable by an MPNN , it may be easy to find distinguishable subgraphs . More generally , instead of encoding multisets of node colors as done in MPNNs and the WL test , we opt for encoding bags ( multisets ) of subgraphs and show that such an encoding can lead to a better expressive power . Following that observation , we advocate representing each graph as a bag of subgraphs chosen according to some predefined policy , e.g. , all graphs that can be obtained by removing one edge from the original graph . Figure 1 illustrates this idea . Bags of subgraphs are highly structured objects whose symmetry arises from both the structure of each constituent graph as well as the multiset on the whole . We propose an equivariant architecture specifically tailored to capture this object ’ s symmetry group . Specifically , we first formulate the symmetry group for a set of graphs as the direct product of the symmetry groups for sets and graphs . We then construct a neural network comprising layers that are equivariant to this group . Motivated by Maron et al . ( 2020 ) , these layers employ two base graph encoders as subroutines : The first encoder implements a Siamese network processing each subgraph independently ; The second acts as an information sharing module by processing the aggregation of the subgraphs . After being processed by several such layers , a set learning module aggregates the obtained subgraph representations into an invariant representation of the original graph that is used in downstream tasks . An integral component of our method , with major impact on its complexity and expressivity , is the subgraph selection policy : a function that maps a graph to a bag of subgraphs , which is then processed by our equivariant neural network . In this paper , we explore four simple — yet powerful — subgraph selection policies : node-deleted subgraphs , edge-deleted subgraphs , and two variants of ego-networks . To alleviate the possible computational burden , we also introduce an efficient stochastic version of our method implemented by random sampling of subgraphs according to the aforementioned policies . We provide a thorough theoretical analysis of our approach . We first prove that our architecture can implement novel and provably stronger variants of the well-known WL test , capable of encoding the multiset of subgraphs according to the base graph encoder ( e.g. , WL for MPNNs ) . Furthermore , we study how the expressive power of our architecture depends on different main design choices like the underlying base graph encoder or the subgraph selection policy . Notably , we prove that our framework can separate 3-WL indistinguishable graphs using only a 1-WL graph encoder , and that it can enhance the expressive power of stronger architectures such as PPGN ( Maron et al. , 2019a ) . We then present empirical results on a wide range of synthetic and real datasets , using several existing GNNs as base encoders . Firstly , we study the expressive power of our approach using the synthetic datasets introduced by Abboud et al . ( 2020 ) and show that it achieves perfect accuracy . We then evaluate ESAN on popular graph benchmarks and show that they consistently outperform their base GNN architectures , and perform better or on-par with state of the art methods . Main contributions . This paper offers the following main contributions : ( 1 ) A general formulation of learning on graphs as learning on bags of subgraphs ; ( 2 ) An equivariant framework for generating and processing bags of subgraphs ; ( 3 ) A comprehensive theoretical analysis of the proposed framework , subgraph selection policies , and their expressive power ; ( 4 ) An in-depth experimental evaluation of the proposed approach , showing noticeable improvements on real and synthetic data . We believe that our approach is a step forward in the development of simple and provably powerful GNN architectures , and hope that it will inspire both theoretical and practical future research efforts . 2 EQUIVARIANT SUBGRAPH AGGREGATION NETWORKS ( ESAN ) . In this section , we introduce the ESAN framework . It consists of ( 1 ) Neural network architectures for processing bags of subgraphs ( DSS-GNN and DS-GNN ) , and ( 2 ) Subgraph selection policies . We refer the reader to Appendix A for a brief introduction to GNNs , set learning , and equivariance . Setup and overview . We assume a standard graph classification/regression setting.1 We represent a graph with n nodes as a tuple G = ( A , X ) where A ∈ Rn×n is the graph ’ s adjacency matrix and X ∈ Rn×d is the node feature matrix . The main idea behind our approach is to represent the graph G as a bag ( multiset ) SG = { { G1 , . . . , Gm } } of its subgraphs , and to make a prediction on a graph based on this subset , namely F ( G ) : = F ( SG ) . Two crucial questions pertain to this approach : ( 1 ) Which architecture should be used to process bags of graphs , i.e. , How to define F ( SG ) ? , and ( 2 ) Which subgraph selection policy should be used , i.e. , How to define SG ? 2.1 BAG-OF-GRAPHS ENCODER ARCHITECTURE . To address the first question , we start from the natural symmetry group of a bag of graphs . We first devise an equivariant architecture called DSS-GNN based on this group , and then propose a particularly interesting variant ( DS-GNN ) obtained by disabling one of the components in DSS-GNN . Symmetry group for sets of graphs . The bag SG = { { G1 , . . . , Gm } } of subgraphs of G can be represented as tensor ( A , X ) ∈ Rn×n×m×Rn×d×m , assuming the number of nodes n and the number of subgraphs m are fixed . Here , A ∈ Rn×n×m represents a set of m adjacency matrices , and X ∈ Rn×d×m represents a set of m node feature matrices . Since the order of nodes in each subgraph , as well as the order of these subgraphs in SG , are arbitrary , our architecture must be equivariant to changes of both of these orders . Subgraph symmetries , or node permutations , are represented by the symmetric group Sn that acts on a graph via ( σ · A ) ij = Aσ−1 ( i ) σ−1 ( j ) and ( σ · X ) il = Xσ−1 ( i ) l. Set symmetries are the permutations of the subgraphs in SG , represented by the symmetric group Sm that acts on sets of graphs via ( τ · A ) ijk = Aijτ−1 ( k ) and ( τ · X ) ilk = Xilτ−1 ( k ) . These two groups can be combined using a direct product into a single group H = Sm × Sn that acts on ( A , X ) in the following way : ( ( τ , σ ) · A ) ijk = Aσ−1 ( i ) σ−1 ( j ) τ−1 ( k ) , ( ( τ , σ ) · X ) ilk = Xσ−1 ( i ) lτ−1 ( k ) In other words , τ permutes the subgraphs in the set and σ permutes the nodes in the subgraphs , as depicted in Figure 2 . Importantly , since all the subgraphs originate from the same graph , we can order their nodes consistently throughout , i.e. , the i-th node in all subgraphs represents the i-th node in the original graph.2 This setting can be seen as a particular instance of the DSS framework ( Maron et al. , 2020 ) applied to a bag of graphs . H-equivariant layers . Our goal is to propose anH-equivariant architecture that can process bags of subgraphs accounting for their natural symmetry ( the product of node and subgraph permutations ) . The main building blocks of such equivariant architectures are H-equivariant layers . Unfortunately , characterizing the set of all equivariant functions can be a difficult task , and most works limit themselves to characterizing the spaces of linear equivariant layers . In our case , Maron et al . ( 2020 ) characterized the space of linear equivariant layers for sets of symmetric elements ( such as graphs ) with a permutation symmetry group P ( Sn in our case ) . Their study shows that each such layer is composed of a Siamese component applying a linear P -equivariant layer to each set element independently and an information sharing component applying a different linear P -equivariant layer to the aggregation of all the elements in the set ( see Appendix A ) . 1The setup can be easily changed to support node and edge prediction tasks . 2This consistency assumption justifies the fact that we apply the same permutation σ to all subgraphs , rather than having a different node permutation ( σ1 , . . . , σm ) for each subgraph . Motivated by the linear characterization , we adopt the same layer structure with Siamese & information sharing components , and parameterize these ones using any equivariant GNN layer such as MPNNs . These layers L : Rn×n×m × Rn×d×m → Rn×n×m × Rn×d′×m map bags of subgraphs to bags of subgraphs , as follows : ( L ( A , X ) ) i = L1 ( Ai , Xi ) + L2 ( ∑m j=1Aj , ∑m j=1Xj ) . ( 1 ) Here , Aj , Xj are the adjacency and feature matrices of the j-th subgraph ( the j-th components of the tensorsA , X , respectively ) , and ( L ( A , X ) ) i denotes the output of the layer on the i-th subgraph . L1 , L2 : Rn×n × Rn×d → Rn×n × Rn×d′ represent two graph encoders and can be any type of GNN layer . We refer to L1 , L2 as the base graph encoders . When MPNN encoders parameterise L1 , L2 , the adjacency matrices of the subgraphs do not change , i.e. , the H-equivariant layer outputs the processed node features with the same adjacency matrix . These layers can also transform the adjacency structure of the subgraphs , e.g. , when using Maron et al . ( 2019a ) as the base encoder . While the encoder L1 acts on every subgraph independently , L2 allows the sharing of information across all the subgraphs in SG . In particular , the L2 module operates on the sum of adjacency and feature matrices across the bag . This is a meaningful operation because the nodes in all the subgraphs of a particular graph are consistently ordered . As in the original DSS paper , this sum aggregator for the adjacency and feature matrices ( ∑m j=1Aj and ∑m j=1Xj ) can be replaced with ( 1 ) Sums that exclude the current subgraph ( ∑m j 6=iAj and ∑m j 6=iXj ) or with ( 2 ) Other aggregators such as max and mean . We note that for the subgraph selection policies we consider in Subsection 2.2 , an entrywise max aggregation over the subgraph adjacency matrices , A = maxj Aj , recovers the original graph connectivity . This is a convenient choice in practice and we use it throughout the paper . Finally , the H-equivariance of the layer follows directly from the Sn-equivariance of the base-encoder , and the Sm-equivariance of the aggregation . Figure 3 ( right panel ) illustrates our suggested H-equivariant layer . DSS-GNN comprises three components : FDSS-GNN = Esets ◦Rsubgraphs ◦ Esubgraphs ( 2 ) The first component is an Equivariant Feature Encoder , Esubgraphs : Rn×n×m × Rn×d×m → Rn×n×m × Rn×d′×m implemented as a composition of several H-equivariant layers ( Equation 1 ) . The purpose of the subgraph encoder is to learn useful node features for all the nodes in all subgraphs . On top of the graph encoder , we apply a Subgraph Readout Layer Rsubgraphs : Rn×n×m × Rn×d′×m → Rd′×m that generates an invariant feature vector for each subgraph independently by aggregating the graph node ( and/or edge ) data , for example by summing the node features . The last layer , Esets : Rd ′×m → Rd′′ , is a universal Set Encoder , for example DeepSets ( Zaheer et al. , 2017 ) or PointNet ( Qi et al. , 2017 ) . Figure 3 ( left panel ) illustrates the DSS-GNN architecture . Intuitively , DSS-GNN encodes the subgraphs ( while taking into account the other subgraphs ) into a set of Sn-invariant representations ( Esubgraphs ◦ Rsubgraphs ) and then encodes the resulting set of subgraph representations into a single H-invariant representation Esets of the original graph G. DS-GNN is a variant of DSS-GNN obtained by disabling the information sharing component , i.e. , by setting L2 = 0 in Equation 1 for each layer in Esubgraphs . As a result , the subgraphs are encoded independently , and Esubgraphs is effectively a Siamese network . DS-GNN can be interpreted as a Siamese network that encodes each subgraph independently , followed by a set encoder that encodes the set of subgraph representations . DS-GNN is invariant to a larger symmetry group obtained from the wreath product Sn o Sm.3 In Section 3 , we will show that both DS-GNN and DSS-GNN are powerful architectures and that in certain cases DSS-GNN has superior expressive power .

Summary: This paper aims to improve GNNs' expressive power based on bag-of-subgraph representation. The authors develop a Equivariant Subgraph Aggregation Network, which uses equivariant layers to encode subgraphs and aggregate subgraph representations. The authors have performed theoretical analysis of the proposed framework, subgraph selection policies, and their expressive power. The experiments show that such an encoding can lead to a better expressive power on several real and synthetic graph classification datasets.