How Should Pre-Trained Language Models Be Fine-Tuned Towards Adversarial Robustness?

1 Introduction . Deep models are well-known to be vulnerable to adversarial examples [ 64 , 19 , 50 , 35 ] . For instance , fine-tuned models pre-trained on very large corpora can be easily fooled by word substitution attacks using only synonyms [ 2 , 58 , 32 , 12 ] . This has raised grand security challenges to modern NLP systems , such as spam filtering and malware detection , where pre-trained language models like BERT [ 11 ] are widely deployed . Attack algorithms [ 19 , 7 , 76 , 40 , 2 ] aim to maliciously generate adversarial examples to fool a victim model , while adversarial defense aims at building robustness against them . Among the defense methods , adversarial training [ 64 , 19 , 48 , 45 ] is the most effective one [ 3 ] . It updates model parameters using perturbed adversarial samples generated on-the-fly and yields consistently robust performance even against the challenging adaptive attacks [ 3 , 68 ] . However , despite its effectiveness in training from scratch , adversarial training may not directly fit the current NLP paradigm , the fine-tuning of pre-trained language models . First , fine-tuning per se suffers from catastrophic forgetting [ 46 , 18 , 34 ] , i.e. , the resultant model tends to over-fit to a small fine-tuning data set , which may deviate too far from the pre-trained knowledge [ 25 , 81 ] . Second , adversarially fine-tuned models are more likely to forget : adversarial samples are usually 1Our code will be available at https : //github.com/dongxinshuai/RIFT-NeurIPS2021 . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . out-of-distribution [ 38 , 63 ] , thus they are generally inconsistent with the pre-trained model . As a result , adversarial fine-tuning fails to memorize all the robust and generic linguistic features already learned during pre-training [ 65 , 57 ] , which are however very beneficial for a robust objective model . Addressing forgetting is essential for achieving a more robust objective model . Conventional wisdom such as pre-trained weight decay [ 8 , 10 ] and random mixout [ 37 ] mitigates forgetting by constraining lp distance between the two models ’ parameters . Though effective to some extent , however , it is limited because change in the model parameter space only serves as an imperfect proxy for that in the function space [ 5 ] . Besides , the extent to which an encoder fails to retain information also depends on the input distribution . Therefore , a better way to encourage memorization is favorable . In this paper , we follow an information-theoretical lens to look into the forgetting problem : we employ mutual information to measure how well an objective model memorizes the useful features captured before . This motivates our novel adversarial fine-tuning method , Robust Informative FineTuning ( RIFT ) . In addition to fitting a down-stream task as conventions , RIFT maximizes the mutual information between the output of an objective model and that of the pre-trained model conditioned on the class label . It encourages an objective model to continuously retain useful information from the pre-trained one throughout the whole fine-tuning process , whereas the conventional one only makes use of the pre-trained weights for initialization . We illustrate the overall objective of RIFT in Figure 1 and summarize the major contributions of this paper as follows : • To the best of our knowledge , we are the first to investigate the intensified phenomenon of catastrophic forgetting in the adversarial fine-tuning of pre-trained language models . • To address the forgetting problem and achieve more robust models , we propose a novel approach named Robust Informative Fine-Tuning ( RIFT ) from an information-theoretical perspective . RIFT enables an objective model to retain robust and generic linguistic features throughout the whole fine-tuning process and thus enhance robustness against adversarial examples . • We empirically evaluate RIFT on two prevailing NLP tasks , sentiment analysis and natural language inference , where RIFT consistently outperforms the state-of-the-arts in terms of robustness , under different attacks across various pre-trained language models . 2 Methodology . 2.1 Notations and Problem Setting . In this paper , we focus on the text classification task to introduce our method , while it can be easily extended to other NLP tasks . We suppose random variables X , Y ∼ pD ( x , y ) , where X represents the textual input , Y represents the class label , pD is the data distribution , and x , y are the observed values . Our goal is to build a classifier q ( y∣Fs ( x ) ) , where Fs ( ⋅ ) , referred to as our objective model ( a ) In adversarial fine-tuning , the relative L2 distance continuously grows as the fine-tuning proceeds . ( b ) Trading clean accuracy off against robustness ( by increasing β in Eq . 2 ) induces increased relative L2 distance at the last epoch . Figure 2 : Relative L2 distance between the pre-trained model and the objective model in the parameter space , under different fine-tuning schemes on IMDB , with small learning rate 2e-5 . in the rest of this paper , is a feature extractor fine-tuned from the encoder of a pre-trained language model , e.g. , Transformer [ 72 ] , and q ( ⋅∣⋅ ) is parameterized by using an MLP with softmaxed output . We favor a classifier that is robust against adversarial attacks [ 64 , 2 ] , i.e. , maintaining high accuracy even given adversarial examples as inputs . 2.2 Adversarial Fine-Tuning Suffers From Forgetting . Adversarial training [ 64 , 19 , 45 ] is currently the most effective defense technique [ 3 ] . Its training objective can be generally formulated as : min [ E x , y∼pD [ max x̂∈B ( x ) L ( x , x̂ , y ) ] ] , ( 1 ) where B ( x ) defines the allowed perturbation space around x , and L is a loss to encourage correct predictions both given vanilla samples and given perturbed examples . Such an objective function can also be applied to the fine-tuning of a pre-trained language model towards robustness , and we refer to it as adversarial fine-tuning . The loss function in Eq . 1 can be specified following Miyato et al . [ 48 ] , Zhang et al . [ 79 ] in a semi-supervised fashion as : − log q ( y∣Fs ( x ) ) + βKL ( q ( ⋅∣Fs ( x ) ) ∣∣q ( ⋅∣Fs ( x̂ ) ) ) , ( 2 ) where the Kullback–Leibler divergence encourages invariant predictions between x and x̂ . Despite its effectiveness in training from scratch , such an adversarial objective function may not directly fit the fine-tuning scenario . To begin with , fine-tuning itself suffers from catastrophic forgetting [ 46 , 18 , 34 ] . During fine-tuning , an objective model Fs ( ⋅ ) tends to continuously deviate away from the pre-trained one to fit a down-stream task [ 25 , 81 ] , and thus the useful information already captured before is less utilized as the fine-tuning proceeds . Further , adversarial fine-tuning suffers even more from the forgetting problem , the reasons of which are given as what follows . ( i ) Adversarial Fine-Tuning Tends to Forget : Adversarial fine-tuning targets at tackling adversarial examples , which are generally out of the manifold [ 38 , 63 ] of the pre-training corpora . To additionally handle them , an objective model would be fine-tuned towards a solution that is far away from the optimization starting point , i.e. , the pre-trained model . Figure 2 ( b ) empirically shows this point : by increasing β in Eq . 2 , we emphasize more on robustness instead of vanilla accuracy , and consequently , at the last epoch the distance between models also increases . Besides , adversarial fine-tuning often entails more iterations to converge ( several times of normal fine-tuning ) , which further intensifies the forgetting problem , as the objective model is continuously deviating away as shown in Figure 2 ( a ) . ( ii ) Adversarial Fine-Tuning Needs to Memorize : Overfitting to training data is a dominant phenomenon throughout the adversarial training/fine-tuning process [ 60 ] . To mitigate overfitting and generalize better , adversarial fine-tuning should retain all the generalizable features already captured by the pre-trained model [ 11 , 41 ] . In addition , adversarial fine-tuning favors an objective model that extracts features invariant to perturbations [ 70 , 28 ] for robustness . As such , all those generalizable and robust linguistic information captured during pre-training [ 11 , 65 ] are particularly beneficial and should be memorized . We empirically validates that encouraging memorization does improve both generalization and robustness in Sec . 3.5 . To address forgetting , previous methods such as pre-trained weight decay [ 8 , 10 ] and random mixout [ 37 ] have shown their effectiveness in stabilizing fine-tuning [ 81 ] . However , they focus only on the parameter space , in which they encourage an object model to be similar to the pre-trained one . Distance in the parameter space can only approximately characterize the change in the function space , and fails to take the data distribution into consideration . A more natural way to capture the extent to which a model memorizes or forgets , should be using the mutual information between outputs of the two models , and we provide our solution as follows . 2.3 Informative Fine-Tuning . We first use an information-theoretical perspective to look into how a pre-trained model should be leveraged throughout the whole fine-tuning process . We define random variable T = Ft ( X ) as the feature of X extracted by the pre-trained language model Ft ( ⋅ ) , and t as the observed value of T . Similarly , we define random variable S = Fs ( X ) as the feature of X extracted by our objective model Fs ( ⋅ ) , and s as the observed value of S. We formulate an overall fine-tuning objective as follows . The motivation is to train Fs ( ⋅ ) such that S is capable of predicting Y , as well as preserving the information from T . max I ( S ; Y , T ) , ( 3 ) i.e. , maximizing the mutual information between ( i ) the feature extracted by the objective model and ( ii ) the class label plus the feature extracted by the pre-trained model . Since the pre-trained model Ft ( ⋅ ) is a fixed deterministic function , the fine-tuning objective in Eq . 3 in essense encourages Fs ( ⋅ ) to output features that contain as much information as possible for predicting Y and T . This enables the objective model to learn from the pre-trained language model via T throughout the whole fine-tuning process , and thus helps address the forgetting problem . However , the objective defined in Eq . 3 is generally hard to optimize directly . Therefore , we decompose Eq . 3 into two terms as follows : I ( S ; Y , T ) = I ( S ; Y ) + I ( S ; T ∣ Y ) , ( 4 ) where I ( S ; Y ) measures how well the output of our objective model can predict the label , and I ( S ; T ∣ Y ) measures when conditioned on the class label , how well the output features of the two models can predict each other . Visualization of each component can be seen in Figure 1 for more intuitive understandings . We next introduce how each term in the right-hand side of Eq . 4 can be transformed into a tractable lower bound for optimization . ( i ) Maximizing I ( S ; Y ) : We treat q ( y∣s ) , which is the classification layer that takes the features extracted by the objective model as input , as a variational distribution of p ( y∣s ) , and derive a variational lower bound on I ( S ; Y ) following variational inference [ 33 ] as follows : I ( S ; Y ) =H ( Y ) −Ex , y∼pD [ − log q ( y∣s ) ] +KL ( p ( ⋅∣s ) ∣∣q ( ⋅∣s ) ) ( 5 ) ≥H ( Y ) −Ex , y∼pD [ − log q ( y∣s ) ] , ( 6 ) whereH ( Y ) is a constant measuring the shannon entropy of Y , and EX , Y [ − log q ( y∣s ) ] is essentially the cross-entropy loss using q ( y∣s ) for classification . Then , the objective of maximizing I ( S ; Y ) can be achieved by minimizing EX , Y [ − log q ( y∣s ) ] instead . ( ii ) Maximizing I ( S ; T| Y ) : The definition of the conditional mutual information I ( S ; T ∣ Y ) is as follows : I ( S ; T ∣ Y ) = Ey∼pD ( y ) [ I ( S ; T ) ∣ Y = y ] = Ey∼pD ( y ) [ Ex∼pD ( x∣y ) [ log p ( s , t∣y ) p ( s∣y ) p ( t∣y ) ] ] . ( 7 ) To achieve a tractable objective for maximizing Eq . 7 , we employ noise contrastive estimation [ 20 , 49 ] and derive a lower bound −Linfo on the conditional mutual information . This can be summarized in the following Lemma ( the proof of which can be found in Appendix A.2 ) : Lemma 1 . Given { xi , y } Ni=1 that is sampled i.i.d . from pD ( x∣y ) , si = Fs ( xi ) , and ti = Ft ( xi ) , I ( S ; T ∣ Y ) is lower bounded by −Linfo = Ey∼pD ( y ) [ E { xi , y } Ni=1 [ 1 N ∑ N i=1 log efy ( si , ti ) ∑Nj=1 e fy ( si , tj ) + logN ] ] , and fy is a score function indexed by y . By leveraging Lemma 1 , Linfo can be computed using a batch of samples and then minimized for maximizing I ( S ; T ∣ Y ) . The score function fy is defined as the inner product after non-linear projections into a space of hyper-sphere following [ 9 ] as fy ( a , b ) = 1τ ⟨g1y ( a ) , g 2 y ( b ) ⟩ ∥g1y ( a ) ∥2∥g 2 y ( b ) ∥2 , where g1y and g2y are parameterized by using MLPs .

Summary: This paper argues that applying adversarial training into the fine-tuning of the pre-trained language model suffers severely from catastrophic forgetting. It then proposes Robust Informative Fine-Tuning (RIFT) to encourage an objective model to retain the features learned from the pre-trained model. The results on sentiment analysis and natural language inference tasks show that the proposed method is able to outperform the strong baseline models in terms of Genetic and PWWS attacks.