BIO ptimus: Pre-training an Optimal Biomedical Language Model with Curriculum Learning for Named Entity Recognition (2023)

Abstract: Using language models (LMs) pre-trained in a self-supervised setting on large corpora and then fine-tuning for a downstream task has helped to deal with the problem of limited label data for supervised learning tasks such as Named Entity Recognition (NER). Recent research in biomedical language processing has offered a number of biomedical LMs pre-trained using different methods and techniques that advance results on many BioNLP tasks, including NER. However, there is still a lack of a comprehensive comparison of pre-training approaches that would work more optimally in the biomedical domain. This paper aims to investigate different pre-training methods, such as pre-training the biomedical LM from scratch and pre-training it in a continued fashion. We compare existing methods with our proposed pre-training method of initializing weights for new tokens by distilling existing weights from the BERT model inside the context where the tokens were found. The method helps to speed up the pre-training stage and improve performance on NER. In addition, we compare how masking rate, corruption strategy, and masking strategies impact the performance of the biomedical LM. Finally, using the insights from our experiments, we introduce a new biomedical LM (BIOptimus), which is pre-trained using Curriculum Learning (CL) and contextualized weight distillation method. Our model sets new states of the art on several biomedical Named Entity Recognition (NER) tasks. We release our code and all pre-trained models.

