FreeLB: Enhanced Adversarial Training for Natural Language Understanding

1 INTRODUCTION . Adversarial training is a method for creating robust neural networks . During adversarial training , mini-batches of training samples are contaminated with adversarial perturbations ( alterations that are small and yet cause misclassification ) , and then used to update network parameters until the resulting model learns to resist such attacks . Adversarial training was originally proposed as a means to enhance the security of machine learning systems ( Goodfellow et al. , 2015 ) , especially for safety-critical systems like self-driving cars ( Xiao et al. , 2018 ) and copyright detection ( Saadatpanah et al. , 2019 ) . In this paper , we turn our focus away from the security benefits of adversarial training , and instead study its effects on generalization . While adversarial training boosts the robustness , it is widely accepted by computer vision researchers that it is at odds with generalization , with classification accuracy on non-corrupted images dropping as much as 10 % on CIFAR-10 , and 15 % on Imagenet ( Madry et al. , 2018 ; Xie et al. , 2019 ) . Surprisingly , people observe the opposite result for language models ( Miyato et al. , 2017 ; Cheng et al. , 2019 ) , showing that adversarial training can improve both generalization and robustness . We will show that adversarial training significantly improves performance of state-of-the-art models for many language understanding tasks . In particular , we propose a novel adversarial training algorithm , called FreeLB ( Free Large-Batch ) , which adds adversarial perturbations to word embeddings and minimizes the resultant adversarial loss around input samples . The method leverages recently proposed “ free ” training strategies ( Shafahi et al. , 2019 ; Zhang et al. , 2019 ) to enrich the training data with diversified adversarial samples under different norm constraints at no extra cost than PGD-based ( Projected Gradient Descent ) adversarial training ( Madry et al. , 2018 ) , which enables us to perform such diversified adversarial training on large-scale state-of-the-art models . We observe improved invariance in the embedding space for models trained with FreeLB , which is positively correlated with generalization . 1Code is available at https : //github.com/zhuchen03/FreeLB . We perform comprehensive experiments to evaluate the performance of a variety of adversarial training algorithms on state-of-the-art language understanding models and tasks . In the comparisons with standard PGD ( Madry et al. , 2018 ) , FreeAT ( Shafahi et al. , 2019 ) and YOPO ( Zhang et al. , 2019 ) , FreeLB stands out to be the best for the datasets and models we evaluated . With FreeLB , we achieve state-of-the-art results on several important language understanding benchmarks . On the GLUE benchmark , FreeLB pushes the performance of the BERT-base model from 78.3 to 79.4 . The overall score of the RoBERTa-large models on the GLUE benchmark is also lifted from 88.5 to 88.8 , achieving best results on most of its sub-tasks . Experiments also show that FreeLB can boost the performance of RoBERTa-large on question answering tasks , such as the ARC and CommonsenseQA benchmarks . We also provide a comprehensive ablation study and analysis to demonstrate the effectiveness of our training process . 2 RELATED WORK . 2.1 ADVERSARIAL TRAINING . To improve the robustness of neural networks against adversarial examples , many defense strategies and models have been proposed , in which PGD-based adversarial training ( Madry et al. , 2018 ) is widely considered to be the most effective , since it largely avoids the the obfuscated gradient problem ( Athalye et al. , 2018 ) . It formulates a class of adversarial training algorithms ( Kurakin et al. , 2017 ) into solving a minimax problem on the cross-entropy loss , which can be achieved reliably through multiple projected gradient ascent steps followed by a SGD ( Stochastic Gradient Descent ) step . Despite being verified by Athalye et al . ( 2018 ) to avoid obfuscated gradients , Qin et al . ( 2019 ) shows that PGD-based adversarial training still leads to highly convolved and non-linear loss surfaces when K is small , which could be readily broken under stronger adversaries . Thus , to be effective , the cost of PGD-based adversarial training is much higher than conventional training . To mitigate this cost , Shafahi et al . ( 2019 ) proposed a “ free ” adversarial training algorithm that simultaneously updates both model parameters and adversarial perturbations on a single backward pass . Using a similar formulation , Zhang et al . ( 2019 ) effectively reduce the total number of full forward and backward propagations for obtaining adversarial examples by restricting most of its adversarial updates in the first layer . 2.2 ADVERSARIAL EXAMPLES IN NATURAL LANGUAGES . Adversarial examples have been explored primarily in the image domain , and received many attention in text domain recently . Previous works on text adversaries have focused on heuristics for creating adversarial examples in the black-box setting , or on specific tasks . Jia & Liang ( 2017 ) propose to add distracting sentences to the input document in order to induce mis-classification . Zhao et al . ( 2018 ) generate text adversaries by projecting the input data to a latent space using GANs , and searching for adversaries close to the original instance . Belinkov & Bisk ( 2018 ) manipulate every word in a sentence with synthetic or natural noise in machine translation systems . Iyyer et al . ( 2018 ) propose a neural paraphrase model based on back-translated data to produce paraphrases that have different sentence structures . Different from previous work , ours is not to produce actual adversarial examples , but only take the benefit of adversarial training for natural language understanding . We are not the first to observe that robust language models may perform better on clean test data . Miyato et al . ( 2017 ) extend adversarial and virtual adversarial training ( Miyato et al. , 2019 ) to the text domain to improve the performance on semi-supervised classification tasks . Ebrahimi et al . ( 2018 ) propose a character/word replacement for crafting attacks , and show employing adversarial examples in training renders the models more robust . Ribeiro et al . ( 2018 ) show that adversarial attacks can be used as a valuable tool for debugging NLP models . Cheng et al . ( 2019 ) also find that crafting adversarial examples can help neural machine translation significantly . Notably , these studies have focused on simple models or text generation tasks . Our work explores how to efficiently use the gradients obtained in adversarial training to boost the performance of state-of-the-art transformer-based models . 3 ADVERSARIAL TRAINING FOR LANGUAGE UNDERSTANDING . Pre-trained large-scale language models , such as BERT ( Devlin et al. , 2019 ) , RoBERTa ( Liu et al. , 2019b ) , ALBERT ( Lan et al. , 2020 ) and T5 ( Raffel et al. , 2019 ) , have proven to be highly effective for downstream tasks . We aim to further improve the generalization of these pre-trained language models on the downstream language understanding tasks by enhancing their robustness in the embedding space during finetuning on these tasks . We achieve this goal by creating “ virtual ” adversarial examples in the embedding space , and then perform parameter updates on these adversarial embeddings . Creating actual adversarial examples for language is difficult ; even with state-of-theart language models as guidance ( e.g. , ( Cheng et al. , 2019 ) ) , it remains unclear how to construct label-preserving adversarial examples via word/character replacement without human evaluations , because the meaning of each word/character depends on the context ( Ribeiro et al. , 2018 ) . Since we are only interested in the effects of adversarial training , rather than producing actual adversarial examples , we add norm-bounded adversarial perturbations to the embeddings of the input sentences using a gradient-based method . Note that our embedding-based adversary is strictly stronger than a more conventional text-based adversary , as our adversary can make manipulations on word embeddings that are not possible in the text domain . For models that incorporate various input representations , including word or subword embeddings , segment embeddings and position embeddings , our adversaries only modify the concatenated word or sub-word embeddings , leaving other components of the sentence representation unchanged . 2 Denote the sequence of one-hot representations of the input subwords as Z = [ z1 , z2 , ... , zn ] , the embedding matrix as V , and the language model ( encoder ) as a function y = fθ ( X ) , whereX = V Z is the subword embeddings , y is the output of the model ( e.g. , class probabilities for classification models ) , and θ denotes all the learnable parameters including the embedding matrix V . We add adversarial perturbations δ to the embeddings such that the prediction becomes y′ = fθ ( X + δ ) . To preserve the semantics , we constrain the norm of δ to be small , and assume the model ’ s prediction should not change after the perturbation . This formulation is analogous to Miyato et al . ( 2017 ) , with the difference that we do not requireX to be normalized . 3.1 PGD FOR ADVERSARIAL TRAINING . Standard adversarial training seeks to find optimal parameters θ∗ to minimize the maximum risk for any δ within a norm ball as : min θ E ( Z , y ) ∼D [ max ‖δ‖≤ L ( fθ ( X + δ ) , y ) ] , ( 1 ) where D is the data distribution , y is the label , and L is some loss function . We use the Frobenius norm to constrain δ . For neural networks , the outer “ min ” is non-convex , and the inner “ max ” is non-concave . Nonetheless , Madry et al . ( 2018 ) demonstrated that this saddle-point problem can be solved reliably with SGD for the outer minimization and PGD ( a standard method for large-scale constrained optimization , see ( Combettes & Pesquet , 2011 ) and ( Goldstein et al. , 2014 ) ) , for the inner maximization . In particular , for the constraint ‖δ‖F ≤ , with an additional assumption that the loss function is locally linear , PGD takes the following step ( with step size α ) in each iteration : δt+1 = Π‖δ‖F≤ ( δt + αg ( δt ) /‖g ( δt ) ‖F ) , ( 2 ) where g ( δt ) = ∇δL ( fθ ( X + δt ) , y ) is the gradient of the loss with respect to δ , and Π‖δ‖F≤ performs a projection onto the -ball . To achieve high-level robustness , multi-step adversarial examples are needed during training , which is computationally expensive . TheK-step PGD ( K-PGD ) requires K forward-backward passes through the network , while the standard SGD update requires only one . As a result , the adversary generation step in adversarial training increases run-time by an order of magnitude—a catastrophic amount when training large state-of-the-art language models . 3.2 LARGE-BATCH ADVERSARIAL TRAINING FOR FREE . In the inner ascent steps of PGD , the gradients of the parameters can be obtained with almost no overhead when computing the gradients of the inputs . From this observation , FreeAT ( Shafahi et al. , 2 “ Subword embeddings ” refers to the embeddings of sub-word encodings such as the popular Byte Pair Encoding ( BPE ) ( Sennrich et al. , 2016 ) . Algorithm 1 “ Free ” Large-Batch Adversarial Training ( FreeLB-K ) Require : Training samples X = { ( Z , y ) } , perturbation bound , learning rate τ , ascent steps K , ascent step size α 1 : Initialize θ 2 : for epoch = 1 . . . Nep do 3 : for minibatch B ⊂ X do 4 : δ0 ← 1√NδU ( − , ) 5 : g0 ← 0 6 : for t = 1 . . .K do 7 : Accumulate gradient of parameters θ 8 : gt ← gt−1 + 1KE ( Z , y ) ∈B [ ∇θ L ( fθ ( X + δt−1 ) , y ) ] 9 : Update the perturbation δ via gradient ascend 10 : gadv ← ∇δ L ( fθ ( X + δt−1 ) , y ) 11 : δt ← Π‖δ‖F≤ ( δt−1 + α · gadv/‖gadv‖F ) 12 : end for 13 : θ ← θ − τgK 14 : end for 15 : end for 2019 ) and YOPO ( Zhang et al. , 2019 ) have been proposed to accelerate adversarial training . They achieve comparable robustness and generalization as standard PGD-trained models using only the same or a slightly larger number of forward-backward passes as natural training ( i.e. , SGD on clean samples ) . FreeAT takes one descent step on the parameters together with each of the K ascent steps on the perturbation . As a result , FreeAT may suffer from the “ stale gradient ” problem ( Dutta et al. , 2018 ) , where in every step t , δt does not necessarily maximize the model with parameter θt since its update is based on ∇δL ( fθt−1 ( X + δt−1 ) , y ) , and vice versa , θt does not necessarily minimize the adversarial risk with adversary δt since its update is based on ∇θL ( fθt−1 ( X + δt−1 ) , y ) . Such a problem may be more significant when the step size is large . Different from FreeAT , YOPO accumulates the gradient of the parameters from each of the ascent steps , and updates the parameters only once after the K inner ascent steps . YOPO also advocates that after each back-propagation , one should take the gradient of the first hidden layer as a constant and perform several additional updates on the adversary using the product of this constant and the Jacobian of the first layer of the network to obtain strong adversaries . However , when the first hidden layer is a linear layer as in their implementation , such an operation is equivalent to taking a larger step size on the adversary . The analysis backing the extra update steps also assumes a twice continuously differentiable loss , which does not hold for ReLU-based neural networks they experimented with , and thus the reasons for the success of such an algorithm remains obscure . We give empirical comparisons between YOPO and our approach in Sec . 4.3 . To obtain better solutions for the inner max and avoid fundamental limitations on the function class , we propose FreeLB , which performs multiple PGD iterations to craft adversarial examples , and simultaneously accumulates the “ free ” parameter gradients ∇θL in each iteration . After that , it updates the model parameter θ all at once with the accumulated gradients . The overall procedure is shown in Algorithm 1 , in which X + δt is an approximation to the local maximum within the intersection of two balls It = BX+δ0 ( αt ) ∩ BX ( ) . By taking a descent step along the averaged gradients atX + δ0 , ... , X + δK−1 , we approximately optimize the following objective : min θ E ( Z , y ) ∼D [ 1 K K−1∑ t=0 max δt∈It L ( fθ ( X + δt ) , y ) ] , ( 3 ) which is equivalent to replacing the original batchX with a K-times larger virtual batch , consisting of samples whose embeddings are X + δ0 , ... , X + δK−1 . Compared with PGD-based adversarial training ( Eq . 1 ) , which minimizes the maximum risk at a single estimated point in the vicinity of each training sample , FreeLB minimizes the maximum risk at each ascent step at almost no overhead . Intuitively , FreeLB could be a learning method with lower generalization error than PGD . Sokolic et al . ( 2017 ) have proved that the generalization error of a learning method invariant to a set of T transformations may be up to √ T smaller than a non-invariant learning method . According to their theory , FreeLB could have a more significant improvement over natural training , since FreeLB enforces the invariance to K adversaries from a set of up to K different norm constraints,3 while PGD only enforces invariance to a single norm constraint . Empirically , FreeLB does lead to higher robustness and invariance than PGD in the embedding space , in the sense that the maximum increase of loss in the vicinity of X for models trained with FreeLB is smaller than that with PGD . See Sec . 4.3 for details . In theory , such improved robustness can lead to better generalization ( Xu & Mannor , 2012 ) , which is consistent with our experiments . Qin et al . ( 2019 ) also demonstrated that PGD-based method leads to highly convolved and non-linear loss surfaces in the vicinity of input samples when K is small , indicating a lack of robustness .

Summary: -	This paper modifies and extends the recent “free” training strategies in adversarial training for representation learning for natural language.  The proposed “Free” Large-Batch Adversarial Training is well motived, in comparison with plain PGD-based adversarial training and the existing methods like FreeAT and YOPO, which virtually enlarges the batch size and minimize maximum risk at every ascent step. The contributions are solid.