From Stars to Subgraphs: Uplifting Any GNN with Local Structure Awareness

1 INTRODUCTION . Graphs are permutation invariant , combinatorial structures used to represent relational data , with wide applications ranging from drug discovery , social network analysis , image analysis to bioinformatics ( Duvenaud et al. , 2015 ; Fan et al. , 2019 ; Shi et al. , 2019 ; Wu et al. , 2020 ) . In recent years , Graph Neural Networks ( GNNs ) have rapidly surpassed traditional methods like heuristically defined features and graph kernels to become the dominant approach for graph ML tasks . Message Passing Neural Networks ( MPNNs ) ( Gilmer et al. , 2017 ) are the most common type of GNNs owing to their intuitiveness , effectiveness and efficiency . They follow a recursive aggregation mechanism where each node aggregates information from its immediate neighbors repeatedly . However , unlike simple multi-layer feedforward networks ( MLPs ) which are universal approximators of continuous functions ( Hornik et al. , 1989 ) , MPNNs can not approximate all permutation-invariant graph functions ( Maron et al. , 2019b ) . In fact , their expressiveness is upper bounded by the first order Weisfeiler-Leman ( 1-WL ) isomorphism test ( Xu et al. , 2018 ) . Importantly , researchers have shown that such 1-WL equivalent GNNs are not expressive , or powerful , enough to capture basic structural concepts , i.e. , counting motifs such as cycles or triangles ( Zhengdao et al. , 2020 ; Arvind et al. , 2020 ) that are shown to be informative for bio- and chemo-informatics ( Elton et al. , 2019 ) . The weakness of MPNNs urges researchers to design more expressive GNNs , which are able to discriminate graphs from an isomorphism test perspective ; Chen et al . ( 2019 ) prove the equivalence between such tests and universal permutation invariant function approximation , which theoretically justifies it . As k-WL is strictly more expressive than 1-WL , many works ( Morris et al. , 2019 ; 2020b ) try to incorporate k-WL in the design of more powerful GNNs , while others approach k-WL expressiveness indirectly from matrix invariant operations ( Maron et al. , 2019a ; b ; Keriven & Peyré , 2019 ) and matrix language perspectives ( Balcilar et al. , 2021 ) . However , they require O ( k ) -order tensors to achieve k-WL expressiveness , and thus are not scalable or feasible for application on large , practical graphs . Besides , the bias-variance tradeoff between complexity and generalization ( Neal et al. , 2018 ) and the fact that almost all graphs ( i.e . O ( 2 ( n 2 ) ) graphs on n vertices , Babai et al . ( 1980 ) ) can be distinguished by 1-WL challenge the necessity of developing such extremely expressive models . In a complementary line of work , Loukas ( 2020a ) sheds light on developing more powerful GNNs while maintaining linear scalability , finding that MPNNs can be universal approximators provided that nodes are sufficiently distinguishable . Relatedly , several works propose to add features to make nodes more distinguishable , such as identifiers ( Loukas , 2020a ) , subgraph counts ( Bouritsas et al. , 2020 ) , distance encoding ( Li et al. , 2020 ) , and random features ( Sato et al. , 2021 ; Abboud et al. , 2021 ) . However , these methods either focus on handcrafted features which lose the premise of automatic learning , or create permutation sensitive features that hurt generalization . Present Work . Our work stands between the two regimes of extremely expressive but unscalable korder GNNs , and the limited expressiveness yet high scalability of MPNNs . Specifically , we propose a general framework that serves as a “ wrapper ” to uplift any GNN . We observe that MPNNs ’ local neighbor aggregation follows a star pattern , where the representation of a node is characterized by applying an injective aggregator function as an encoder to the star subgraph ( comprised of the central node and edges to neighbors ) . We propose a design which naturally generalizes from encoding the star to encoding a more flexibly defined subgraph , and we replace the standard injective aggregator with a GNN : in short , we characterize the new representation of a node by using a GNN to encode a locally induced encompassing subgraph , as shown in Fig.1 . This uplifts GNN as a base model in effect by applying it on each subgraph instead of the whole input graph . This generalization is close to Convolutional Neural Networks ( CNN ) in computer vision : like the CNN that convolves image patches with a kernel to compute new pixel embeddings , our designed wrapper convolves subgraphs with a GNN to generate new node embeddings . Hence , we name our approach GNN-AK ( GNN As Kernel ) . We show theoretically that GNN-AK is strictly more powerful than 1 & 2-WL with any MPNN as base model , and is not less powerful than 3-WL with PPGN ( Maron et al. , 2019a ) used . We also give sufficient conditions under which GNN-AK can successfully distinguish two non-isomorphic graphs . Given this increase in expressive power , we discuss careful implementation strategies for GNN-AK , which allow us to carefully leverage multiple modalities of information from subgraph encoding , and resulting in an empirically more expressive version GNN-AK+ . As a result , GNN-AK and GNN-AK+ induce a constant factor overhead in memory . To amplify our method ’ s practicality , we further develop a subgraph sampling strategy inspired by Dropout ( Srivastava et al. , 2014 ) to drastically reduce this overhead ( 1-3× in practice ) without hurting performance . We conduct extensive experiments on 4 simulation datasets and 5 well-known real-world graph classification & regression benchmarks ( Dwivedi et al. , 2020 ; Hu et al. , 2020 ) , to show significant and consistent practical benefits of our approach across different MPNNs and datasets . Specifically , GNN-AK+ sets new state-of-the-art performance on ZINC , CIFAR10 , and PATTERN – for example , on ZINC we see a relative error reduction of 60.3 % , 50.5 % , and 39.4 % for base model being GCN ( Kipf & Welling , 2017 ) , GIN ( Xu et al. , 2018 ) , and ( a variant of ) PNA ( Corso et al. , 2020 ) respectively . To summarize , our contributions are listed as follows : • A General GNN-AK Framework . We propose GNN-AK ( and enhanced GNN-AK+ ) , a general framework which uplifts any GNN by encoding local subgraph structure with a GNN . • Theoretical Findings . We show that GNN-AK ’ s expressiveness is strictly better than 1 & 2-WL , and is not less powerful than 3-WL . We analyze sufficient conditions for successful discrimination . • Effective and Efficient Realization . We present effective implementations for GNN-AK and GNN-AK+ to fully exploit all node embeddings within a subgraph . We design efficient online subgraph sampling to mitigate memory and runtime overhead while maintaining performance . • Experimental Results . We show strong empirical results , demonstrating both expressivity improvements as well as practical performance gains where we achieve new state-of-the-art performance on several graph-level benchmarks . Our implementation is easy-to-use , and directly accepts any GNN from PyG ( Fey & Lenssen , 2019 ) for plug-and-play use . See code at https : //github.com/GNNAsKernel/GNNAsKernel . 2 RELATED WORK . Improving Expressiveness of GNNs : Several works other than those mentioned in Sec.1 tackle expressive GNNs . Murphy et al . ( 2019 ) achieve universality by summing permutation-sensitive functions across a combinatorial number of permutations , limiting feasibility . Dasoulas et al . ( 2020 ) adds node indicators to make them distinguishable , but at the cost of an invariant model , while Vignac et al . ( 2020 ) further addresses the invariance problem , but at the cost of quadratic time complexity . Corso et al . ( 2020 ) generalizes MPNN ’ s default sum aggregator , but is still limited by 1-WL . Beani et al . ( 2021 ) generalizes spatial and spectral aggregation with > 1-WL expressiveness , but using expensive eigendecomposition . Recently , Bodnar et al . ( 2021b ) introduce MPNNs over simplicial complexes that shares similar expressiveness as GNN-AK . Ying et al . ( 2021 ) studies transformer with above 1-WL expressiveness . Azizian & Lelarge ( 2021 ) surveys GNN expressiveness work . Leveraging Substructures in Learning : Exploiting subgraph information in GNNs is not new ; in fact , k-WL considers all k node subgraphs . Monti et al . ( 2018 ) ; Lee et al . ( 2019 ) exploit motif information within aggregation , and others ( Bouritsas et al. , 2020 ; Barceló et al. , 2021 ) augment MPNN features with handcrafted subgraph based features . MixHop ( Abu-El-Haija et al. , 2019 ) directly aggregates k-hop information by using adjacency matrix powers , ignoring neighbor connections . Towards a meta-learning goal , G-meta ( Huang & Zitnik , 2020 ) applies GNNs on rooted subgraphs around each node to help transferring ability . Tahmasebi & Jegelka ( 2020 ) only theoretically justifies subgraph convolution with GNN by showing its ability in counting substructures . Zhengdao et al . ( 2020 ) also represent a node by encoding its local subgraph , however using non-scalable relational pooling . k-hop GNN ( Nikolentzos et al. , 2020 ) uses k-egonet in a specially designed way : it encodes a rooted subgraph via sequentially passing messages from k-th hops in the subgraph to k−1 hops , until it reaches the root node , and use the root node as encoding of the subgraph . Ego-GNNs ( Sandfelder et al. , 2021 ) computes a context encoding with SGC ( Wu et al. , 2019 ) as the subgraph encoder , and only be studied on node-level tasks . Both k-hop GNN and Ego-GNNs can be viewed as a special case of GNN-AK . You et al . ( 2021 ) designs ID-GNNs which inject node identity during message passing with the help of k-egonet , which is the receptive field of each node in k-layer GNN ( Hamilton et al. , 2017 ) – this differs from GNN-AK , which encodes each node by encoding its rooted subgraph . Unlike GNN-AK which uses rooted subgraphs , Fey et al . ( 2020 ) ; Thiede et al . ( 2021 ) ; Bodnar et al . ( 2021a ) design GNNs to use certain subgraph patterns ( like cycles and paths ) in message passing , however their preprocessing requires solving the subgraph isomorphism problem . To summarize , our work differs by ( i ) proposing a general subgraph encoding framework for uplifting GNNs , and ( ii ) addressing scalability issues involved with using subgraphs , which poses significant challenges for subgraph-based methods in practice . 3 GENERAL FRAMEWORK AND THEORY . We first introduce our setting and formalisms . Let G = ( V , E ) be a graph with node features xi ∈ Rd , ∀i ∈ V . We consider graph-level problems where the goal is to classify/regress a target yG by learning a graph-level representation hG . Let Nk ( v ) be the set of nodes in the k-hop egonet rooted at node v. N ( v ) = N1 ( v ) \v denotes the immediate neighbors of node v. For S ⊆ V , let G [ S ] be the induced subgraph : G [ S ] = ( S , { ( i , j ) ∈ E|i ∈ S , j ∈ S } ) . Then G [ Nk ( v ) ] denotes the k-hop egonet rooted at node v. We also define Star ( v ) = ( N1 ( v ) , { ( v , j ) ∈ E|j ∈ N ( v ) } ) be the induced star-like subgraph around v. We use { · } denotes multiset , i.e . set that allows repetition . Before presenting GNN-AK , we highlight the insights in designing GNN-AK and driving the expressiveness boost . Insight 1 : Generalizing star to subgraph . In MPNNs , every node aggregates information from its immediate neighbors following a star pattern . Consequently , MPNNs fail to distinguish any non-isomorphic regular graphs where all stars are the same , since all nodes have the same degree . Even simply generalizing star to the induced , 1-hop egonet considers connections among neighbors , enabling distinguishing regular graphs . Insight 2 : Divide and conquer . When two graphs are non-isomorphic , there exists a subgraph where this difference is captured ( see Figure 2 ) . Although a fixed-expressiveness GNN may not distinguish the two original graphs , it may distinguish the two smaller subgraphs , given that the required expressiveness for successful discrimination is proportional to graph size ( Loukas , 2020b ) . As such , GNN-AK divides the harder problem of encoding the whole graph to smaller and easier problems of encoding its subgraphs , and “ conquers ” the encoding with the base GNN .

Summary: This manuscript presents a general-purpose technique to improve GNN expressivity, dubbed as GNN-AK, that replaces the conventional neighbourhood/star-like aggregation (multi-set of neighbour embeddings) with an ego-net level aggregation. In particular, at each layer of GNN-AK, the authors first extract all $|\mathcal{V}|$ ego-nets of the original graph, then they apply a GNN on each one of them, and finally, they collect their outputs into node-wise representations.  The expressive power of different variants of their model is theoretically analysed (more expressive than 1-WL when the ego-net aggregation is as expressive as 1-WL, and no less powerful than 3-WL, when the ego-net aggregation is as expressive as 3-WL) and a series of design choices for a practical instantiation is provided. The authors extensively evaluate their method on both synthetic and real-world datasets and ablate the influence of some of the moving parts in the overall performance.