T a CL : Improving BERT Pre-training with Token-aware Contrastive Learning (2022)

Abstract: Masked language models (MLMs) such as BERT have revolutionized the field of Natural Language Understanding in the past few years. However, existing pre-trained MLMs often output an anisotropic distribution of token representations that occupies a narrow subset of the entire representation space. Such token representations are not ideal, especially for tasks that demand discriminative semantic meanings of distinct tokens. In this work, we propose TaCL (Token-aware Contrastive Learning), a novel continual pre-training approach that encourages BERT to learn an isotropic and discriminative distribution of token representations. TaCL is fully unsupervised and requires no additional data. We extensively test our approach on a wide range of English and Chinese benchmarks. The results show that TaCL brings consistent and notable improvements over the original BERT model. Furthermore, we conduct detailed analysis to reveal the merits and inner-workings of our approach.

## introduction since the rising of bert @xcite , masked language models (mlms) have become the de facto backbone for almost all natural language understanding (nlu) tasks. despite their clear success, many existing language models pre-trained with mlm objective suffer from the anisotropic problem @xcite . that is, their token representations reside in a narrow subset of the representation space, therefore being less discriminative and less powerful in capturing the semantic differences of distinct tokens. recently, great advancement has been made in continually training mlms with unsupervised sentence-level contrastive learning, aiming at creating more discriminative sentence-level representations @xcite @xcite figure 1 : an overview of tacl. the student learns to make the representation of a masked token closer to its "reference" representation produced by the teacher (solid arrow) and away from the representations of other tokens in the same sequence (dashed arrows). @xcite @xcite . however, such representations are only evaluated as sentence embeddings and there is no evidence that they will benefit other well-established nlu tasks. we show that these approaches hardly bring any benefit to challenging tasks like squad @xcite . in this paper, we argue that the key of obtaining more discriminative and transferrable representations lies in learning contrastive and isotropic token-level representations. to this end, we propose tacl (token-aware contrastive learning), a new continual pre-training approach that encourages bert to learn discriminative token representations. specifically, our approach involves two models (a student and a teacher) that are both initialized from the same pre-trained bert. during the learning stage, we freeze the parameters of the teacher and continually optimize the student model with (1) the original bert pre-training objectives (masked language modelling and next sentence prediction) and (2) a newly proposed tacl objective. the tacl loss is obtained by contrasting the student representations of masked tokens against the "reference" representations produced by the teacher without masking the input tokens. in figure 1 , we provide an overview of our approach. we extensively test our approach on a wide range of english and chinese benchmarks and illustrate that tacl brings notable performance improvements on most evaluated datasets ( §3.1.1). these results validate that more discriminative and isotropic token representations lead to better model performances. additionally, we highlight the benefits of using our token-level method compared to current state-of-the-art sentence-level contrastive learning techniques on nlu tasks ( §3.2.1). we further analyze the inner workings of tacl and its impact on the token representation space ( §3.2.2). our work, to the best of our knowledge, is the first effort on applying contrastive learning to improve token representations of transformer models. we hope the findings of this work could facilitate further development of methods on the intersection of contrastive learning and representation learning at a more fine-grained granularities. note that the learning of the student is fully unsupervised and can be realized using the original pretraining corpus. after the learning is completed, we fine-tune the student model on downstream tasks. ## experiment we test our approach on a wide range of benchmarks in two languages. for english benchmarks, we evaluate the bert base and bert large models. for chinese benchmarks, we test the bert base model. 2 after initializing the student and teacher, we continually pre-train the student on the same wikipedia corpus as in @xcite for 150k steps. the training samples are truncated with a maximum length of 256 and the batch size is set as 256. the temperature τ in eq. ( 1 ) is set as 0.01. same as @xcite , we optimize the model with adam optimizer (kingma and ba, 2015) with weighted decay, and an initial learning rate of 1e-4 (with warm-up ratio of 10%). ## conclusion in this work, we proposed tacl, a novel approach that applies token-aware contrastive learning for the continual pre-training of bert. extensive experiments are conducted on a wide range of english and chinese benchmarks. the results show that our approach leads to notable performance improvement across all evaluated benchmarks. we then delve into the inner-working of tacl and demonstrate that our performance gain comes from a more discriminative distribution of token representations.