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: Data augmentation remains under-explored on graph-structured data. In this work, the authors introduce local augmentation that enhances the feature of each node by its local subgraph structure. Specifically, they model the data augmentation as a feature generation process: given a node’s feature, they learn the conditional distribution of its neighbors’ features to make a generative augmentation model. Experiments show that the proposed local augmentation makes performance improvement on real-world data.