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 proposes to improve the expressivity of MPNN by using H-Equivariant layers to process bags of subgraphs accounting for their natural symmetry. The proposed DS(S)-GNN is proven to be equivalently powerful as DS(S)-WL, which can be strictly more powerful than 1-WL. It also provides theoretical results about how design choices such as the subgraph selection policy and equivariant neural architecture affect the architecture’s expressive power. In addition, a stochastic sampling scheme is proposed to mitigate the computational overhead.