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 studies the catastrophic forgetting problem when fine-tuning pre-trained language models towards adversarial robustness. They argue that existing robust training methods, such as adversarial training, make the objective model deviated too much from the pre-trained model, resulting in bad generalization. Starting from an information-theoretical perspective, the authors propose RIFT, a fine-tuning method that maximizes the mutual information between the outputs of the objective model and the outputs of the pre-trained model when conditioned on the labels. The experimental results show that RIFT has a better performance than other baselines, which supports the better generalization as the authors claim. In addition, they give explanations why considering I(S;T|Y) in the objective is better than considering I(S;T).   ==== After Rebuttal ====  Thanks for the response. All my questions are properly resolved.