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: The paper deals with supervised learning with graphs, specifically graph-level tasks. The paper proposes an algorithm to overcome the expressive limits of the standard GNNs, which are upper-bounded by the 1-WL. The main idea of the paper is to, for each node v, extract the subgraph induced by nodes of at most distance k to node v, and then deploy a GNN on top of each of these subgraphs. Further, the authors study the gains of the expressive power of this architecture compared to the 1-, 2-, 3-WL, and k-WL.    Moreover, the authors propose a subgraph sampling strategy to speed up the computation.   The proposed method is evaluated on large benchmark datasets, mostly stemming from the molecular domain, reporting good performance boosts compared to standard GNN architectures.