Better Fine-Tuning by Reducing Representational Collapse

1 INTRODUCTION . Pre-trained language models ( Radford et al. , 2019 ; Devlin et al. , 2018 ; Liu et al. , 2019 ; Lewis et al. , 2019 ; 2020 ) have been shown to capture a wide array of semantic , syntactic , and world knowledge ( Clark et al. , 2019 ) , and provide the defacto initialization for modeling most existing NLP tasks . However , fine-tuning them for each task is a highly unstable process , with many hyperparameter settings producing failed fine-tuning runs , unstable results ( considerable variation between random seeds ) , over-fitting , and other unwanted consequences ( Zhang et al. , 2020 ; Dodge et al. , 2020 ) . Recently , trust region or adversarial based approaches , including SMART ( Jiang et al. , 2019 ) and FreeLB ( Zhu et al. , 2019 ) , have been shown to increase the stability and accuracy of fine-tuning by adding additional constraints limiting how much the fine-tuning changes the initial parameters . However , these methods are significantly more computationally and memory intensive than the more commonly adopted simple-gradient-based approaches . This paper presents a lightweight fine-tuning strategy that matches or improves performance relative to SMART and FreeLB while needing just a fraction of the computational and memory overhead and no additional backward passes . Our approach is motivated by trust region theory while also reducing to simply regularizing the model relative to parametric noise applied to the original pre-trained representations . We show uniformly better performance , setting a new state of the art for RoBERTa fine-tuning on GLUE and reaching state of the art on XNLI using no novel pre-training approaches ( Liu et al. , 2019 ; Wang et al. , 2018 ; Conneau et al. , 2018 ) . Furthermore , the low overhead of our family of fine-tuning methods allows our method to be applied to generation tasks where we consistently outperform standard fine-tuning , setting state of the art on summarization tasks . We also introduce a new analysis to motivate the use of trust-region-style methods more generally , by defining a new notion of representational collapse and introducing a new methodology for measuring it during fine-tuning . Representational collapse is the degradation of generalizable representations of pre-trained models during the fine-tuning stage . We empirically show that standard fine-tuning degrades generalizable representations through a series of probing experiments on GLUE tasks . Furthermore , we attribute this phenomenon to using standard gradient descent algorithms for the fine-tuning stage . We also find that ( 1 ) recently proposed fine-tuning methods rooted in trust region , i.e. , SMART , can alleviate representation collapse , and ( 2 ) our methods alleviate representational collapse to an even greater degree , manifesting in better performance across almost all datasets and models . Our contributions in this paper are the following . • We propose a novel approach to fine-tuning rooted in trust-region theory , which we show directly alleviates representational collapse at a fraction of the cost of other recently proposed fine-tuning methods . • Through extensive experimentation , we show that our method outperforms standard finetuning methodology following recently proposed best practices from Zhang et al . ( 2020 ) . We improve various SOTA models from sentence prediction to summarization , from monolingual to cross-lingual . • We further define and explore the phenomena of representational collapse in fine-tuning and directly correlate it with generalization in tasks of interest . 2 LEARNING ROBUST REPRESENTATIONS THROUGH REGULARIZED FINE-TUNING . We are interested in deriving methods for fine-tuning representations that provide guarantees on the movement of representations , in the sense that they do not forget the original pre-trained representations when they are fine-tuned for new tasks ( see Section 4 for more details ) . We introduce a new fine-tuning method rooted in an approximation to trust region , which provides guarantees for stochastic gradient descent algorithms by bounding some divergence between model at update t and t+ 1 ( Pascanu & Bengio , 2013 ; Schulman et al. , 2015b ; Jiang et al. , 2019 ) . Let f : Rm×n → Rp be a function which returns some pre-trained representation parameterized by θf from m tokens embedded into a fixed vector of size n. Let the learned classification head g : Rp → Rq be a function which takes an input from f and outputs a valid probability distribution parameterized by θg in q dimensions and let X be our dataset . In the case of generation , we can assume the classification head is simply an identity function or softmax depending on the loss function . Let L ( θ ) denote a loss function given by θ = [ θf , θg ] . We are interested in minimizing L with respect to θ such that each update step is constrained by movement in the representational density space p ( f ) . More formally given an arbitrary arg min ∆θ L ( θ + ∆θ ) s.t . KL ( p ( f ( · ; θf ) ) ||p ( f ( · ; θf + ∆θf ) ) ) = ( 1 ) This constrained optimization problem is equivalent to doing natural gradient descent directly over the representations ( Pascanu & Bengio , 2013 ) . Unfortunately , we do not have direct access to the density of representations ; therefore , it is not trivial to directly bound this quantity . Instead , we propose to do natural gradient over g · f with an additional constraint that g is at most 1-Lipschitz ( which naturally constrains change of representations , see Section A.1 in the Appendix ) . Traditional computation of natural gradient is computationally prohibitive due to the need for inverting the Hessian . An alternative formulation of natural gradient can be stated through mirror descent , using Bregmann divergences ( Raskutti & Mukherjee , 2015 ; Jiang et al. , 2019 ) . This method primarily serves as a robust regularizer by preventing large updates in the model ’ s probability space . This family of methods is classically known as trust-region methods ( Pascanu & Bengio , 2013 ; Schulman et al. , 2015a ) . LSMART ( θ , f , g ) = L ( θ ) + λEx∼X [ sup x∼ : |x∼−x|≤ KLS ( g · f ( x ) ‖ g · f ( x∼ ) ) ] ( 2 ) However , the supremum is computationally intractable . An approximation is possible by doing gradient ascent steps , similar to finding adversarial examples . This was first proposed by SMART with a symmetrical KLS ( X , Y ) = KL ( X||Y ) +KL ( Y ||X ) term ( Jiang et al. , 2019 ) . We propose an even simpler approximation which does not require extra backward computations and empirically works as well as or better than SMART . We altogether remove the adversarial nature from SMART and instead optimize for a smoothness parameterized by KLS . Furthermore , we optionally also add a constraint on the smoothness of g by making it at most 1-Lipschitz , the intuition being if we can bound the volume of change in g we can more effectively bound f . LR3 ( f , g , θ ) = L ( θ ) + λEx∼X [ KLS ( g · f ( x ) ‖ g · f ( x+ z ) ) ] R3F Method ( 3 ) s.t . z ∼ N ( 0 , σ2I ) or z ∼ U ( −σ , σ ) ( 4 ) s.t . Lip { g } ≤ 1 Optional R4F Method ( 5 ) where KLS is the symmetric KL divergence and z is a sample from a parametric distribution . In our work we test against two distributions , normal and uniform centered around 0 . We denote this as the Robust Representations through Regularized Finetuning ( R3F ) method . Additionally we propose an extension to R3F ( R4F ; Robust Representations through Regularized and Reparameterized Finetuning , which reparameterizes g to be at most 1-Lipschitz via Spectral Normalization ( Miyato et al. , 2018 ) . By constraining g to be at most 1-Lipschitz , we can more directly bound the change in representation ( Appendix Section A.1 ) . Specifically we scale all the weight matrices of g by the inverse of their largest singular values WSN : = W/σ ( W ) . Given that spectral radius σ ( WSN ) = 1 we can bound Lip { g } ≤ 1 . In the case of generation , g does not have any weights therefore we can only apply the R3F method . 2.1 RELATIONSHIP TO SMART AND FREELB . Our method is most closely related to the SMART algorithm , which utilizes an auxiliary smoothness inducing regularization term , which directly optimizes the Bregmann divergence mentioned above in Equation 2 ( Jiang et al. , 2019 ) . SMART solves the supremum by using an adversarial methodology to ascent to the largest KL divergence with an −ball . We instead propose to remove the ascent step completely , optionally fixing the smoothness of the classification head g. This completely removes SMART ’ s adversarial nature and is more akin to optimizing the smoothness of g · f directly . Another recently proposed adversarial method for fine-tuning , FreeLB optimizes a direct adversarial loss LFreeLB ( θ ) = sup∆θ : |∆θ|≤ L ( θ + ∆θ ) through iterative gradient ascent steps . This is similar to SMART in the sense that both are adversarial and require gradient ascent steps . Unfortunately , the need for extra forward-backward passes can be prohibitively expensive when fine-tuning large pre-trained models ( Zhu et al. , 2019 ) . Our method is significantly more computationally efficient than adversarial based fine-tuning methods , as seen in Table 1 . We show that this efficiency does not hurt performance ; we can match or exceed FreeLB and SMART on a large number of tasks . In addition , the relatively low costs of our methods allow us to improve over fine-tuning on an array of generation tasks . 3 EXPERIMENTS . We will first measure performance by fine-tuning on a range of tasks and languages . The next sections report why methods rooted in trust region , including ours , outperform standard fine-tuning . We aimed for fair comparisons throughout all of our experiments by using fixed budget hyperparameters searches across all methods . Furthermore , for computationally tractable tasks , we report median/max numbers as well as show distributions across a large number of runs . 3.1 SENTENCE PREDICTION . GLUE . We will first test R3F and R4F on sentence classification tasks from the GLUE benchmark ( Wang et al. , 2018 ) . We select the same subset of GLUE tasks that have been reported by prior work in this space ( Jiang et al. , 2019 ) : MNLI ( Williams et al. , 2018 ) , QQP ( Iyer et al. , 2017 ) , RTE ( Bentivogli et al. , 2009 ) , QNLI ( Rajpurkar et al. , 2016 ) , MRPC ( Dolan & Brockett , 2005 ) , CoLA ( Warstadt et al. , 2018 ) , SST-2 ( Socher et al. , 2013 ) .1 Consistent with prior work ( Jiang et al. , 2019 ; Zhu et al. , 2019 ) , we focus on improving the performance of RoBERTa-Large based models in the single-task setting ( Liu et al. , 2019 ) . We report the performance of all models on the GLUE development set . Method 5000 6000 7000 8000 9000 Se co nd s SST-2 Walltime Analysis Finetuning Method Standard++ SMART R3F R3F R4F R4F ods across ten seeds to demonstrate the stability properties of individual methods in Figure 2 . R3F and R4F unanimously improve over Standard and Standard++ fine-tuning . Furthermore , our methods match or exceed adversarial methods such as SMART/FreeLB at a fraction of the computational cost when comparing median runs . We show computational cost in Figure 1 for a single task , but the relative behavior of wall times is consistent across all other GLUE tasks . We note that we could not find a discernable difference in the experimental setting , which would make the selection between R3F vs. R4F trivial . 1We do not test against STS-B because it is a regression task where our KL divergence is not defined ( Cer et al. , 2017 ) .

Summary: The paper proposes a method for finetuning pre-trained models that ensures the generalization ability of the representation is maintained. The key innovation is that the computationally expensive ascent step in the mirror descent method of SMART can be replaced by simply injecting noise. The results support the hypothesis that this works well for keeping the generalization-ability of the model. The authors also define the degradation of the generalizability of the representation during finetuning as “representational collapse”.