A Convergence Analysis of Gradient Descent on Graph Neural Networks

1 ) ) iterations . 1 Introduction . In the last decade , deep neural networks have been successfully used for a variety of machine learning tasks . In particular , optimization of various deep neural network architectures using gradient descent has been empirically ( and in some cases , theoretically ) shown to work surprisingly well , with an ability to obtain small training error on a variety of problems and datasets . In this paper , we focus on a particular class of neural network architectures known as Graph Neural Networks ( GNNs ) . Graph Neural Networks ( GNNs ) have become a popular choice for machine learning tasks with graph structured data [ Hamilton et al. , 2017 , Kipf and 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 embedding vector . After a few rounds of message passing , the computed node embeddings from all the nodes are aggregated to compute the final output [ Gilmer et al. , 2017 ] . This leads to a simple and elegant architecture for a variety of graph-related problems in practice . While theoretically understanding the optimization , learning , and generalization properties of deep neural networks is a challenging task in itself , our current understanding of GNNs lags significantly behind their more popular counterparts such as fully connected feed-forward networks and convolutional neural networks ( CNNs ) . In the context of GNNs recent theoretical works have focused on questions such as as the representation power of various GNN architectures [ Xu et al. , 2019b , a , Loukas , 2020a , b ] and establishing generalization bounds [ Garg et al. , 2020 ] . However , the optimization properties of GNN architectures are less explored . These involve identifying conditions under which algorithms like gradient descent can be shown to be provably effective in optimizing GNNs . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Such analysis has been carried out in recent years for fully connected feed-forward networks and CNNs , mainly restricted to depth-2 non-linear networks or deep linear networks [ Ge et al. , 2017 , Andoni et al. , 2014 , Bartlett et al. , 2017 , Arora et al. , 2018a , Hardt and Ma , 2016 ] . However , the message passing nature of GNNs makes it harder to theoretically analyze the optimization problem in GNNs , compared to other neural network architectures . We provide the first convergence analysis of gradient descent for GNNs , and provably show that gradient descent can recover near optimal solutions . We first consider the case of a one-round ( or one hidden-layer ) GNN with ReLU activations . This scenario already captures much of the complexity of the problem , in particular allowing for multiple local optima . Following standard assumptions made in the case of analyzing feed-forward networks , we show that when the input data distribution is the standard Gaussian distribution , and the training objective is to minimize least square regression loss , gradient descent converges to the population error of in O ( 1 2 log ( 1 ) ) iterations . In particular , our analysis applies to the standard practice of initializing the network with weights from a Gaussian distribution , and we do not perform any complicated pre-processing step to initialize the network . We then extend our result to the multi-round ( or multi-layer ) case , but for the case of linear activations ( analogous to the setting for deep linear networks [ Arora et al. , 2018a , Bartlett et al. , 2018 ] ) . In this case , we show a stronger convergence result ; gradient descent converges to an error in O ( log ( 1/ ) ) iterations . All our results are for the realizable setting , i.e , we assume that the data is indeed generated by an ( unknown ) underlying GNN , as is typical in analysis of gradient descent for other network architectures ( see the discussion of related work in Section A ) . The rest of the paper is organized as follows . In Section 2 we set up notation and discuss the model of GNNs that we consider . In Section 3 we present our first main convergence theorem for the case of one layer/round GNNs with ReLU activations . This is followed by our second main result on analyzing deep linear GNNs in Section 4 . We conclude with discussion and open problems in Section 6 . We discuss the most relevant related works in individual sections ( Section 3 and Section 4 ) , and provide a more comprehensive survey of related work in Section A . 2 Preliminaries . Graph Neural Networks operate via passing and aggregating messages among the nodes of a graph . Given an undirected unweighted graph G = ( V , E ) , let n = |V | be the number of nodes in the graph and m = |E| be the number of edges . Initially each node i ∈ [ n ] is initialized with its own private embedding xi ∈ Rr , where r denotes the embedding size . Then the computation in a GNN proceeds in a number of rounds where at round ` , the new embedding x ` i for vertex i is obtained via a combination of aggregate and combine steps Gilmer et al . [ 2017 ] as shown below . a ( ` ) i = AGGREGATE ( { x ` −1 j : j ∈ N ( i ) } ) ( 1 ) x ( ` ) i = COMBINE ( x ( ` −1 ) i , a ` i ) . ( 2 ) Here N ( i ) is the neighborhood of vertex i . When comparing to standard architectures such as feed-forward networks , it is useful to think of a round of message passing as computation performed by one layer of a more standard architecture . Hence an ` -round GNN shouldbe thought of as a network with ` hidden layers . Different choices of the aggregate and combine operations lead to different versions of GNNs such as graph convolutional networks ( GCNs ) [ Kipf and Welling , 2016 ] , GraphSAGE [ Hamilton et al. , 2017 ] , and graph isomorphism networks ( GINs ) [ Xu et al. , 2019b ] to name a few . The aggregate operation is typically a simple pooling operation such as the sum or average , and the combine operation is implemented via a low depth neural network . Furthermore , in the most popular implementation , i.e , GCNs , the network parameters are shared across the different rounds . This setting will also be the focus of study in our work . In particular we consider two problem settings . In the first case ( Section 3 ) we assume that the number of rounds in the GNN is one ( i.e . one hidden layer ) and that the combine operation is a depth one network with ReLU activations . In the second setting ( Section 4 ) we consider arbitrary round L GNNs but restrict the combine operation to be a linear network . In both the settings we will assume that the aggregate operation is a sum , and that the initial input embeddings for the nodes are drawn from the standard Gaussian distributions , i.e. , N ( 0 , Ir×r ) . Finally , our analysis for the case of ReLU activations will rely crucially on the notion of dual activations and their properties that we recap below . Definition 1 ( [ Daniely et al. , 2016 ] ) . The dual activation of σ is the function σ̂ : [ −1 , 1 ] 7→ R defined as σ̂ ( ρ ) = E [ σ ( X ) σ ( Y ) ] , ( 3 ) where X and Y are jointly Gaussian random variables with mean , zero variance one , and covariance ρ . The work of Daniely et al . [ 2016 ] showed that dual activations satisfy many nice properties such as continuity in [ −1 , 1 ] and the fact that they are convex in [ 0 , 1 ] . For a more extensive list of the properties of the dual activations please refer to Lemma 11 in Daniely et al . [ 2016 ] . 3 One round GNNs with ReLU activations . In ths section we present our main result on convergence of gradient descent for learning an unknown one round GNN with ReLU activations . We consider a graphG = ( V , E ) with n nodes and maximum degree d. Furthermore , we assume that there is an unknown GNN generating outcomes as y = n∑ i=1 σ ( W ∗x̄i ) , ( 4 ) where x̄i = ∑ j∈N ( i ) xj and each xj is drawn i.i.d . from N ( 0 , Ir×r ) , and σ ( t ) = √ 2 max ( t , 0 ) . In other words , the aggregate operation is a sum , and the combine operation is a depth-1 neural network with ReLU activations that produces an h-dimensional embedding for each node , where h is the number of hidden units ( inner dimensionality of W ∗ ) . Finally , the embeddings of all the nodes are summed up to produce an h-dimensional output for the graph G. The √ 2 factor multiplication to the ReLU function is for technical convenience and does not affect our results . The ground truth matrix W ∗ is an h× r matrix where we denote w∗1 , . . . , w∗h to be the rows of W ∗ . Without loss of generality we will assume that ‖w∗j ‖ = 1 , and again our results extend easily to the case when ‖wj‖ is bounded for all j ( see Appendix B ) . We train another one round GNN with unknown parameter matrix W to minimize the following loss via gradient descent . L ( W ) = 1 2 E [ ‖ŷ − y‖2 ] ( 5 ) = 1 2 E [ ∥∥ n∑ i=1 σ ( Wx̄i ) − n∑ i=1 σ ( W ∗x̄i ) ∥∥2 ] ( 6 ) = h∑ j=1 Lj ( wj ) : = h∑ j=1 1 2 E [ ( n∑ i=1 σ ( w > j x̄i ) − n∑ i=1 σ ( w∗j > x̄i ) ) 2 ] . ( 7 ) We will analyze the following gradient descent updates . Wt+1 = Wt − η∇L ( Wt ) . ( 8 ) Here∇L ( W ) denotes the gradient of the population loss and as is common in practice , we assume that the entries of the matrix W are initialized i.i.d . from a Gaussian distribution . We are now ready to state our main theorem below . Theorem 1 . Let W ∗ be the unknown parameter for a one-round GNN in Eq . ( 4 ) with ‖w∗j ‖ = 1 for j ∈ [ h ] , and let L ( W ) denote the population loss at W as defined in Eq . ( 5 ) . If the degree d of the graph is o ( √ n ) , then for any ∈ ( 0 , 1 ) , if W0 ∼ N ( 0 , I ) , with probability at least 1− h · e−c ′rwe have that L ( WT ) ≤ 2 provided T ≥ c · n 4h2 ( d+1 ) 2 log ( nh ) for absolute constants c , c ′ > 0 . Several remarks are in order regarding the above theorem . Note that the convergence rate is O ( 1 2 log ( 1 ) ) as a function of the desired error , and is polynomial in the other problem parameters such as the number of nodes , the maximum degree of the graph , and the dimensionality of the W ∗ matrix . In the context of feedforward neural networks , gradient descent based algorithms have been analyzed primarily in the Neural Tangent Kernel ( NTK ) regime where an unknown target is learned by performing gradient descent on a highly over-parameterized network [ Jacot et al. , 2018 , Du et al. , 2018 , Daniely , 2017 , Allen-Zhu et al. , 2019 ] . The advantage of such an analysis is that it can be carried out for large depth neural networks ( at least for smooth activations ) . However , due to the massive overparameterization the gradient descent updates move extremely slowly and the dynamics correspond to performing kernel regression in a very high dimensional space [ Lee et al. , 2019b , Arora et al. , 2019b ] . This is quite far from the behavior of gradient descent in realistic settings . We , on the other hand are interested in the setting where there is an unknown network generating the data and one would like to learn it via gradient descent with only mild over-parameterization . In particular , even the simple problem setting we consider in this section captures important aspects of the complexity of the general problem , in particular allowing for multiple local minima . There have been recent efforts in going beyond the NTK regime for the case of feed-forward networks . Similar to our setting , these works assume that the input is generated from the Gaussian distribution . Furthermore , these works still either need some over-parameterization to argue convergence [ Li et al. , 2020 ] , or need to add appropriate regularizers to the squared loss objective [ Ge et al. , 2017 ] . In contrast , an interesting aspect of our result is that we do not need any additional overparameterization and learn the unknown GNN using another network of the same size . The recent work of Zhang et al . [ 2020 ] studied a model for GNNs similar to our setting in Eq . ( 4 ) and designed a learning algorithm that first initializes the network weights based on a tensor decomposition subroutine , followed by an accelerated gradient descent procedure . We on the other hand do not perform any special initialization based on techniques such as spectral methods and analyze gradient descent updates starting from Gaussian initialization . The model of GNN we consider in Eq . ( 4 ) also bears similarity to the model of overlapping patches considered in the work of Du et al . [ 2017 ] in the context of convolutional neural networks . The authors analyze the convergence of gradient descent under certain strong assumptions on the correlation among the different patch distributions . We on the other hand make no such assumptions . Furthermore , even in the case of feed-forward networks , current analysis in the non-NTK regime do not extend beyond one hidden layer . Hence our work brings the understanding of gradient descent for GNNs on par with that of feed-forward networks . We next present the main ideas and key lemmas behind the proof of Theorem 1 .

Summary: This is a theoretical paper analyzing the optimization of graph neural networks (GNN). It assumes that the data is generated from a GNN with unknown weights, and tries to recover it by the gradient descent on the population loss. The convergence rate is established for two settings: one-layer GNNs with ReLU activations, and deep linear GNNs.