Representation Projection Invariance Mitigates Representation Collapse (2023)

Abstract: Fine-tuning contextualized representations learned by pre-trained language models remains a prevalent practice in NLP. However, fine-tuning can lead to representation degradation (also known as representation collapse), which may result in instability, sub-optimal performance, and weak generalization. In this paper, we propose Representation Projection Invariance (REPINA), a novel regularization method to maintain the information content of representation and reduce representation collapse during fine-tuning by discouraging undesirable changes in the representations. We study the empirical behavior of the proposed regularization in comparison to 5 comparable baselines across 13 language understanding tasks (GLUE benchmark and six additional datasets). When evaluating in-domain performance, REPINA consistently outperforms other baselines on most tasks (10 out of 13). Additionally, REPINA improves out-of-distribution performance. We also demonstrate its effectiveness in few-shot settings and robustness to label perturbation. As a by-product, we extend previous studies of representation collapse and propose several metrics to quantify it. Our empirical findings show that our approach is significantly more effective at mitigating representation collapse.

## introduction fine-tuning pre-trained language models has been shown to achieve remarkable performance on a variety of natural language processing (nlp) tasks @xcite @xcite . a standard fine-tuning strategy involves adapting the pre-trained model to a supervised downstream task (fig 1 ; left). such a procedure can result in representation collapse @xcite , distorting the pre-trained representations that limits 1 code is available at https://github.com/arazd/repina . figure 1 : fine-tuning the whole architecture (left) generally leads to good performance though it distorts the pre-trained representations. minimal tuning (of classification head, for example; middle) mitigates the representation collapse but limits the model performance. our proposal repina (right) leads to good performance while mitigating the representation collapse. we introduce representation projection invariance (repina), a regularization objective that prevents undesirable changes in the representations (fig 2a ). our regularization applies an invariance loss on a tunable projection of the representation. this regularization allows the underlying representation to change mildly (e.g., shift and scaling) while not losing its expressivity (fig 2b ). our regularization objective provides a knob that controls the amount of loss-free transformations allowed during fine-tuning. we compare our method against several established regularization approaches which explicitly or implicitly address the issue of representation degradation (section 5.1). we show that our approach consistently outperforms major fine-tuning methods across 13 tasks (fig 3 ; left) and improves generalizability on out-of-distribution (ood) data (section 5.4). our approach is particularly effective in scenarios where data is limited (such as with only 250, 500, or 1000 examples), as the model is more likely to overfit and memorize the training data in these cases (section 5.6). furthermore, we thoroughly investigate fine-tuning under label perturbation (from 5% to 30% label noise) and observe that our approach is robust to incorrect labels, exceeding the performance of standard finetuning procedure and common baseline methods (section 5.5). finally, we quantify how much different methods mitigate the degradation of representations (section 6). we use previously explored probing experiments @xcite , and propose a new set of metrics that quantify representation collapse objectively, without requiring extra datasets/training. we observe that repina shows the strongest resistance to representation degradation among all methods. ## repina: representation projection invariance our method avoids representation collapse by preventing undesirable changes in representations dur-ing the fine-tuning process. a straightforward implementation would anchor representations during fine-tuning to their pre-trained values. that is, the final loss l would combine the standard fine-tuning objective and regularizer of the deviation in representations: various transformations of a representation maintain its expressivity (such as linear shift; fig 2b ). while such transformations do not change the information content of a representation, they incur a high regularization loss based on equation 1. to address this issue and allow flexibility in representations while preserving their expressive capacity, we propose representation projection invariance regularization (repina): here φ is a class of dimension preserving functions chosen before the fine-tuning process and defines the strength of regularization. the intuition behind the regularization objective is to incentivize the representations to be invariant under some projection ϕ ∈ φ; pre-trained representations can be constructed from the fine-tuned representations by applying a function ϕ ∈ φ. for instance, if φ is the set of linear functions {ϕ | ∃w, b : ϕ(z) = w z + b} , then we bias finetuned representations to be linear transformations of the fine-tuned representations. thus, regularization loss in case of fig 2a would be zero since there exists a linear mapping from fine-tuned representations to the pre-trained representations. however, regularization loss for fig 2a would be high as there does not exist such a linear mapping from fine-tuned to the pre-trained representations. choice of the class of functions φ: φ defines the strength of the regularizer. for instance, a singleton φ containing an identity function is the strongest regularizer which keeps fine-tuned representations close to the pre-trained representations (equivalent to equation 1). conversely, for a rich φ, e.g., deep and wide neural networks, ϕ can be chosen to reconstruct the lost information in representation even if there is a severe degradation. thus, it provides a weak regularization. choice of φ then depends on how prone fine-tuning is to overfitting and how strong of a regularization method is needed. for instance, few-shot setting may require the strongest regularization and larger training datasets may require milder regularization. in this paper, we experiment with φ containing identity function (repina i ) and shallow multilayer perceptrons (repina mlp ). choosing the right representation: normally, sentence-level representations are obtained from the final encoder blocks of a transformer model. however, it may be more beneficial to use representations from the lower layers. in fact, @xcite show that re-initializing weights of the top layers of encoder improves fine-tuning performance, suggesting that representation consistency may not be desirable for the top layers. thus, we consider a variant of regularization that uses representations from the intermediate layers of encoder. explaining representation invariance regularization as implicitly learning multiple tasks: consider the overfitting in fig 2a again. it can be prevented by fine-tuning the representations for multiple tasks simultaneously instead of a single task. this multi-tasking is a well-known way to prevent overfitting. it not only prevents overfitting of representations but can also improves generalization performance for all of the tasks. we show that repina's regularization objective (equation 2) is equivalent to fine-tuning on multiple hypothetical tasks. due to space constraints, we defer further discussion on the connection and the formal proof of equivalence in appendix b. mitigating representation collapse: aghajanyan et al. (2020) study representation collapse and propose two methods, r3f and r4f, to address it. r3f induces bias towards a solution with a locally smooth prediction function, and r4f extends r3f by adding a lipschitzness constraint on the top classification layer. some of the other methods which implicitly target representation collapse are freelb and smart. freelb uses adversarial training to improve fine-tuning @xcite and smart is a trust region-based method that avoids aggressive updates during fine-tuning @xcite . r3f @xcite has been shown to outperform both of these methods. thus, we only include r3f in our set of baselines. a method that specifically targets representations during fine-tuning is supervised contrastive learning (scl). scl induces representations of examples with the same label to be close to each other and far from the examples of other classes @xcite . a major disadvantage of scl is a requirement for large mini-batch size and, hence, heavy memory consumption. we implement a memory-efficient scl version but exclude the original implementation from the baselines due to computational cost (see appendix h). another method which can be potentially useful for mitigating representation collapse and form part of our baseline is data augmentation. it can be done via back translation, synonymn replacement, random deletion, and synthetic noise and is known to improve generalization performance @xcite . catastrophic forgetting @xcite is a phenomenon closely related to representation collapse. in a sequential training setting, it refers to forgetting information learnt from previous tasks while being trained on a new task. in our context, this means forgetting the pre-training language modeling task while training for the fine-tuning task. in contrast to catastrophic forgetting, we measure representation collapse as the loss in expressive power of the representations irrespective of the performance on pre-training task. a method known to alleviate catastrophic forgetting is weight consolidation @xcite . it regularizes the fine-tuning process by encouraging fine-tuned weights to be close to the pre-trained weights. in contrast to weight consolidation, repina does not put direct constraints on weight updates, but tries to control the structural changes in representations. due to our limited space, we discuss further details on related works in appendix a. ## results: generalization performance in this section, we present experimental results with the baselines introduced in the earlier section. ## degradation of representations: analytical perspective here we quantify the representation collapse. ## conclusion in this paper, we propose a novel representation invariance regularizer targeted at avoiding representation degradation during finetuning. it has a knob that can control strength of regularization. we experiment with two choices of this knob, re-pina i and repina mlp and show that they both achieve significant performance gain across 13 tasks, including few-shot and label noise settings, and improve generalization performance. we also study the degradation of representations during fine-tuning, representation collapse, and propose new metrics to quantify it. our methods reduce representation collapse and improve ood robustness.