Beyond GNNs: A Sample Efficient Architecture for Graph Problems

1 INTRODUCTION . In recent years Graph Neural Networks ( GNNs ) have become the predominant paradigm for learning problems over graph structured data ( Hamilton et al. , 2017 ; Kipf & Welling , 2016 ; Veličković et al. , 2017 ) . Computation in GNNs is performed by each node sending and receiving messages along the edges of the graph , and aggregating messages from its neighbors to update its own embedding vector . After a few rounds of message passing , the computed node embeddings are aggregated to compute the final output ( Gilmer et al. , 2017 ) . The analogy to message passing leads to a simple and elegant architecture for learning functions on graphs . On the other hand , from a theoretical and practical perspective , we also need these architectures to be sample efficient , i.e. , learnable from a small number of training examples , where each training example corresponds to a graph . Recent works have shown that generalization in GNNs depends upon the depth of the architecture , i.e. , the number of rounds of message passing , as well as the embedding size for each node in the graph ( Garg et al. , 2020 ) . However , this requirement is in fundamental conflict with the message passing framework . In particular , using GNNs to compute several fundamental graph problems such as shortest paths , minimum spanning tree , min cut etc. , necessarily requires the product of the depth of the GNN and the embedding size to scale as √ n where n is the size of the graph ( Loukas , 2020 ) . This in turn places a significant statistical burden when learning these fundamental problems on large scale graphs . The above raises the the following question : Can one develop sample efficient architectures for graph problems while retaining the simplicity of the message passing framework ? Several recent works have tried to address the above question by proposing extensions to the basic GNN framework by augmenting various pooling operations in conjunction with message passing rounds to capture more global structure ( Ying et al. , 2018 ; Simonovsky & Komodakis , 2017 ; Fey et al. , 2018 ) . While these works demonstrate an empirical advantage over GNNs , we currently do not know of a general neural architecture that is versatile enough to provably encode the solution space of a variety of graph problems such as shortest paths and minimum spanning trees , while being significantly superior to GNNs in terms of statistical efficiency . In this work we propose a theoretically principled architecture , called GNN+ networks for learning graph problems . While the basic GNN framework is inspired from classical message passing style models studied in distributed computing , we borrow from two fundamental paradigms in graph algorithm design namely , sketching and parallel computation , to design GNN+ networks . As a result of combining these two powerful paradigms , we get a new neural architecture that simultaneously achieve low depth and low embedding size for many fundamental graph problems . As a result our proposed GNN+ architecture have a significantly smaller number of parameters that provably leads to better statistical efficiency than GNNs . Before we present our improved architecture , we briefly describe the standard GNN framework . Model for GNNs . In this work we will study GNNs that fall within the message passing framework and using notation from previous works we denote such networks as GNNmp ( Loukas , 2020 ) . A GNNmp network operates in the AGGREGATE and COMBINE model ( Gilmer et al. , 2017 ) that captures many popular variants such as GraphSAGE , Graph Convolutional Networks ( GCNs ) and GIN networks ( Hamilton et al. , 2017 ; Kipf & Welling , 2016 ; Xu et al. , 2019a ) . Given a graph G = ( V , E ) , let x ( k ) i denote the feature representation of node i at layer k. Then we have a ( k−1 ) i = AGGREGATE ( { x ( k−1 ) j : j ∈ N ( i ) } ) ( 1 ) x ( k ) i = COMBINE ( x ( k−1 ) i , a ( k−1 ) i ) . ( 2 ) Here N ( i ) is the set of neighbors for node i . Typically the aggregation and combination is performed via simple one or two layer full connected networks ( FNNs ) , also known as multi layer perceptrons ( MLPs ) . In the rest of the paper we will use the two terms interchangeably . GNN+ Networks . Our proposed GNN+ networks consist of one or more layers of a GNN+ block shown in Figure 1 . The GNN+ block comprises of r parallel GNNmp networks follows by s parallel fully connected network modules for pooling where r and s are the hyperparameters of the architecture . More importantly we restrict the r GNNmp modules to share the same set of weights . Hence the parallel GNNmp modules only differ in the way the node embeddings are initialized . Furthermore , we restrict each GNNmp to be of low depth . In particular , for degree-d graphs of diameter D , over n nodes , we will restrict the GNNmp to be of depth O ( ( d + D ) · polylog ( n ) ) . Similarly , we require the s fully connected networks to be of depth O ( ( d + D ) · polylog ( n ) ) and share the network weights . We connect the outputs of the GNNmp modules to the fully connected pooling networks in a sparse manner and restrict the input size of each fully connected network to be O ( ( d + D ) · polylog ( n ) ) . Stacking up L layers of GNN+ blocks results in a GNN+ network that is highly parameter efficient and in total has O ( ( d + D ) L · polylog ( n ) ) parameters . For such a network we call the depth as the total number of message passing rounds and the number of MLP layers used across all the L stacks . Since we restrict our MLPs and GNNmp blocks inside a GNN+ network to be of low depth , we will often abuse notation and refer to a GNN+ architecture with L stacks of GNN+ blocks as a depth L architecture . Our proposed design lets us alternate between local computations involving multiple parallel GNN blocks and global post-processing stages , while still being sample efficient due to the enforced parameter sharing . We will show via several applications that optimal or near-optimal solutions to many popular graph problems can indeed be computed via a GNN+ architecture . Below we briefly summarize our main results . Overview of Results . To demonstrate the generality of our proposed GNN+ architecture , we study several fundamental graph problems and show how to construct efficient GNN+ networks to compute optimal or near optimal solutions to these problems . In particular , we will focus on degree-d graphs , i.e. , graphs of maximum node degree d , with n nodes and diameter D and will construct GNN+ networks of depth polylog ( n ) and O ( ( D + d ) polylog ( n ) ) total parameters . Shortest Paths . The first problem we consider is the fundamental graph problem of computing ( approximate ) all pairs shortest paths in undirected graphs . Given a graph G = ( V , E ) , let dG ( u , v ) be the shortest path between nodes u and v. We say that an output { d̃G ( u , v ) : u , v ∈ V } is an α-approximate solution if for all u 6= v it holds that dG ( u , v ) ≤ d̃G ( u , v ) ≤ αdG ( u , v ) . We construct efficient GNN+ networks for all pairs shortest paths with the following guarantee . Theorem 1 ( Informal Theorem ) . For any constant c > 1 , there is a depth O ( D log d+ log n ) GNN+ network with O ( ( n 2 c + d ) polylog ( n ) ) parameters that computes ( 4c − 2 ) -approximate all pairs shortest paths in the undirected unweighted degree-d graphs over n nodes . On the other hand , computing a c-approximate shortest paths using GNNmp networks requires a network of depth Ω ( n ) . From the above theorem we see that by setting c = O ( log n ) we can encode a c-approximate solution using an O ( D log d+ log n ) GNN+ network with only O ( d ·polylog ( n ) ) parameters . This is in stark contrast with the depth requirement of Ω ( n ) for the traditional GNNmp networks . Connectivity Measures . Next we consider computing various graph connectivity measures . We first study the popular measure based on graph effective resistances ( Chandra et al. , 1996 ) . Definition 1 ( Effective Resistance ) . Let G be a weighted undirected graph G with adjacency matrix A and the associated Laplacian L = D − A . Given an edge u , v , the effective resistance between u , v is defined as Ru , v = ξ > u , vL †ξu , v . Here ξu , v is an n dimensional vector with +1 at position u , −1 at position v and zeros everywhere . L† refers to the matrix pseudo-inverse . We also study the following connectivity measure that was proposed by Panigrahy et al . ( 2012 ) in the context of web graphs . Given an undirected graph G , let Gp be the random graph obtained by sampling each edge with probability p. Definition 2 ( Affinity ) . For any two vertices u , v and for p ∈ [ 0 , 1 ] , define Ap ( u , v ) to be the probability that u , v are connected in Gp . Then the affinity between u and v is defined as A ( u , v ) = Ep [ Ap ( u , v ) ] where the expectation is taken over p drawn from the uniform distribution in [ 0 , 1 ] . For the above measures we show the following Theorem 2 ( Informal Theorem ) . There exists a GNN+ architecture with O ( D log ( nd ) ) parameters , and depth O ( D log ( nd ) ) on graphs of diameter D with n nodes and maximum degree d , that approximate the above connectivity measures up to constant factors . On the other hand using GNNmp networks to compute the above measures , even approximately , necessarily requires a network of depth Ω ( √ n ) . Clustering , Minimum Cuts and Minimum Spanning Trees . Finally , we showcase the power of a GNN+ architecture for computing other fundamental graph problems . Given an undirected graph G , the spectral clustering of G corresponds to the cut obtained by taking the sign of the eigenvector v corresponding to the second smallest eigenvalue λ2 ( L ) , where L is the graph Laplacian . For computing the spectral clustering via GNN+ networks we show the following Theorem 3 ( Informal Theorem ) . There is a GNN+ network of depth ` = O ( 1λ2 ( L ) 2 log n ) , with O ( d ) parameters that computes an -approximate spectral clustering on graphs of degree d. On the other hand , using GNNmp networks to even approximately compute the spectral clustering requires depth Ω ( √ n ) . Next we consider the classical problems of computing a global minimum cut and minimum spanning trees in undirected graphs . Theorem 4 ( Informal Theorem ) . There exist GNN+ networks of of depth O ( ( D+ log n ) log n ) , and O ( d ) parameters for computing a global minimum cut ( MINCUT ) and minimum spanning tree ( MST ) in degree d graphs of diameter D. Furthermore , using GNNmp networks to compute these primitives ( even approximately ) necessarily requires depth Ω ( √ n ) . Generalization Bounds . Our final result concerns the generalization properties of a depth L GNN+ architecture . For ease of exposition , we state here the results for the case when the GNN+ architecture produces a one dimensional output . More general results are presented in Appendix D. Our generalization bounds depend on the depth L and the total number of parameters P in the GNN+ network . Following recent work on providing generalization bounds for fully connected and convolutional neural networks ( Bartlett et al. , 2017 ; Long & Sedghi , 2020 ) that are based on distance to initialization , we consider the class Fβ of depth L GNN+ networks with P parameters that are at a distance β from a reference parameter configuration ( typically the parameters at random initialization ) . Let y ∈ R denote the output of the network and consider a Lipschitz loss function ` ( y , ŷ ) . Then , we provide following guarantee . Theorem 5 ( Informal Theorem ) . Let ` ( ŷ , y ) be a Lipschitz loss function bounded in [ 0 , B ] . Then , given m i.i.d . samples ( G1 , y1 ) , ( G2 , y2 ) , . . . ( Gm , ym ) generated from a distribution D , with probability at least 2/3 , it holds that for all f ∈ Fβ , ∣∣∣ÊD [ ` f ] − ED [ ` f ] ∣∣∣ ≤ O ( B√P ( β + L ) m ) . We refer the reader to Theorem 16 in Appendix D for a formal statement and the proof . Notice that the above theorem implies that our proposed GNN+ architecture for the above graph problems can indeed be trained using very few samples as opposed to the traditional GNNmp networks since the GNN+ network requires much fewer parameters and depth . Furthermore , since a GNNmp network is a special case of a GNN+ architecture , our analysis also leads to an improved bound on the generalization guarantees for GNNmp networks as well . In particular , the above theorem improves upon the recent work of Garg et al . ( 2020 ) that provides generalization guarantees for training GNNs that scale with the branching factor of the graph . Using our improved analysis we are able to remove this dependence on the branching factor . See Appendix D for details .

Summary: This paper proposes a new building block for GNNs, called GNN$^+$. This building block trades of depth for width and involves multiple parallel regular GNN processing units. Using the GNN$^+$ architecture, the authors establish bounds for the required network depth (and total parameters) for several combinatorial problems over graphs.