Dirichlet Energy Constrained Learning for Deep Graph Neural Networks

1 Introduction . Graph neural networks ( GNNs ) [ 1 ] are promising deep learning tools to analyze networked data , such as social networks [ 2 , 3 , 4 ] , academic networks [ 5 , 6 , 7 ] , and molecular graphs [ 8 , 9 , 10 , 11 ] . Based on spatial graph convolutions , GNNs apply a recursive aggregation mechanism to update the representation of each node by incorporating representations of itself and its neighbors [ 12 ] . A variety of GNN variations have been explored for different real-world networks and applications [ 13 , 14 ] . A key limitation of GNNs is that when we stack many layers , the performance would decrease significantly . Experiments show that GNNs often achieve the best performance with less than 3 layers [ 15 , 13 ] . As the layer number increases , the node representations will converge to indistinguishable vectors due to the recursive neighborhood aggregation and non-linear activation [ 16 , 17 ] . ∗Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Such phenomenon is recognized as over-smoothing issue [ 18 , 19 , 20 , 21 , 22 ] . It prevents the stacking of many layers and modeling the dependencies to high-order neighbors . A number of algorithms have been proposed to alleviate the over-smoothing issue and construct deep GNNs , including embedding normalization [ 23 , 24 , 25 ] , residual connection [ 26 , 27 , 28 ] and random data augmentation [ 29 , 30 , 31 ] . However , some of them are motivated directly by techniques in convolutional neural networks ( CNNs ) [ 32 ] , such as the embedding normalization and residual connection . Others are based on heuristic strategies , such as random embedding propagation [ 30 ] and dropping edge [ 29 ] . Most of them only achieve comparable or even worse performance compared to their shallow models . Recently , a metric of Dirichlet energy has been applied to quantify the over-smoothing [ 33 ] , which is based on measuring node pair distances . With the increasing of layers , the Dirichlet energy converges to zero since node embeddings become close to each other . But there is a lack of empirical methods to leverage this metric to overcome the over-smoothing issue . Therefore , it remains a non-trivial task to train a deep GNN architecture due to three challenges . First , the existing efforts are developed from diverse perspectives , without a generalizable principle and analysis . The abundance of these components also makes the design of deep GNNs challenging , i.e. , how should we choose a suitable one or combinations for real-world scenarios ? Second , even if an effective indicator of over-smoothing is given , it is hard to theoretically analyze the bottleneck and propose a generalizable principle to guide the training of deep GNNs . Third , even if theoretical guidance is given , it may be difficult to be utilized and implemented to train GNNs empirically . To this end , in this paper , we target to develop a generalizable framework with a theoretical basis , to handle the over-smoothing issue and enable effective deep GNN architectures . In particular , we will investigate two research questions . 1 ) Is there a theoretical and generalizable principle to guide the architecture design and training of deep GNNs ? 2 ) How can we develop an effective architecture to achieve state-of-the-art performance by stacking a large number of layers ? Following these questions , we make three major contributions as follows . • We propose a generalizable principle – Dirichlet energy constrained learning , to guide the training of deep GNNs by regularizing Dirichlet energy . Without proper training , the Dirichlet energy would be either too small due to the over-smoothing issue , or too large when the node embeddings are over-separating . Our principle carefully defines an appropriate range of Dirichlet energy at each layer . Being regularized within this range , a deep GNN model could be trained by jointly optimizing the task loss and energy value . • We design a novel deep architecture – Energetic Graph Neural Networks ( EGNN ) . It follows the proposed principle and could efficiently learn an optimal Dirichlet energy . It consists of three components , i.e. , orthogonal weight controlling , lower-bounded residual connection , and shifted ReLU ( SReLU ) activation . The trainable weights at graph convolutional layers are orthogonally initialized as diagonal matrices , whose diagonal values are regularized to meet the upper energy limit and eliminate the over-separating . The residual connection strength is determined by the lower energy limit to avoid the over-smoothing . While the widely-used ReLU activation causes the extra loss of Dirichlet energy , the linear mapping worsens the learning ability of GNNs . We apply SReLU with a trainable shift to provide a trade-off between the non-linear and linear mappings . • We show that the proposed principle and EGNN can well explain most of the existing techniques for deep GNNs . Empirical results demonstrate that EGNN could be easily trained to reach 64 layers and achieves surprisingly competitive performance on benchmarks . 2 Problem Statement . Notations . Given an undirected graph consisting of n nodes , it is represented as G = ( A , X ) , where A ∈ Rn×n denotes the adjacency matrix and X ∈ Rn×d denotes the feature matrix . Let Ã : = A+ In and D̃ : = D+ In be the adjacency and degree matrix of the graph augmented with selfloops . The augmented normalized Laplacian is then given by ∆̃ : = In− P̃ , where P̃ : = D̃− 1 2 ÃD̃− 1 2 is an augmented normalized adjacency matrix used for the neighborhood aggregation in GNN models . Node classification task . GNNs have been adopted in many applications [ 6 , 11 , 34 ] . Without loss of generality , we take node classification as an example . Given a graph G = ( A , X ) and a set of its nodes with labels for training , the goal is to predict the labels of nodes in a test set . We now use the graph convolutional network ( GCN ) [ 15 ] as a typical example , to illustrate how traditional GNNs perform the network analysis task . Formally , the layer-wise forward-propagation operation in GCN at the k-th layer is defined as : X ( k ) = σ ( P̃X ( k−1 ) W ( k ) ) . ( 1 ) X ( k ) and X ( k−1 ) are node embedding matrices at layers k and k − 1 , respectively ; W ( k ) ∈ Rd×d denotes trainable weights used for feature transformation ; σ denotes an activation function such as ReLU ; X ( 0 ) = X at the initial layer of GCN . The embeddings at the final layer are optimized with a node classification loss function , e.g. , cross-entropy loss . The recursive neighborhood aggregation in Eq . ( 1 ) will make node embeddings similar to each other as the number of layer k increases . This property , i.e. , over-smoothing , prevents traditional GNNs from exploring neighbors many hops away . In practice , the dependencies to high-order neighbors are important to the node classification . The traditional shallow GNNs may have sub-optimal performances in the downstream tasks [ 16 , 28 ] . 3 Dirichlet Energy Constrained Learning . In this paper , we aim to develop an effective principle to alleviate the over-smoothing issue and enable deep GNNs to leverage the high-order neighbors . We first theoretically analyze the over-smoothing issue , and then provide a principle to explain the key constraint in training deep GNNs . Node pair distance has been widely adopted to quantify the over-smoothing based on embedding similarities [ 19 , 23 ] . Among the series of distance metrics , Dirichlet energy is simple and expressive for the over-smoothing analysis [ 33 ] . Thus , we adopt Dirichlet energy and formally define it as below . Definition 1 . Given node embedding matrix X ( k ) = [ x ( k ) 1 , · · · , x ( k ) n ] > ∈ Rn×d learned from GCN at the k-th layer , the Dirichlet energy E ( X ( k ) ) is defined as follows : E ( X ( k ) ) = tr ( X ( k ) > ∆̃X ( k ) ) = 1 2 ∑ aij || x ( k ) i√ 1 + di − x ( k ) j√ 1 + dj ||22 , ( 2 ) where tr ( · ) denotes trace of a matrix ; aij is edge weight given by the ( i , j ) -th element in matrix A ; di is node degree given by the i-th diagonal element in matrixD . Dirichlet energy reveals the embedding smoothness with the weighted node pair distance . While a smaller value of E ( X ( k ) ) is highly related to the over-smoothing , a larger one indicates that the node embeddings are over-separating even for those nodes with the same label . Considering the node classification task , one would prefer to have an appropriate Dirichlet energy at each layer to separate the nodes of different classes while keeping those of the same class close . However , under some conditions , the upper bound of Dirichlet energy is theoretically proved to converge to 0 in the limit of infinite layers [ 33 ] . In other words , all nodes converge to a trivial fixed point in the embedding space . Based on the previous analysis , we derive the corresponding lower bound and revisit the oversmoothing/separating problem from the model design and training perspectives . To simplify the derivation process , we remove the non-linear activation σ , and re-express GCN as : X ( k ) = P · · ·PXW ( 1 ) · · ·W ( k ) . The impact of non-linear function will be considered in the model design . Lemma 1 . The Dirichlet energy at the k-th layer is bounded as follows : ( 1− λ1 ) 2s ( k ) minE ( X ( k−1 ) ) ≤ E ( X ( k ) ) ≤ ( 1− λ0 ) 2s ( k ) maxE ( X ( k−1 ) ) . ( 3 ) The detailed proof is provided in the Appendix . λ1 and λ0 are the non-zero eigenvalues of matrix ∆̃ that are most close to values 1 and 0 , respectively . s ( k ) min and s ( k ) max are the squares of minimum and maximum singular values of weight W ( k ) , respectively . Note that the eigenvalues of ∆̃ vary with the real-world graphs , and locate within range [ 0 , 2 ) . We relax the above bounds as below . Lemma 2 . The lower and upper bounds of Dirichlet energy at the k-th layer could be relaxed as : 0 ≤ E ( X ( k ) ) ≤ s ( k ) maxE ( X ( k−1 ) ) . ( 4 ) Besides the uncontrollable eigenvalues determined by the underlying graph , it is shown that the Dirichlet energy can be either too small or too large without proper design and training on weight W ( k ) . On one hand , based on the common Glorot initialization [ 35 ] and L2 regularization , we empirically find that some of the weight matrices approximate to zero in a deep GCN . The corresponding square singular values are hence close to zero in these intermediate layers . That means the Dirichlet energy will become zero at the higher layers of GCN and causes the over-smoothing issue . On the other hand , without the proper weight initialization and regularization , a large s ( k ) max may lead to the energy explosion and the over-separating . The Dirichlet energy plays a key role in training a deep GNN model . However , the optimal value of Dirichlet energy varies in the different layers and applications . It is hard to be specified ahead and then enforces the node representation learning . Therefore , we propose a principle – Dirichlet energy constrained learning , defined in Proposition 1 . It provides appropriate lower and upper limits of Dirichlet energy . Regularized by such a given range , a deep GNN model could be trained by jointly optimizing the node classification loss and Dirichlet energy at each layer . Proposition 1 . Dirichlet energy constrained learning defines the lower & upper limits at layer k as : cminE ( X ( k−1 ) ) ≤ E ( X ( k ) ) ≤ cmaxE ( X ( 0 ) ) . ( 5 ) We apply the transformed initial feature through trainable function f : X ( 0 ) = f ( X ) ∈ Rn×d . Both cmin and cmax are positive hyperparameters . From value interval ( 0 , 1 ) , hyperparameter cmin is selected by satisfying constraint of E ( X ( k ) ) ≥ ckminE ( X ( 0 ) ) > 0 . In such a way , the over-smoothing is overcome since the Dirichlet energies of all the layers are larger than appropriate limits related to ckmin . Compared with the initial transformed feature X ( 0 ) , the intermediate node embeddings of the same class are expected to be merged closely to have a smaller Dirichlet energy and facilitate the downstream applications . Therefore , we exploit the upper limit cmaxE ( X ( 0 ) ) to avoid overseparating , where cmax is usually selected from ( 0 , 1 ] . In the experiment part , we show that the optimal energy accompanied with the minimized classification loss locates within the above range at each layer . Furthermore , hyperparameters cmin and cmax could be easily selected from the large and appropriate value scopes , which do not affect the model performance . Given both the low and upper limits , an intuitive solution to search the optimal energy is to train GNNs by optimizing the following constrained problem : min Ltask + γ ∑ k ||W ( k ) ||F , s.t . cminE ( X ( k−1 ) ) ≤ E ( X ( k ) ) ≤ cmaxE ( X ( 0 ) ) , for k = 1 , · · · , K. ( 6 ) Ltask denotes the cross-entropy loss of node classification task ; K is layer number of GNN ; || · ||F denotes Frobenius norm of a matrix ; and γ is loss hyperparameter . Note that Dirichlet energy has also been adopted to regularize the node representation learning in shallow neural networks [ 36 , 37 , 38 ] . We instead focus on optimizing deep GNNs as shown in Eq . ( 6 ) , where K is often large .

Summary: This paper proposes a Dirichlet energy-based method to alleviate the over-smoothing and over-separating problems which are commonly observed in deep GNNs. Different from existing methods that rely on techniques from CNNs or heuristic strategies, this work provides theoretical proof of the lower and upper bound of the Dirichlet energy in GNNs and proposes EGNN which enables deeper GNNs. The effectiveness of EGNN is evaluated on commonly used benchmarks.