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: In this paper, the authors present a new adversarial training algorithm and apply it to the fintuning stage large scale language models BERT and RoBERTa. They find that with FreeLB applied to finetuning, both BERT and RoBERTa see small boosts in performance on GLUE, ARC, and CommonsenseQA. The gains they see on GLUE are quite small (0.3 on the GLUE test score for RoBERTa) but the gains are more substantial on ARC and CommonsenseQA. The paper also presents some ablation studies on the use of the same dropout mask across each ascent step of FreeLB, empirically seeing gains by using the same mask. They also present some analysis on robustness in the embedding space, showing that FreeLB leads to greater robustness than other adversarial training methods