COCO-LM: Correcting and Contrasting Text Sequences for Language Model Pretraining

1 Introduction . Pretrained language models ( PLMs ) have reshaped the way AI systems process natural language [ 11 , 36 , 39 , 40 ] . Before task-specific training , it is now a common practice to first pretrain the deep neural networks , often Transformers [ 53 ] , via a self-supervised token-level language modeling task [ 29 , 31 , 40 ] . Whether it is autoregressive [ 39 ] , permutational [ 62 ] , or masked language modeling ( MLM ) [ 11 ] , the Transformer networks are pretrained to recover some omitted tokens using the rest of input texts . Then the language semantics captured during pretraining are conveyed to downstream tasks via the pretrained Transformer parameters [ 5 , 8 , 44 ] . Recent research [ 14 , 16 , 25 , 43 ] observed several challenges in this self-supervised learning framework . One challenge is its efficiency . After pretrained for a while with the standard token-level language modeling , the networks have already captured the basic language patterns , making a large fraction of pretraining signals no longer informative . Linear improvement in the model effectiveness often requires exponentially more pretraining compute and parameters [ 25 ] , which is unsustainable . Another challenge is the anisotropy of text representations from pretrained models . The sequence representations from many pretrained models are quite irregular [ 30 , 43 ] and require dedicated fine-tuning approaches to be useful in sequence-level applications [ 32 , 60 ] . Clark et al . [ 7 ] proposed a new pretraining strategy , ELECTRA , that uses an auxiliary language model ( “ generator ” ) to replace tokens in input texts and pretrains the main Transformer ( “ discriminator ” ) to ∗Part of this work was done while Yu was interning at Microsoft . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . detect replaced tokens . This improves the pretraining efficiency and effectiveness , but pretraining via binary classification hinders the model ’ s usage on applications requiring language modeling capability ( e.g. , prompt-based learning [ 15 , 28 , 46 ] ) . It could further distort the representation space as the Transformers are pretrained to output the same “ non-replacement ” label for all actual tokens . In this paper , we present a new self-supervised learning approach , COCO-LM , that pretrains Language Models by COrrecting and COntrasting corrupted text sequences . Following ELECTRA-style pretraining , COCO-LM employs an auxiliary model to corrupt the input texts , upon which it introduces two new pretraining tasks for the main Transformer , one at token level and one at sequence level . The token-level task , corrective language modeling ( CLM ) , pretrains the main Transformer to detect and correct the tokens in the corrupted sequences . It uses a multi-task setup to combine the benefits of replaced token detection and language modeling . The sequence-level task , sequence contrastive learning ( SCL ) , pretrains the model to align text sequences originated from the same source sequence and enforce uniformity of the representation space . In our experiments on GLUE [ 54 ] and SQuAD [ 41 ] benchmarks , COCO-LM not only outperforms state-of-the-art pretraining approaches in effectiveness , but also significantly improves the pretraining efficiency . Under the same setting , COCO-LM matches the MNLI accuracy of RoBERTa and ELECTRA with 60 % and 50 % of their GPU hours in pretraining , respectively . When pretrained with the same number of steps , COCO-LM outperforms the previous best models by 1+ GLUE average points under the standard base/large-sized model evaluations . With 367 million parameters , COCOLMLarge++ reaches the MNLI accuracy of Megatron3.9B [ 49 ] , one of the largest BERT-style model with 3.9 billion parameters . Our analyses provide further insights on the advantage of CLM in learning token representations and its effectiveness in prompted-based fine-tuning , as well as the benefit of SCL in ensuring alignment and uniformity in the representation space for better generalization1 . 2 Related Work . Various token-level tasks have been used to pretrain language models . The most classic auto-regressive language modeling is to predict a token given all the previous tokens , or all subsequent ones [ 36 , 39 ] . BERT uses masked language modeling ( MLM ) that recovers randomly masked tokens using the rest input . XLNet proposes permutation language modeling that conducts MLM in an autoregressive manner [ 62 ] . UniLM uses pseudo MLM which unifies autoregressive and MLM tasks [ 1 , 13 ] . Sequence-level tasks are also explored , which often pretrain the model to predict certain cooccurrences of sequence pairs . For example , next sentence prediction [ 11 ] , sentence ordering [ 27 ] and previous sentence prediction [ 56 ] concatenate two sentences ( either correlated or random ) , and train the Transformer to classify the pair . Empirically , MLM is still among the most effective tasks to pretrain encoders [ 29 , 31 , 40 ] . RoBERTa [ 31 ] found the sentence-level task in BERT not benefitial and discarded it . BART [ 29 ] and T5 [ 40 ] both observed that MLM is often the most effective task . The empirical advantages of other pretraining tasks are more task-specific , for example , entity related masks for knowledge intensive applications [ 20 , 24 ] , and sequence-level tasks for long form text modeling [ 42 ] . Instead of randomly altering texts , ELECTRA [ 7 ] uses a smaller auxiliary Transformer pretrained by MLM to replace some tokens in the text sequences using its language modeling probability , and pretrains the main Transformer to detect the replaced tokens . ELECTRA achieves state-of-the-art accuracy in many language tasks [ 7 ] . Later , Clark et el . [ 6 ] developed ELECTRIC , which pretrains encoders by contrasting original tokens against negatives sampled from a cloze model . ELECTRIC re-enables the language modeling capability but underperforms ELECTRA in downstream tasks . Our work is also related to contrastive learning which has shown great success in visual representation learning [ 4 , 22 , 34 ] . Its effectiveness of in language is more observed in the fine-tuning stage , for example , in sentence representation [ 16 ] , dense retrieval [ 60 ] , and GLUE fine-tuning [ 19 ] . 3 Method . We present the preliminaries of PLMs , their challenges , and the new COCO-LM framework . 1Code and pretrained models can be found at https : //github.com/microsoft/COCO-LM . 3.1 Preliminary on Language Model Pretraining . In this work we focus on pretraining BERT-style bidirectional Transformer encoders [ 11 ] that are widely used in language representation tasks . We first recap the masked language modeling ( MLM ) task introduced by BERT [ 11 ] and then discuss the pretraining framework of ELECTRA [ 7 ] . BERT Pretraining uses the masked language modeling task ( MLM ) [ 11 ] , which is to take an input sequenceXorig = [ xorig1 , . . . , x orig i , . . . , x orig n ] , with 15 % random tokens replaced by [ MASK ] symbols ( e.g. , the i-th token ) , and train the model to predict the original tokens at the masked positions : [ xorig1 , . . . , [ MASK ] i , . . . , x orig n ] Transformer−−−−−−→H MLM Head−−−−−−→ pMLM ( x|hi ) , where the Transformer generates contextualized representations H = { hi } ni=1 . The MLM Head predicts the masked token from the vocabulary V using the hidden representation hi and token embeddings x . The pretraining minimizes the MLM loss on the set of masked positionsM . Specifically , pMLM ( x|hi ) = exp ( x > hi ) ∑ xt∈V exp ( x > t hi ) ; LMLM = E ( − ∑ i∈M log pMLM ( xorigi ∣∣hi ) ) . ELECTRA Pretraining uses two Transformers , a “ generator ” pretrained by MLM , and a “ discriminator ” pretrained using the generator ’ s outputs . We refer them as auxiliary and main Transformers , as the former is discarded after pretraining and the latter may be trained by “ generative ” tasks too . The auxiliary model outputs a corrupted sequence XMLM by sampling from its predicted probability : xMLMi ∼ pMLM ( x|hi ) , if i ∈M ; xMLMi = x orig i , else . ( 1 ) The masked positions are replaced by sampled tokens considered plausible in context by the auxiliary Transformer , which are more deceiving than random replacements . ELECTRA uses a skinnier auxiliary network ( e.g. , hidden dimension is 1/3 of the main model ) to control the signal difficulty . The main Transformer takes XMLM and classifies the replaced tokens : XMLM Main Transformer−−−−−−−−−→H RTD Head−−−−−→ pRTD ( 1 ( xMLMi = x orig i ) ∣∣hi ) , where 1 ( · ) is the indicator function . The Replaced Token Detection ( RTD ) head uses a sigmoid linear layer to output the binary probability , and the main Transformer is trained with binary cross entropy loss . The RTD task is trained on all tokens instead of masked ones and improves efficiency . The two Transformers are pretrained jointly . The auxiliary model gradually generates more realistic replacement tokens and the main model learns to better detect them . This forms a natural learning curriculum and significantly improves ELECTRA ’ s accuracy in downstream tasks [ 7 ] . 3.2 Challenges of ELECTRA-Style Pretraining . Missing Language Modeling Benefits . The classification task in ELECTRA is simpler and more stable [ 61 ] , but raises two challenges . The first is the lack of language modeling capability which is a necessity in some tasks [ 6 ] . For example , prompt-based learning requires a language model to generate labels [ 15 , 33 , 45 , 46 ] . The second is that the binary classification task may not be sufficient to capture certain word-level semantics that are critical for token-level tasks . Squeezing Representation Space . Another challenge is that the representations from Transformer-based language models often reside in a narrow cone , where two random sentences have high similarity scores ( lack of uniformity ) , and closely related sentences may have more different representations ( lack of alignment ) [ 14 , 16 , 30 ] . Figure 1 illustrates such behaviors with random sentence pairs ( from pretraining corpus ) and semantically similar pairs ( those annotated with maximum similarity from STS-B [ 3 ] ) . With RoBERTa , the cosine similarities of most random sentence pairs are near 0.8 , bigger than many semantically similar pairs . The representation space from ELECTRA is even more squeezed . Nearly all sentence pairs , both random and similar ones , have around 0.9 cosine similarity . This may not be surprising as ELECTRA is pretrained to predict the same output ( “ non-replacement ” ) for all tokens in these sequences . The irregular representation space raises the risk of degeneration [ 37 , 55 ] and often necessitates sophisticated post-adjustment or fine-tuning to improve the sequence representations [ 16 , 30 , 32 , 60 ] . 3.3 COCO-LM Pretraining . COCO-LM also employs an auxiliary Transformer to construct the corrupted text sequence , as in Eqn . ( 1 ) , but it introduces two new pretraining tasks upon the corrupted sequences to address the challenges previously described . In the rest of this section , we present these two tasks and then the detailed configurations of COCO-LM . Its framework is illustrated in Figure 2 . Corrective Language Modeling ( CLM ) trains the main Transformer to recover the original tokens , given the corrupted text sequence XMLM : XMLM Main Transformer−−−−−−−−−→H CLM Head−−−−−−→ pCLM ( x|hi ) . The CLM Head uses the hidden representations H to output a language modeling probability , instead of a binary classification score . The forward pass of the CLM Head is the same as All-Token MLM , a variation of ELECTRA [ 7 ] that consists of a language modeling layer and a binary classification layer for the copy mechanism : pLM ( xi|hi ) = 1 ( xi = x MLM i ) pcopy ( 1|hi ) + pcopy ( 0|hi ) exp ( x > i hi ) ∑ xt∈V exp ( x > t hi ) , pcopy ( yi|hi ) = exp ( yi ·w > copyhi ) / ( exp ( w > copyhi ) + 1 ) , where wcopy is a learnable weight and pcopy ( yi|hi ) is the copy mechanism ( yi = 1 when the input token is original and can be directly copied to the output ; yi = 0 when the input token needs to be corrected to another token from the vocabulary ) . In ELECTRA , All-Token MLM performs worse than RTD [ 7 ] . Language modeling on the corrupted text sequence XMLM is hard as the replaced tokens from the auxiliary model are more deceiving than [ MASK ] . To improve the language model learning , different from All-Token MLM , CLM employs a multi-task setup that combines the RTD task to explicitly train the copy mechanism pcopy ( · ) : Lcopy =− E ( n∑ i=1 1 ( xMLMi = x orig i ) log pcopy ( 1|hi ) + 1 ( xMLMi 6= x orig i ) log pcopy ( 0|hi ) ) , ( 2 ) LLM =− E ( ∑ i∈M log pLM ( xorigi ∣∣hi ) ) =− E ( ∑ i∈M log ( 1 ( xMLMi = x orig i ) psgcopy ( 1|hi ) + psgcopy ( 0|hi ) exp ( x > i hi ) ∑ xt∈V exp ( x > t hi ) ) ) , LCLM =λcopyLcopy + LLM . The hyperparameter λcopy balances the weights of the two tasks . The binary cross entropy loss in Eqn . ( 2 ) explicitly trains the copy probability . We also use stop gradient ( sg ) to decouple the gradient backpropagation to pcopy ( · ) from the LM task . This way , the main Transformer first learns the easier classification task and then uses it to help learn the harder LM task . The binary classification task is trained on all tokens while the language modeling task is trained only on masked positions . CLM combines the advantages of MLM and ELECTRA : The main Transformer is trained on all tokens with the help of the binary classification task while also being able to predict words , thus enjoying the efficiency benefits of ELECTRA and preserving the language modeling benefits . Sequence Contrastive Learning ( SCL ) forms a contrastive learning objective upon the sequence embeddings to learn more robust representations . Broadly , contrastive learning is to align a positive pair of instances , often different views of the same information [ 4 , 34 ] , in contrast to unrelated negative instances [ 22 , 60 ] . The different views are often obtained by applying data augmentations on the same input , for example , rotation , cropping , and blurring on visual representations [ 4 , 34 ] , so that the neural networks can learn representations robust to these data alterations . In COCO-LM , the corrupted sequence XMLM already provides a form of data augmentation . We pair it with another augmentation , Xcrop , a randomly cropped contiguous span of Xorig ( the length of Xcrop is 90 % of Xorig so that the major sequence meaning is preserved ) , to construct the positive pair and to contrast with random negatives . Specifically , a training batch B in SCL includes a random set of corrupted and cropped sequences : B = { ( XMLM1 , X crop 1 ) , . . . , ( X MLM N , X crop N ) } , withXMLMk andX crop k originated fromX orig k . A positive contrastive pair ( X , X+ ) consists of either ( XMLMk , X crop k ) or ( X crop k , X MLM k ) ( symmetrical contrast ) . The negative instances are all the remaining sequences in the batch B− = B \ { ( X , X+ ) } . The contrastive loss is formulated as : LSCL = −E ( log exp ( cos ( s , s+ ) /τ ) exp ( cos ( s , s+ ) /τ ) + ∑ X−∈B− exp ( cos ( s , s− ) /τ ) ) , = −E ( cos ( s , s+ ) /τ − log ( exp ( cos ( s , s+ ) /τ ) + ∑ X−∈B− exp ( cos ( s , s− ) /τ ) ) ) , ( 3 ) where s , s+ , s− are the representations of X , X+ , X− , respectively , from the main Transformer ( i.e. , h [ CLS ] ) . The similarity metric is cosine similarity ( cos ) and the temperature τ is set to 1 . As shown in Wang et al . [ 55 ] , the first term in Eqn . ( 3 ) ( cos ( s , s+ ) ) improves alignment of the space . It encourages representations to be robust to the corruptions and the alterations on the original text . The second term in Eqn . ( 3 ) promotes uniformity . It pushes unrelated sequences apart in the representation space and ensures low cosine similarity between random data points . Several studies have observed improved generalization ability from better alignment and uniformity [ 16 , 37 , 55 ] . AligningXMLM withXcrop requires the main Transformer to produce sequence representations robust to both token-level ( i.e. , MLM replacements ) and sequence-level ( i.e. , cropping ) alterations . The model is thus encouraged to reason more using partially altered sequences to recover the original information . Overall Training . COCO-LM uses the following loss function : LCOCO-LM = LAux.MLM + LMainCLM + LMainSCL . ( 4 ) The auxiliary Transformer is pretrained by masked language modeling ( MLM ) and generates corrupted sequences . The main Transformer is pretrained to correct the corruption ( CLM ) and to contrast the corrupted sequences with the cropped sequences ( SCL ) . The two Transformers are pretrained jointly with the loss in Eqn . ( 4 ) . The main Transformer is used in downstream applications . Network Configurations . Similar to ELECTRA , the auxiliary Transformer is smaller than the main model , but we use different configurations in the auxiliary model : ( 1 ) We reduce the number of layers to 1/3 or 1/4 ( under base or large model setup , respectively ) but keep its hidden dimension the same with the main model , instead of shrinking its hidden dimensions ; ( 2 ) We disable dropout in it when sampling replacement tokens . We find such configurations empirically more effective and use them as the backbone of COCO-LM . The main Transformer follows the standard architecture of BERT/ELECTRA and can be easily adopted by downstream application pipelines with almost no changes .

Summary: The paper proposes a pre-trained technique based on correcting tokens and contrasting sequences of text. The correcting objective is based on ELECTRA with a modification in the formulation of multi-task setup where the one task is to identify the corrupted tokens and the other task is to correct it. The contrasting objective aims to maximize the similarity between similar sentences while pushing dissimilar sentences away in cosine space. With the two modifications, COCO-LM is trained much faster than ELECTRA. The paper also presents a comprehensive evaluation on GLUE and SQUAD 2.0 benchmarks and ablation study of their method.