Adversarial Weight Perturbation Improves Generalization in Graph Neural Networks

1 INTRODUCTION . Simply minimizing the standard cross-entropy loss for highly non-convex and non-linear models such as ( deep ) neural networks is not guaranteed to obtain solutions that generalize well , especially for today ’ s overparamatrized networks . The key underlying issue is that these models have many different local minima which can have wildly different generalization properties despite having nearly the same performance on training and validation data . Naturally , there is a rich literature that studies the properties of well-behaving local minima , as well as the design choices that improve our chances of finding them ( Stutz et al. , 2021 ) . The notion of flatness which measure how quickly the loss changes in a neighbourhood around a given local minimum has been empirically shown to correlate with generalization among a variety of different measures ( Jiang et al. , 2019 ) . In addition , generalization bounds based on the PAC-Bayes framework ( McAllester , 1999 ; Foret et al. , 2021 ) provide theoretical insights that corroborate the mounting empirical data . Since the evidence implies that flatter minima tend to generalize better , the obvious question is how to efficiently find them . Not only do flat minima improve generalization from the training to the test data , i.e . the clean accuracy ( Foret et al. , 2021 ; Zheng et al. , 2021 ; Kwon et al. , 2021 ) , but they also improve generalization to adversarial examples , i.e . the robust accuracy ( Wu et al. , 2020a ; Stutz et al. , 2021 ) . Improving adversarial robustness is important , especially for models deployed in safety-critical domain or models deployed in the real-world , since most standard ( undefended ) models are vulnerable to adversarial attacks . Attackers can easily craft deliberate and unnoticeable input perturbations that change the prediction of the classifier . Flat minima show higher resistance to such adversarially perturbed inputs while maintaining good clean accuracy ( Stutz et al. , 2021 ) . Among the variety of techniques for finding flat minima Adversarial Weight Perturbation ( AWP ) ( Wu et al. , 2020a ) , and the closely-related ( adaptive ) sharpness-aware minimization ( Foret et al. , 2021 ; Kwon et al. , 2021 ) and adversarial model perturbation ( Zheng et al. , 2021 ) , seems to be quite effective in practice . The key idea is to minimize the loss w.r.t . a bounded worst-case perturbation of the model parameters , i.e . minimize a local notion of sharpness . The benefits of this approach , and more generally the correlation between flatness and ( clean/robust ) generalization , have been extensively studied for i.i.d . data such as images . In this paper we study this phenomenon for graph data for the first time . Concretely , we analyze and improve the robustness of Graph Neural Networks ( GNNs ) which have become a fundamental building block ( in addition to CNNs and RNNs ) . Blindly applying existing weight perturbation techniques to GNNs is unfortunately not effective in practice due to a vanishing-gradient issue . Intuitively , the adversarially perturbed weights tend to have a higher norm which in turn leads to a saturation in the last layer where that logits for one class are on a significantly larger scale compared to the rest . Even though this limitation plagues all formulations of AWP , for both GNNs and other models ( e.g . ResNets ) , it has gone unnoticed so far . To address it we propose Weighted Truncated Adversarial Weight Perturbation ( WT-AWP ) where rather than directly minimizing the ( robust ) AWP loss we use it as a regularizer in addition to the standard cross-entropy loss . Moreover , we propose to abstain from perturbation in the last layer ( s ) of the network for a more fine-grained control of the training dynamics . These two modifications are simple , but necessary and effective . With our resulting formulation the models can obtain useful gradient signals for training even when the perturbed weights have a high norm , mitigating the gradient-vanishing issue . Furthermore , we theoretically study the AWP learning objective and show its invariance for local extrema . We can summarize our contributions as follows : • We provide a theoretical analysis of AWP on non-i.i.d . graph data and identify a vanishinggradient issue that plagues all previous AWP variants . Based on this analysis we propose Weighted Truncated Adversarial Weight Perturbation ( WT-AWP ) that mitigates this issue . • We study for the first time the connections between flatness and generalization for Graph Neural Networks . We show that GNNs trained with our WT-AWP formulation have simultaneously improved natural and robust generalization . The improvement is statistically significant and consistent across tasks ( node-level and graph-level classification ) and across models ( standard and robustness-aware GNNs ) . All this at a negligible computational cost . 2 BACKGROUND AND RELATED WORK . Adversarial Weight Perturbation for Images . AWP is motivated by the connection between the flatness of the loss landscape and model generalization . Given a learning objective L ( · ) and an image classification model with parameters θ , the generalization gap ( Wu et al. , 2020a ) , also named the sharpness term ( Foret et al. , 2021 ) , which measures the worst-case flatness of the loss landscape , is defined by [ max||δ||≤ρ L ( θ + δ ) − L ( θ ) ] . This gap is known to control a PAC-Bayes generalization bound ( Neyshabur et al. , 2017 ) , with a smaller gap implying better generalization . The AWP objective optimizes the generalization gap and the loss function simultaneously via minθ [ ( max||δ||≤ρ L ( θ + δ ) − L ( θ ) ) + L ( θ ) ] = minθ max||δ||≤ρ L ( θ + δ ) . Providing further theoretical justification for the effectiveness of the AWP , Zheng et al . ( 2021 ) prove that this objective favors solutions corresponding to flatter local minima assuming that the loss surface can be approximated as an inverted Gaussian surface . Relatedly , they show that AWP penalizes the gradient-norm . Keskar et al . ( 2016 ) show that large-batch training may reach sharp minima , however GNNs usually use a small batch size . In some cases we can rescale the weights to achieve arbitrarily sharp minima that also generalize well ( Dinh et al. , 2017 ) . Investigating this issue for GNNs is out of our scope . GNNs , Graph attacks , and Graph defenses . Graph Neural Networks ( GNNs ) are emerging as a fundamental building block . They have achieved spectacular results on a variety of graph learning tasks across many high-impact domains ( see survey ( Wu et al. , 2020b ) ) . Despite their success , it has been demonstrated that GNNs suffer from evasion attacks at test time ( Zügner et al. , 2018 ) and poisoning attacks at training time ( Zügner and Günnemann , 2019 ) . Meanwhile , a series of methods have been developed to improve their robustness . For example , GCNJaccard ( Wu et al. , 2019 ) drops dissimilar edges in the graph , as it found that attackers tend to add edges between nodes with different features . GCNSVD ( Entezari et al. , 2020 ) replaces the adjacency matrix with its low-rank approximation motivated by the observation that mostly the high frequency spectrum of the graph is affected by the adversarial perturbations . We also have provable defenses that provide robustness certificates ( Bojchevski et al. , 2020 ) . Both heuristic defenses ( e.g . GCNJaccard and GCNSVD ) and certificates are improved with our WT-AWP . For an overview of attacks and defenses see Sun et al . ( 2018 ) . 3 ADVERSARIAL WEIGHT PERTURBATION ON GRAPH NEURAL NETWORKS . To simplify the exposition we focus on the semi-supervised node classification task . Nonetheless , in Sec . 5.8 we show that AWP also improves graph-level classification . LetG = ( A , X ) be a given ( attributed ) graph whereA is the adjacency matrix andX contains the node attributes . Let V be the set of all nodes . Normally we optimize minθ Ltrain ( θ ; A , X ) , where Ltrain = ∑ v∈Vtrain l ( fθ ( A , X ) , yv ) , f is a GNN parametrized by weights θ = ( θ1 , ... , θk ) , yv is the ground-truth label for node v , and l is some loss function ( e.g . cross-entropy ) applied to each node in the training set Vtrain ⊂ V . In AWP we first find the worst-case weight perturbation δ∗ ( θ ) that maximizes the loss . Then we minimize the loss with the perturbed weights . The worst-case perturbation for a given θ is defined as δ∗ ( θ ) : = arg max { δ| ||δi||2≤ρ ( θi ) , i∈ [ k ] } Ltrain ( θ + δ ; A , X ) , ( 1 ) where ρ ( θ ) is the strength of perturbation . Then the AWP learning objective is min θ max { δ| ||δi||2≤ρ ( θi ) , i∈ [ k ] } Ltrain ( θ + δ ; A , X ) = min θ Ltrain ( θ + δ ∗ ( θ ) ; A , X ) . ( 2 ) Since the PAC-Bayes bound proposed by McAllester ( 1999 ) only holds for i.i.d . training tasks and semi-supervised node classification is a non-i.i.d . task , analysis in Wu et al . ( 2020a ) and Foret et al . ( 2021 ) could not be naturally extended to node classification tasks . We derived a new generalization bound for node classification tasks on GNNs based on the sub-group generalization ( Ma et al. , 2021 ) . Theorem 1 . ( informal ) . Let Lall ( θ ; A , X ) be the loss on all nodes , including the unseen test nodes . It is bounded by the adversarially weight perturbed loss on the training nodes as follows : Lall ( θ ; A , X ) ≤ max { δ| ||δi||2≤ρ ( θi ) , i∈ [ k ] } Ltrain ( θ + δ ; A , X ) + h ( ||θ||22/ρ ( θ ) 2 ) ( 3 ) The formal version , the details for h ( ) , and the proof are in Appendix E. This bound justifies the use of AWP since the perturbed loss on training nodes bounds the standard loss on all nodes . Moreover , as h ( ||θ||22/ρ ( θ ) 2 ) is monotonously decreasing with ρ ( θ ) , increasing the perturbation strength ρ can make the bound in Eq . 3 sharper , i.e . the resulting AWP objective should lead to better generalization . Since finding the optimal perturbation ( Eq . 1 ) is intractable , we approximate it with a one-step projected gradient descent as in previous work ( Wu et al. , 2020a ; Foret et al. , 2021 ; Zheng et al. , 2021 ) , δ̂∗ ( θ ) : = ΠB ( ρ ( θ ) ) ( ∇θLtrain ( θ ; A , X ) ) , ( 4 ) where B ( ρ ( θ ) ) is an l2 ball with radius ρ ( θ ) and ΠB ( ρ ( θ ) ) ( · ) is a projection operation , which projects the perturbation back to the surface ofB ( ρ ( θ ) ) when the perturbation is out of the ball . The maximum perturbation norm ρ ( θ ) could either be a constant ( Foret et al. , 2021 ; Zheng et al. , 2021 ) or layer dependent ( Wu et al. , 2020a ) . We specify a layer-dependent norm constraint ρ ( θ ) : = ρ||θ||2 because the scales of different layers in a neural network vary greatly . With the approximation δ̂∗ ( θ ) , the definition of the final AWP learning objective is given by min θ Lawp ( θ ) : = Ltrain ( θ + ΠB ( ρ ( θ ) ) ( ∇θLtrain ( θ ; A , X ) ) ; A , X ) , ( 5 ) If Ltrain ( θ ; A , X ) is smooth enough , ∇θLtrain ( θ ; A , X ) = 0 when θ∗ is a local extremum . In this case Lawp ( θ ) = Ltrain ( θ ; A , X ) . A natural question is whether θ∗ will also be the extremum of Lawp ( θ ) ? We show that Lawp ( θ ) keeps the local extremum of Ltrain ( θ ; A , X ) unchanged . Theorem 2 . ( Invariant of local minimum and maximum ) With the AWP learning objective in Eq . 5 , and for continuous Ltrain ( θ ; A , X ) , ∇θLtrain ( θ ; A , X ) , ∆θLtrain ( θ ; A , X ) , if θ∗ is a local minimum of Ltrain ( θ ; A , X ) and the Hessian matrix ∆θLtrain ( θ ; A , X ) |θ∗ is positive definite , θ∗ is also a local minimum of Lawp ( θ ) . The proof is provided in Appendix A . The exact gradient of this new objective is ∇θLtrain ( θ + δ̂∗ ( θ ) ; A , X ) = ∇θLtrain ( θ ; A , X ) |θ+δ̂∗ ( θ ) +∇θδ̂ ∗ ( θ ) ∇θLtrain ( θ ; A , X ) |θ+δ̂∗ ( θ ) ( 6 ) Since∇θδ̂∗ ( θ ) includes second and higher order derivative of θ , which are computationally expensive , they are omitted during training , obtaining the following approximate gradient of the AWP loss ∇θLtrain ( θ ; A , X ) |θ+δ̂∗ ( θ ) ( 7 ) Foret et al . ( 2021 ) show the models trained with the exact gradient ( Eq . 6 ) have almost the same performance as model trained with the estimated first-order gradient ( Eq . 7 ) .

Summary: The authors extend the line of work related to adversarial weight perturbations: first showing that a vanishing gradient issue in standard AWP can hinder training. To remedy this, some natural tweaks are applied to AWP. The authors then focus on using AWP to train graph neural networks, demonstrating a minor boost in both clean and robust accuracy.