Supervised Contrastive Learning for Pre-trained Language Model Fine-tuning

1 INTRODUCTION . State-of-the-art for most existing natural language processing ( NLP ) classification tasks is achieved by models that are first pre-trained on auxiliary language modeling tasks and then fine-tuned on the task of interest with cross-entropy loss ( Radford et al. , 2019 ; Howard & Ruder , 2018 ; Liu et al. , 2019 ; Devlin et al. , 2019 ) . Although ubiquitous , the cross-entropy loss – the KL-divergence between one-hot vectors of labels and the distribution of model ’ s output logits – has several shortcomings . Cross entropy loss leads to poor generalization performance ( Liu et al. , 2016 ; Cao et al. , 2019 ) , and it lacks robustness to noisy labels ( Zhang & Sabuncu , 2018 ; Sukhbaatar et al. , 2015 ) or adversarial examples ( Elsayed et al. , 2018 ; Nar et al. , 2019 ) . Effective alternatives have been proposed to modify the reference label distributions through label smoothing ( Szegedy et al. , 2016 ; Müller et al. , 2019 ) , Mixup ( Zhang et al. , 2018 ) , CutMix ( Yun et al. , 2019 ) , knowledge distillation ( Hinton et al. , 2015 ) or self-training ( Yalniz et al. , 2019 ; Xie et al. , 2020 ) . Fine-tuning using cross entropy loss in NLP also tends to be unstable across different runs ( Zhang et al. , 2020 ; Dodge et al. , 2020 ) , especially when supervised data is limited , a scenario in which pre-training is particularly helpful . To tackle the issue of unstable fine-tuning and poor generalization , recent works propose local smoothness-inducing regularizers ( Jiang et al. , 2020 ) and regularization methods inspired by the trust region theory ( Aghajanyan et al. , 2020 ) to prevent representation collapse . Empirical evidence suggests that fine-tuning for more iterations , reinitializing top few layers ( Zhang et al. , 2020 ) , and using debiased Adam optimizer during fine-tuning ( Mosbach et al. , 2020 ) can make the fine-tuning stage more stable . Inspired by the learning strategy that humans utilize when given a few examples , we seek to find the commonalities between the examples of each class and contrast them with examples from other classes . We hypothesize that a similarity-based loss will be able to hone in on the important dimensions of the multidimensional hidden representations hence lead to better few-shot learning results and be more stable while fine-tuning pre-trained language models . We propose a novel objective for fine-tuning that includes a supervised contrastive learning ( SCL ) term that pushes the examples from the same class close and the examples from different classes further apart . The SCL ∗Work done during Facebook AI research internship , correspondence to bgunel @ stanford.edu . term is similar to the contrastive objectives used in self-supervised representation learning across image , speech , and video domains . ( Sohn , 2016 ; Oord et al. , 2018 ; Wu et al. , 2018 ; Bachman et al. , 2019 ; Hénaff et al. , 2019 ; Baevski et al. , 2020 ; Conneau et al. , 2020 ; Tian et al. , 2020 ; Hjelm et al. , 2019 ; Han et al. , 2019 ; He et al. , 2020 ; Misra & Maaten , 2020 ; Chen et al. , 2020a ; b ) . Unlike these methods , however , we use a contrastive objective for supervised learning of the final task , instead of contrasting different augmented views of examples . In few-shot learning settings ( 20 , 100 , 1000 labeled examples ) , the addition of the SCL term to the finetuning objective significantly improves the performance on several natural language understanding classification tasks from the popular GLUE benchmark ( Wang et al. , 2019 ) over the very strong baseline of fine-tuning RoBERTa-Large with cross-entropy loss only . Furthermore , pre-trained language models fine-tuned with our proposed objective are not only robust to noise in the fine-tuning training data , but can also exhibit improved generalization to related tasks with limited labeled task data . Our approach does not require any specialized network architectures ( Bachman et al. , 2019 ; Hénaff et al. , 2019 ) , memory banks ( Wu et al. , 2018 ; Tian et al. , 2020 ; Misra & Maaten , 2020 ) , data augmentation of any kind , or additional unsupervised data . To the best of our knowledge , our work is the first to successfully integrate a supervised contrastive learning objective for fine-tuning pre-trained language models . We empirically demonstrate that the new objective has desirable properties across several different settings . Our contributions in this work are listed in the following : • We propose a novel objective for fine-tuning pre-trained language models that includes a supervised contrastive learning term , as described in Section 2 . • We obtain strong improvements in the few-shot learning settings ( 20 , 100 , 1000 labeled examples ) as shown in Table 2 , leading up to 10.7 points improvement on a subset of GLUE benchmark tasks ( SST-2 , QNLI , MNLI ) for the 20 labeled example few-shot setting , over a very strong baseline – RoBERTa-Large fine-tuned with cross-entropy loss . • We demonstrate that our proposed fine-tuning objective is more robust , in comparison to RoBERTa-Large fine-tuned with cross-entropy loss , across augmented noisy training datasets ( used to fine-tune the models for the task of interest ) with varying noise levels as shown in Table 3 – leading up to 7 points improvement on a subset of GLUE benchmark tasks ( SST-2 , QNLI , MNLI ) across augmented noisy training datasets . We use a backtranslation model to construct the augmented noisy training datasets of varying noise levels ( controlled by the temperature parameter ) , as described in detail in Section 4.2 . • We show that the task-models fine-tuned with our proposed objective have improved generalizability to related tasks despite having limited availability of labeled task data ( Table 7 ) . This led to a 2.9 point improvement on Amazon-2 over the task model fine-tuned with cross-entropy loss only . Moreover , it considerably reduced the variance across few-shot training samples , when transferred from the source SST-2 sentiment analysis task model . 2 APPROACH . We propose a novel objective that includes a supervised contrastive learning term for fine-tuning pre-trained language models . The loss is meant to capture the similarities between examples of the same class and contrast them with the examples from other classes . For a multi-class classification problem with C classes , we work with a batch of training examples of size N , { xi , yi } i=1 , ... N . Φ ( · ) ∈ Rd denotes an encoder that outputs the l2 normalized final encoder hidden layer before the softmax projection ; Nyi is the total number of examples in the batch that have the same label as yi ; τ > 0 is an adjustable scalar temperature parameter that controls the separation of classes ; yi , c denotes the label and ŷi , c denotes the model output for the probability of the ith example belonging to the class c ; λ is a scalar weighting hyperparameter that we tune for each downstream task and setting . The overall loss is then given in the following : L = ( 1− λ ) LCE + λLSCL ( 1 ) LCE = − 1 N N∑ i=1 C∑ c=1 yi , c · logŷi , c ( 2 ) LSCL = N∑ i=1 − 1 Nyi − 1 N∑ j=1 1i6=j1yi=yj log exp ( Φ ( xi ) · Φ ( xj ) /τ ) ∑N k=1 1i 6=k exp ( Φ ( xi ) · Φ ( xk ) /τ ) ( 3 ) The overall loss is a weighted average of CE and the proposed SCL loss , as given in the equation ( 1 ) . The canonical definition of the multi-class CE loss that we use is given in equation ( 2 ) . The novel SCL loss is given in the equation ( 3 ) . This loss can be applied using a variety of encoders Φ ( · ) ∈ Rd – for example a ResNet for a computer vision application or a pre-trained language model such as BERT for an NLP application . In this work , we focus on fine-tuning pre-trained language models for single sentence and sentence-pair classification settings . For single sentence classification , each example xi consists of sequence of tokens prepended with the special [ CLS ] token xi = [ [ CLS ] , t1 , t2 , . . . , tL , [ EOS ] ] . The length of sequence L is constrained such that L < Lmax . Similarly , for sentence-pair classification tasks , each example xi is a concatenation of two sequences of tokens [ t1 , t2 , . . . tL ] and [ s1 , s2 , . . . , sM ] corresponding to the sentences with special tokens delimiting them : xi = [ [ CLS ] , t1 , t2 , . . . , tL , [ SEP ] , s1 , s2 , . . . , sM , [ EOS ] ] . The length of concatenated sequences is constrained such that L+M < Lmax . In both cases , Φ ( xi ) ∈ Rd uses the embedding of [ CLS ] token as the representation for example xi . These choices follow standard practices for fine-tuning pre-trained language models for classification ( Devlin et al. , 2019 ; Liu et al. , 2019 ) . Empirical observations show that both l2 normalization of the encoded embedding representations and an adjustable scalar temperature parameter τ improve performance . Lower temperature increases the influence of examples that are harder to separate , effectively creating harder negatives . Using hard negatives has been previously shown to improve performance in the context of margin-based loss formulations such as triplet loss ( Schroff et al. , 2015 ) . The empirical behavior of the adjustable temperature parameter is consistent with the observations of previous work related to supervised contrastive learning . ( Chen et al. , 2020a ; Khosla et al. , 2020 ) . Relationship to Self-Supervised Contrastive Learning Self-supervised contrastive learning has shown success in learning powerful representations , particularly in the computer vision domain . ( Chen et al. , 2020a ; He et al. , 2020 ; Tian et al. , 2020 ; Mnih & Kavukcuoglu , 2013 ; Gutmann & Hyvärinen , 2012 ; Kolesnikov et al. , 2019 ) Self-supervised learning methods do not require any labeled data ; instead they sample a mini batch from unsupervised data and create positive and negative examples from these samples using strong data augmentation techniques such as AutoAugment ( Cubuk et al. , 2019 ) or RandAugment ( Cubuk et al. , 2020 ) for computer vision . Positive examples are constructed by applying data augmentation to the same example ( cropping , flipping , etc . for an image ) , and negative examples are simply all the other examples in the sampled mini batch . Intuitively , selfsupervised contrastive objectives are learning representations that are invariant to different views of positive pairs ; while maximizing the distance between negative pairs . The distance metric used is often the inner product or the Euclidean distance between vector representations of the examples . For a batch of size N , self-supervised contrastive loss is defined as : Lself = 2N∑ i=1 − log exp ( Φ ( x′2i−1 ) · Φ ( x′2i ) /τ ) ∑2N k=1 1i 6=k exp ( Φ ( x ′ i ) · Φ ( x′k ) /τ ) ( 4 ) where Φ ( · ) ∈ Rd denotes an encoder that outputs the l2 normalized final encoder hidden layer before the softmax projection ; τ > 0 is a scalar temperature parameter . A is defined as a data augmentation block that generates two randomly generated augmented examples , x′2i and x ′ 2i−1 from the original example xi : A ( { xi , yi } i=1 , ... N ) = { x′i , y′i } i=1 , ... 2N . As an example , A can be RandAugment for a computer vision application ; or it could be a back-translation model for an NLP application .

Summary: The paper proposes a new training objective for fine-tuning pre-trained models: a weighted sum of the classical cross-entropy (CE) and a new supervised contrastive learning term (SCP). The latter uses the (negated) softmax over the embedding distances (i.e. dot products) between a training instance and all other instances in the batch with the same label. In contrast to the more traditional self-supervised contrastive learning (where positive pairs are obtained by applying transformations to the original data instance), there is no data augmentation; two examples with the same label constitute a positive pair.