## introduction since the introduction of transformer architecture @xcite , transfer learning has gained immense popularity in natural language processing (nlp). pre-training lms on a large corpus (devlin et al., 2019; liu et al., 2019) allows a largescale knowledge extraction. in the current pipeline of knowledge transfer in nlp, the pre-training stage can be viewed as preparing the model with broad general knowledge @xcite that can later be fine-tuned to new tasks or adapted to different domains. however, applying these models in a new domain directly without additional pretraining steps cannot achieve high performance because of a word distribution shift from a general domain vocabulary to a new domain vocabulary. the importance of a pre-trained model for the biomedical domain is even more significant due to the difficulties of constructing supervised training data because it requires expert knowledge for quality annotations. the biomedical domain has greatly benefitted from domain-specified lms. to date, various methods have been developed to pre-train biomedical language models: pre-training from scratch @xcite , continual pretraining @xcite , and hybrid approaches @xcite @xcite . nonetheless, more research is needed to compare the optimal ways to pre-train the biomedical lms. comparison of different pre-training approaches and techniques primarily targets general domain lms @xcite @xcite . what still needs to be clarified is the effects of these methods on the pre-training biomedical lm. typically, general domain lms are pre-trained on wikipedia, book corpora, and common-crawl @xcite @xcite @xcite . these corpora are considerably different from the biomedical corpora in terms of domain specifics (use of biomedical terminology) and the style of scientific text (most of the models are trained on pubmed). thus, the methods and techniques that work well for the general domain are not guaranteed to work the same for the biomedical domain. figure 1 illustrates the difference in the raw counts of word frequencies between wikipedia and pubmed, which contains only abstracts (pubmed corpus). the corpora are of similar sizes (3.1 and 3.6 billion tokens, correspondingly). we can observe that, for example, the count of rare words noticeably differs between two corpora and twice less for the pubmed corpus. @xcite and @xcite showed how the presence of rare words in the pre-training corpus might greatly impact the performance of the lm. this work aims to compare different techniques and methods for creating biomedical lms. we compare such approaches as pre-training from scratch, continued pre-training with no specialized vocabulary, and continued pretraining with biomedical vocabulary. our evaluations show that a biomedical lm pre-trained in a continued setting speeds up the pre-training. to account for the lack of specialized vocabulary that inevitably accompanies models pre-trained following a continued approach, we propose a new method to initialize weights for new tokens by distilling existing weights from the bert model inside the context where they were found. furthermore, we compare how masking rate, corruption strategy, and masking strategies impact the performance of the biomedical lm. we observe that the corruption strategy of "80-10-10" @xcite may play a role in the representation degeneration problem of lm by increasing the degree of anisotropy in the contextualized word representations. with relation to the word representation, anisotropy characterizes an embedding space that is restricted into a narrow cone-like shape @xcite . this phenomenon leads to a decrease in word embeddings' expressiveness and is referred to as the repre-sentation degeneration problem @xcite . finally, based on the results of our experiments, we introduce a model which pre-trained with a cl schedule. cl is an easy-to-hard training strategy @xcite . though very successful and widely adopted in nlp, it is still under-explored in pre-training lms. to the best of our knowledge, cl was not used before to pre-train the biomedical lm. we propose the cl method based on the task complexity of predicting masked labels. our model (bioptimus), pre-trained with the contextualized weight distillation and following our cl method, sets new states of the art on several biomedical named entity recognition (ner) tasks. in summary, our contributions are: in addition, we experiment with masking ratio, corruption strategy, and masking strategies, testing what works more optimally to prepare a language model for the biomedical domain. ## prior literature recent studies indicate that adapting existing lms pre-trained on general corpora for a new domain or pre-training a new lm from scratch for a specific domain gives a substantial gain in terms of performance for a downstream task. the authors of biobert @xcite performed domain adaptation by continuing pre-training bert @xcite on a large corpus of biomedical data. the model showed notable achievement on downstream tasks, outperforming bert and the state-of-the-art models. @xcite further pre-train biobert using clinical corpora to address clinical text's challenges and released clin-icalbert. @xcite showed the benefits of continued domain pre-training followed by task-adaptive pre-training (tapt) of roberta on four domains, including biomedical. one of the downsides of continued pre-training is the lack of domain-specific vocabulary. the authors of the scibert model @xcite constructed a new scivocab using the sentencepiece library, introducing 58% of new tokens, which shows how much scientific vocabulary differs from a general vocabulary of bert-base. the pre-training corpus comprises 82% of the biomedical domain; the rest is the computer science domain. training from scratch solely on the biomedical domain showed further gain and was presented by the authors of pubmedbert @xcite . they construct a new domain-specific vocabulary, proving its effectiveness when fine-tuning for downstream tasks. another approach to pre-train the biomedical lm from scratch using citation links is the linkbert model, which achieved new state-of-the-art on several bionlp tasks @xcite . to mitigate the high costs of pre-training from scratch, tai et al. (2020) introduced extension vocabulary that is not found in the vocabulary of bert. the weights of the embedding layer of the extension vocabulary are randomly initialized, then further pre-trained on the biomedical corpus. to avoid randomly initializing weights for new tokens and speed up pre-training, poerner et al. (2020) used word2vec. word2vec vectors are trained on the biomedical domain and then aligned with wordpiece vectors from bert. sachidananda et al. (2021) experimented with subword-based initialization using the mean of roberta fixed subword embeddings. curriculum learning curriculum learning @xcite has been successfully explored in nlp to solve different tasks. it was widely adopted in machine translation @xcite @xcite . some works have been done in the area of answer generation @xcite , relation extraction @xcite , reading comprehension ( tay et al., 2019 ## method in continued biomedical domain adaptation, pretraining using the checkpoint of the model pre-trained on the general domain can significantly speed up training. moreover, the benefits of continued domain adaptation are highly aligned with the present trends of nlp communities @xcite of reducing the environmental impact and making knowledge readily available in limited resources settings. one of the problems of continued pre-training of lms for domain adaptation is, loosely speaking, the absence of domainspecific vocabulary @xcite . the issue of out-of-vocabulary (oov) words is more elegantly solved with subword tokenization such as byte-pair encoding (bpe) @xcite , wordpiece @xcite , and sentencepiece @xcite . nevertheless, the presence of a whole word or more meaningful subwords in domain-specific vocabulary greatly enhances performance on downstream tasks @xcite @xcite . there are ways to address this issue by incorporating domain-specific tokens and assigning new token weights without initializing them from scratch @xcite . nevertheless, none of the current approaches leverage the advantage of contextualized embeddings of transformer models. to boost knowledge transfer, we introduce the new vocabulary by distilling existing weights from the bert model inside the context of the domain where they were found. our approach uses wordpiece tokenization @xcite to tokenize pubmed corpus foot_1 (biomedtokenizer) with a vocabulary size of 30522. we contextualize tokens in the following way (see figure 2 ): in that case, the token will be broken into six pieces: ['bro','##nch','##oco','##nst','##ric', '##tion'] (see figure 2 ). we use the mean operation of constituting tokens to compute a single representation for a new token (distilled representations): the variable k represents the number of the subtokens that make up the token of interest. we average token weights of only the last layer of the bert model as it is more expressive of context and domain information @xcite . the size of sentences sampled per token m may vary. we sample uniformly at random, setting the upper bound equal to 20 and the lower bound to 1. in case the token is not found in the corpus (such tokens represent less than 1/10 of the whole vocabulary), we assign distilled representation to this token. we use mean operation, which was found to be the most efficient operation for aggregating across several contexts @xcite . the averaging of the weights of subtokens constitutes more meaning when placed in context. moreover, contextualization of tokens' embeddings leverages the ability of the bert model to create domain-specific word embeddings that align with the corpus where they were found. @xcite . we call this approach contextualized weight distillation. ## models to perform a more rigorous comparison and find an optimal approach to the pre-train biomedical lm, we pre-train four models (see ## experimental set-up we pre-train all models with the same hyperparameters (details can be found in appendix a). we further fine-tune each model for five ner datasets and report results in moreover, we experiment with other pre-training techniques. the above-described models are pretrained by masking random tokens. @xcite , and @xcite showed that pretraining with whole-word masking (wwm) can be more beneficial. @xcite experimented with different percentages of masking rate and found out that for the bert-base, 0.2 percent may be more optimal to prompt the model to learn better. thus, we pre-train two more models using our method but with different techniques, one model using wwm and the second using wwm and 0.2% masking rate. results are presented in ## analysis figure 3 shows the models' performance on the dev set after each epoch. we can see that bm continued learns quickly and performs better on jnlpba and bc5-chem datasets. however, the model shows the lowest learning curve on the ncbi-disease dataset, which the absence of specialized vocabulary can explain. figure 3 also illustrates the importance of domain-specific vocabulary, and we can see that all the models with specialized vocabulary learn quickly and perform relatively well. overall, introducing domain-specific vocabulary is beneficial. however, it is less efficient to pre-train a new model entirely from scratch, initializing all weights randomly. in the case of bioptimus 0.1, adding new domain-specific tokens alongside the corresponding weights initialized in a way that leverages the contextualizing ability of transformers gives noticeable gains. additionally, it is important to compare how biomedical lms respond to pre-training on wwm against pre-training with masking only tokens/subtokens. ## conclusion this paper presented a new method to initialize tokens' weight for new biomedical vocabulary when pre-training from the existing checkpoint (continued approach). we also compared this method of token weight initialization with other pre-training methods (see comparing pre-training techniques showed that wwm is not the best-performing approach for all ner tasks, and masking only tokens/subtokens shows competitive performance. increasing the masking rate and removing the corruption strategy are generally beneficial techniques for pretraining biomedical lm. finally, we introduced the cl method based on the task complexity to pre-train lms. the "easy-to-hard" cl method introduces the biomedical lm to a broader scope of language experience, speeds up pre-training, and enhances performance on downstream tasks like ner. it is important to highlight that our model bioptimus 0.4 achieves high performance with the pre-training time reduced by at least half, proving the pre-training approach's efficiency.