Local Augmentation for Graph Neural Networks

1 INTRODUCTION . Graph Neural Networks ( GNNs ) and their variants ( Abu-El-Haija et al. , 2019 ; Kipf & Welling , 2017 ; Veličković et al. , 2018 ) have achieved state-of-the-art performance for many tasks on graphs such as recommendation system ( Ying et al. , 2018 ) and traffic prediction ( Guo et al. , 2019 ) . However , most of the GNN models , such as GCN ( Kipf & Welling , 2017 ) and GAT ( Veličković et al. , 2018 ) , learn the node representations by aggregating information over only the 2-hop neighborhood . Such shallow architectures limit their ability to extract information from higher-layer neighborhoods ( Wang & Derr , 2021 ) . But deep GNNs are prone to over-smoothing ( Li et al. , 2018 ) , which suggests the node representations tend to converge to a certain vector and thus become indistinguishable . One solution to address this problem is to preserve the locality of node representations when increasing the number of layers . For example , JKNet ( Xu et al. , 2018 ) densely connects ( Huang et al. , 2017 ) each hidden layer to the final layer . GCNII ( Chen et al. , 2020 ) employs an initial residual to construct a skip connection from the input layer . Besides , Zeng et al . ( 2021 ) pointed out that the key for GNN is to smooth the local neighborhood into informative representation , no matter how deep it is . And they decouple the depth and scope of GNNs to help capture local graph structure . Prior works have emphasized the importance of local information , but one property of the graph is that the number of nodes in the local neighborhood is far fewer than higher-order neighbors . And this property limits the expressive power of GNNs due to the limited neighbors in the local structure . A very intuitive idea is to use data augmentation to increase the number of nodes in the local substructure . However , existing graph data augmentation methods ignore the importance of local information and only perturb at the topology-level and feature-level from a global perspective , which can be divided into two categories : topology-level augmentation ( Rong et al. , 2020 ; Wang et al. , 2020b ; Zhao et al. , 2021 ) and feature-level augmentation ( Deng et al. , 2019 ; Feng et al. , 2019 ; Kong et al. , 2020 ) . Topology-level augmentation perturbs the adjacency matrix , yielding different graph structures . On the other hand , existing feature-level augmentation mainly exploits perturbation of node attributes guided by adversarial training ( Deng et al. , 2019 ; Feng et al. , 2019 ; Kong et al. , 2020 ) . These augmentation techniques have two drawbacks . 1 ) Some of they employ full-batch training for augmentation , which is computationally expensive , and introduce some additional side effects such as over-smoothing . 2 ) The type of feature-level augmentation is coarse-grained , which focuses on global augmentation and overlooks the local information of the neighborhood . Moreover , to our best knowledge , none of the existing approaches combines both the feature representations and the graph topology , especially the local subgraph structures , for graph-level data augmentation . In this work , we propose a framework : Local Augmentation for Graph Neural Networks ( LA-GNNs ) , to further enhance the locality of node representations based on both the topology-level and featurelevel information in the substructure . The term `` local augmentation '' refers to the generation of neighborhood features via a generative model conditioned on local structures and node features . Specifically , our proposed framework learns the conditional distribution of the connected neighbors ’ representations given the representation of the central node , bearing some similarities with the Skipgram ( Mikolov et al. , 2013 ) and Deepwalk Perozzi et al . ( 2014 ) , with the difference that our method does not base on word or graph embedding . The motivation behind this work concludes three-fold . 1 ) Existing feature-level augmentation works primarily pay attention to global augmentation without considering the informative neighborhood . 2 ) The distributions of the representations of the neighbors are closely connected to the central node , making ample room for feature augmentation . 3 ) Preserving the locality of node representations is key to avoiding over-smoothing ( Xu et al. , 2018 ; Klicpera et al. , 2019 ; Chen et al. , 2020 ) . And there are several benefits in applying local augmentation for the GNN training . First , local augmentation is essentially a data augmentation technique that can improve the generalization of the GNN models and prevent over-fitting . Second , we can recover some missing contextual information of the local neighborhood in an attributed graph via the generative model ( Jia & Benson , 2020 ) . Third , our proposed framework is flexible and can be applied to various popular backbone networks such as GCN ( Kipf & Welling , 2017 ) , GAT ( Veličković et al. , 2018 ) , GCNII ( Chen et al. , 2020 ) , and GRAND ( Feng et al. , 2020 ) to enhance their performance . Extensive experimental results demonstrate that our proposed framework could improve the performance of GNN variants on 7 benchmark datasets . 2 BACKGROUND . Notations . Let G = ( V , E ) represent the graph , where V is the set of vertices { v1 , · · · , vN } with |V | = N and E is the set of edges . The adjacency matrix is defined as A ∈ { 0 , 1 } N×N , and Aij = 1 if and only if ( vi , vj ) ∈ E. Let Ni = { vj |Aij = 1 } denotes the neighborhood of node vi and D denote the diagonal degree matrix , where Dii = ∑n j=1 Aij . The feature matrix is denoted as X ∈ RN×F where each node v is associated with a F -dimensional feature vector Xv . Y ∈ { 0 , 1 } N×C denote the one-hot label matrix , where Yi ∈ { 0 , 1 } C is a one-hot vector and∑C j=1 Yij = 1 for any vi ∈ V . GNN . Graph Neural Network ( GNN ) is a type of neural network that directly operates on the graph structure , such as GCN and GAT ( Kipf & Welling , 2017 ; Veličković et al. , 2018 ) , that capture the dependence of graphs via message passing between the nodes of a graph as H ( ` ) = f ( A , H ( ` −1 ) ) , ( 1 ) where f denotes the specific GNN layer for different models , H ( ` ) are the hidden vectors of the ` -th layer and H ( 0 ) = X . For example , f ( A , H ) = σ ( ÂHW ) for GCN , where Â = D̃− 1 2 ÃD̃− 1 2 , D̃ is the degree matrix of Ã , i.e. , D̃ii = ∑ j Ãij , and Ã = A+ I. Topology-level Augmentation . Topology-level augmentation usually perturbs A to generate different graph structures , which can be formulated as A′ = F ( A , X ) , where F ( · ) is a structure perturbation function . For example , DropEdge ( Rong et al. , 2020 ) considers F ( A , X ) = A−As which is independent of X , where As is a sparse matrix consists of a subset of the original edges E. GAUG-O ( Zhao et al. , 2021 ) leverages their proposed neural edge predictors to produce a different structure A′ where A′ij = ⌊ 1 1+e− ( log Pij+G ) /τ + 12 ⌋ , Pij = αMij + ( 1 − α ) Aij , M = σ ( ZZT ) , Z = f ( A , f ( A , X ) ) , τ is the temperature of Gumbel-Softmax distribution , G ∼ Gumbel ( 0 , 1 ) is a Gumbel random variate , and α is a hyperparameter mediating the influence of edge predictor on the original graph . Feature-level Augmentation . Besides , feature-level augmentation function can be defines as X′ = H ( A , X ) , where H ( · ) is a feature perturbation function . FLAG ( Kong et al. , 2020 ) defines the perturbation function as H ( A , X ) = X + δ where perturbation δ is updated iteratively during the adversarial training phase . G-GCN ( plain ) ( Zhu et al. , 2020 ) obtains the global attribute feature matrix X ( a ) ∈ RN×da through minimizing the objective ∏ v∈V ∏ a∈CA ( v ) exp ( X ( a ) v ·Va ) ∑ k∈U exp ( X ( a ) v ·Vk ) where U is the set of all attributes , CA ( v ) is the sampled context attributes of v , and V ∈ Rda×F denotes the parameters . Obviously , the perturbation function of G-GCN has no close-form solution . In this work , we propose a novel feature-level augmentation method , named local augmentation . And the comparison of the details of various graph data augmentation techniques can be found in Table 1 . 3 LOCAL AUGMENTATION . In this section , we describe details of the proposed method . The local augmentation framework consists of three modules : learning the conditional distribution via a generative model , the active learning trick , and the downstream GNN models , as illustrated in Figure 1 . Note that the proposed algorithm enhances the locality of node representations through augmenting 1-hop neighbors in a generative way . Specifically , we exploit a generative model to learn the conditional distribution of the connected neighbors ’ representations given the representation of a node . We describe the details of learning the conditional distribution and the motivation for why local augmentation is able to improve the performance in a probabilistic view in Sec . 3.1 , detail the architecture of downstream GNN models in Sec . 3.2 . We finally elaborate the training procedure of both the generative model and the downstream GNN models with the active learning trick in Sec . 3.3 . 3.1 LEARNING THE CONDITIONAL DISTRIBUTION . We start by reviewing the semi-supervised learning of GNNs in a probabilistic view . Most existing GNN models ( Kipf & Welling , 2017 ; Veličković et al. , 2018 ) are viewed as a classification function to predict the class labels of the graph nodes . In this work , we use a GNN classification estimator Pθ ( Y|A , X ) ( θ is the parameter ) to model the conditional distribution of label Y with respect to the graph structure A and feature matrix X . Given training samples { A , X , Y } , the parameter θ can be estimated using Maximum Likelihood Estimation ( MLE ) , by optimizing the following likelihood function : max ∏ k∈K Pθ ( Yk|A , X ) , ( 2 ) where K is the set of node indices of the training dataset whose labels are visible during the semi-supervised training . To further boost the performance of GNN , we introduce a new model Pθ ( Y , X|A , X ) , where X is generated features by feature-level augmentation . For this model , the MLE method needs to optimize a marginalized probability Pθ over the generated feature matrix X : max ∏ k∈K ∫ X Pθ ( Yk , X|A , X ) . ( 3 ) For Bayesian tractability , we decompose Pθ in Eq . ( 3 ) as a product of two posterior probabilities : Pθ , φ ( Yk , X|A , X ) : = Pθ ( Yk|A , X , X ) Qφ ( X|A , X ) , ( 4 ) where Pθ ( Yk|A , X , X ) and Qφ ( X|A , X ) denote the probabilistic distributions approximated by the downstream GNN and the ( feature-level augmentation ) generator respectively , parameterized by θ and φ . There are two benefits in the decomposition above . First , it allows us to decouple the training of the downstream predictor Pθ and the generator Qφ , enabling the generator to easily generalize to other downstream tasks . Moreover , inspired by the successes of data augmentation via deep-learning-based generative modeling ( Antoniou et al. , 2017 ) , the representation power of Eq . ( 4 ) is superior than that of a single predictor Pθ ( Yk|A , X ) without data augmentation . Consequently , once a generator Qφ is trained very well , our training procedure can optimize Pθ ( Yk|A , X , X ) with samples X drawn from the fixed conditional distribution Qφ . Now , we show how to train the generator as follows . Generator To learn a feature augmentation generator , a naive solution is to learn one single distribution for all the neighbors using the MLE method , i.e. , solving the following optimization problem max ψ ∑ j∈Ni log pψ ( Xj |Xi ) = max ψ log ∏ j∈Ni pψ ( Xj |Xi ) , ( 5 ) where { Xj|j∈Ni , Xi } . Then pψ can be used to augment features for all the neighbors . However , this method ignores the differences between all the neighbors , which may induce severe noise . To overcome the limitation , we assume that each neighbor satisfies a different conditional distribution . Specifically , there exists a conditional distribution p ( ·|Xi , zj ) with latent random variable zj , such that we have Xj ∼ p ( X|Xi , zj ) for Xj|j∈Ni . Once we obtain p ( ·|Xi , zj ) in some way , we can generate augmented features X , and then we can train Pθ ( Yk|A , X , X ) instead of Pθ ( Yk|A , X ) to improve the final performance of Pθ . Below , we will present how to find p ( ·|Xi , zj ) , which will produce the generator Qφ . To achieve our purpose , a suitable method is the conditional variational auto-encoder ( CVAE ) ( Kingma & Welling , 2013 ; Sohn et al. , 2015 ) , which can help learn the distribution of the latent variable zj , and the conditional distribution p ( ·|Xi , zj ) . So , a CVAE model Qφ ( X|A , X ) is adopted as our generator , where φ = { ϕ , ψ } , ϕ denotes the variational parameters and ψ represents the generative parameters . To derive the optimization problem for CVAE , log pψ ( Xj |Xi ) can be written with latent variables z as follows , following previous work ( Pandey & Dukkipati , 2017 ; Sohn et al. , 2015 ) : log pψ ( Xj |Xi ) = ∫ qϕ ( z|Xj , Xi ) log pψ ( Xj , z|Xi ) qϕ ( z|Xj , Xi ) dz+KL ( qϕ ( z|Xj , Xi ) ‖pψ ( z|Xj , Xi ) ) ≥ ∫ qϕ ( z|Xj , Xi ) log pψ ( Xj , z|Xi ) qϕ ( z|Xj , Xi ) dz , and the evidence lower bound ( ELBO ) can be written as : L ( Xj , Xi ; ψ , ϕ ) = −KL ( qϕ ( z|Xj , Xi ) ‖pψ ( z|Xi ) ) + ∫ qϕ ( z|Xj , Xi ) log pψ ( Xj |Xi , z ) dz , ( 6 ) where the encoder qϕ ( z|Xj , Xi ) = N ( f ( Xj , Xi ) , g ( Xj , Xi ) ) and decoder pψ ( Xj |Xi , z ) = N ( h ( Xi , z ) , cI ) . The encoder is a two-layer MLP . f and g share the first layer , and their second layers employ different parameters . The decoder h is two-layer MLP . For simplicity and tractability , the implemented generator Q ( X|A , X ) uses the same parameters across all nodes vi ∈ V . Optimization of the MLE Now , we present how to optimize the MLE Eq . ( 4 ) using the feature matrix produced from the generator . Once the augmented feature matrix can be sampled from the generator , we can optimize the parameters of Eq . ( 4 ) in the following way . Firstly , the parameter φ = { ψ , ϕ } can be optimized by maximizing the ELBO of the generator ( 6 ) , i.e. , we train the generator . Secondly , the parameter θ is optimized by maximizing the MLE Eq . ( 4 ) with φ fixed , which is the conditional distribution of Yk given A , X , and X , i.e. , we train the downstream GNN model . In this paper , the MLE is formulated by a downstream GNN model as follows : Pθ ( Yk | A , X , X ) ∝ −L ( θ|A , X , X , φ ) , ( 7 ) where L ( θ|A , X , X , φ ) = − ∑ k∈T ∑C f=1 Ykf ln ( softmax ( GNN ( A , X , X ) ) kf ) .

Summary: This paper studies the problem of feature augmentation for training the graph neural networks. Specifically, given the central nodes features and local structure, the proposed approach estimates the distribution of the node features of neighbors. Experimental results show that the proposed approach can further improve the baseline backbones.