E ntity BERT : Entity-centric Masking Strategy for Model Pretraining for the Clinical Domain (2021)

Abstract: Transformer-based neural language models have led to breakthroughs for a variety of natural language processing (NLP) tasks. However, most models are pretrained on general domain data. We propose a methodology to produce a model focused on the clinical domain: continued pretraining of a model with a broad representation of biomedical terminology (PubMedBERT) on a clinical corpus along with a novel entity-centric masking strategy to infuse domain knowledge in the learning process. We show that such a model achieves superior results on clinical extraction tasks by comparing our entity-centric masking strategy with classic random masking on three clinical NLP tasks: cross-domain negation detection, document time relation (DocTimeRel) classification, and temporal relation extraction. We also evaluate our models on the PubMedQA dataset to measure the models’ performance on a non-entity-centric task in the biomedical domain. The language addressed in this work is English.

## introduction transformer-based neural language models, such as bert @xcite , have achieved state-of-the-art performance for a variety of natural language processing (nlp) tasks. since most are pre-trained on large general domain corpora, many efforts have been made to continue pretaining general-domain language models on clinical/biomedical corpora to derive domain-specific language models @xcite @xcite . yet, as @xcite pointed out, in specialized domains such as biomedicine, continued pretraining from generic language models is inferior to domain-specific pretraining from scratch. continued pre-training from a generic model would break down many of the domain specific terms into sub-words through the byte-pair encoding (bpe) @xcite or variants like wordpiece tokenization @xcite because these specific terms are not in the vocabulary of the generic pretrained model. a clinical domain-specific pretraining from scratch would derive an in-domain vocabulary as many of the biomedical terms, such as diseases, signs/symptoms, medications, anatomical sites, procedures, would be represented in their original form. such an improved word-level representation is expected to bring substantial performance gains in clinical domain tasks because the model would learn the characteristics of the term along with its surrounding context as one unit. in our preliminary work on a clinical relation extraction task, we observed a performance gain with the pubmedbert model @xcite which outperformed biobert @xcite , clinicalbert @xcite , and even some larger general domain models like roberta @xcite and bart-large @xcite . the performance gain was primarily attributed to pubmedbert's in-domain vocabulary as we observed that pubmedbert kept 30% more in-domain words in its vocabulary than bert. when we swapped pubmedbert tokenization with bert or roberta tokenization, the performance of pubmedbert degraded. thus, pubmedbert appears to provide a vocabulary that is helpful to the clinical domain. however, the language of biomedical literature is different from the language of the clinical documents found in electronic medical records (emrs). in general, a clinical document is written by physicians who have very limited time to express the numerous details of a patient-physician encounter. many nonstandard expressions, abbreviations, assumptions and domain knowledge are used in clinical notes which makes the text hard to understand outside of the clinical community and presents challenges for automated systems. pretraining a language model specific to the clinical domain requires large amounts of unlabeled clinical text on par with what the generic models are trained on. unfortunately, such data are not available to the community. the only available such corpus is mimic iii used to train clinicalbert @xcite and bluebert @xcite , but it is magnitudes smaller and represents one specialty in medicine -intensive care. pretraining is agnostic to downstream tasks: it learns representations for all words using a selfsupervised data-rich task. yet, not all words are important for downstream fine-tuning tasks. numerous pretrained words are not even used in the fine-tuning step, while important words crucial for the downstream task are not well represented due to insufficient amounts of labeled data. many clinical nlp tasks are centered around entities: clinical named entity recognition aims to detect clinical entities @xcite @xcite , clinical negation extraction decides if a certain clinical entity is negated @xcite @xcite , clinical relation discovery extracts relations among clinical entities @xcite , etc. though various masking strategies have been employed during pretraining -masking contiguous spans of text @xcite , varying masking ratios @xcite , building additional neural models to predict which words to mask @xcite , incorporating knowledge graphs @xcite , masking entities for a named entity recognition task @xcite ) -none of the masking techniques so far have investigated and focused on clinical entities. besides transformer-based models, there are other efforts @xcite to characterize the biomedical/clinical entities at the word embedding level. there are also other statistical methods applied to the downstream tasks. we do not include these efforts in our discussion because the focus of our paper is the investigation of a novel entity-based masking strategy in a transformer-based setting. in this paper, we propose a methodology to produce a model focused on clinical entities: continued pretraining of a model with a broad representation of biomedical terminology (the pubmedbert model) on a clinical corpus, along with a novel entity-centric masking strategy to infuse domain knowledge in the learning process 1 . we show that such a model achieves superior results on clinical extraction tasks by comparing our entity-centric masking strategy with classic random masking on three clinical nlp tasks: cross-domain negation detetction @xcite , document time relation (doctimerel) classification @xcite , and temporal relation extraction @xcite . the contributions of this paper are: (1) a continued pretraining methodology for clinical domain specific neural language models, (2) a novel entitycentric masking strategy to infuse domain specific knowledge, (3) evaluation of the proposed strategies on three clinical tasks: cross-domain negation detection, doctimerel classification, and temporal relation extraction, and (4) evaluation of our models on the pubmedqa @xcite dataset to measure the models' performance on a non-entitycentric task in the biomedical domain. ## methods in this section, we first describe our clinical text datasets and related nlp tasks, the details of our entity-centric masking strategy, and finally the settings we used for both pretraining and fine-tuning. ## results and words are 40% and 12%, respectively. the table shows only the most successful rates; we considered more entity rates (20%, 40%, 60%, 80%, 100%) and word rates (0%, 8%, 10%, 12%, 14%, 16%). we found that (1) masking non-entity words in addition to masking entities is important as nonentity words capture semantic/syntactic information, and (2) masking too many tokens may make the reconstruction task too hard. for our experiments on the downstream tasks, we choose the entitybert model continuously pretrained on mimic-small with random seed 42 (0.651 f1) and the randmask model continuously pretrained on mimic-big with random seed 12 (0.641 f1) because of their best performance. for randmask models that all get 0.641 f1, we pick the one continuously pretrained on mimic-big. we fine-tuned them for the specific tasks. the detailed model performance on all tlink categories is in the bottom two rows in ## discussion the benefit of an in-domain vocabulary. to study the in-domain vocabulary's contribution to a clinical task, we extract all 3,471 gold standard events in the thyme colon cancer training set and feed them into the pubmedbert, roberta, and bert tokenizers. these events are all singletoken events. figure 4 shows the histogram of tokens per event after tokenization (x-axis shows the number of tokens each event is represented by). we see that pubmedbert keeps the majority of the events (2,330) as one unit instead of breaking them into multiple sub-words. the bert tokenizer keeps 1,729 events as one unit. the 601 events that pubmedbert recognizes but bert breaks into word pieces are of importance for the tlink task. if we remove these 601 events from the pub-medbert vocabulary -forcing them to be broken down into word pieces -the model performance on the tlink task drops from 0.638 f1 ( what makes a difference? by comparing the tlink predictions (without applying temporal closure) produced by the best entitybert (0.651 f1) and by the best randmask (0.641 f1), we found that entitybert predicted 4,924 correct tlinks, while randmask predicted 4,778 correct tlinks ( mask. some entities only appear in entitybertidentified relations, e.g. staging, hemoglobin, finding, consideration, consider, develops, request, treatment, neuropathy, carcinoma, metastasis, injection, resected, and staged are involved in multiple relations. entity-centric masking masks more entities than random masking so that those clinical entities can be better represented by the language model in terms of their semantic and syntactic usage. when the model is fine-tuned for an entity in figure 5 we visualize with bertviz (vig, 2019) the attention weights of head zero from the last layer of the fine-tuned randmask and entitybert models on the tlink task for a relation that enti-tybert correctly predicted but randmask missed. the context is he has had @xmath4. a plausible explanation is that because the key entities, injection and date, are not well represented in randmask model, the [cls] token of randmask model (figure 5 (a) ) focuses on entity markers, <e>, </e>, <t>, @xmath5. it may figure out this is an event-time relation but incorrectly infers its type. the [cls] token of en-titybert (figure 5 (b) ) bakes in representations of all tokens with knowledge that injection is related to steroid and date is related @xmath6, which shows the key entities are well represented. combining entity-centric masking with another masking strategy. some other pre-trained language models like bart @xcite and spanbert @xcite are not pretrained on clinical/biomedical corpora. yet, they are suitable for clinical tasks in that they mask contiguous random spans instead of individual tokens/word pieces during pretraining -in the clinical domain there are a lot of events and entities that span multiple tokens (e.g., ascending colon cancer, march 11, 2011). even without any continued pretraining on a clinical corpus, bart-large achieves 0.628 f1 on the tlink task ( the strength and limitations of entity-bert: entitybert assumes that clinical entities are important words, thus if a clinical language model can represent clinical entities better, it will benefit downstream clinical entity-centric tasks. therefore, such a masking strategy increases the entity concentrations in the masked words during the model pretraining, but does not increase the overall computational loads either for pre-training or for fine-tuning since the overall total number of masked items is similar to random word masking. this is unlike building an additional neural network for selective masking @xcite or incorporating knowledge graphs @xcite . the better representation of clinical entities is not only beneficial in an in-domain setting, e.g., the tlink task, but also effective in a cross-domain setting, e.g., the negation and doctimerel tasks. for the doctimerel task, both entitybert and randmask achieve very good in-domain performance of 0.92 f1 (see we acknowledge some limitations of the current entitybert model. first, it is pretrained with a relatively small block size (100 tokens) which is sufficient for a sentence-or a short-paragraphlevel reasoning tasks but may be not sufficient for document-level tasks. second, entitybert aims to improve the performance of entity-centric clinical tasks. for tasks that may not directly leverage entities, such as question answering or document classification, entities may still play a supporting role but may not prove as effective. however, we hypothesize that even in those cases its performance would be on-par with randmask because of its in-domain vocabulary and continued training on a clinical corpus. based on the results of mimic-big vs. mimic-small: rand-mask and entitybert models pretrained on mimic-small perform almost on par with models pretrained on the much bigger corpus, mimic-big ( mimic-big was created by filtering sentences with fewer than two entities with the goal of capturing pair-wise interactions between events in the language model. one of the limitations of our architecture is its block size. perhaps with a model that can effectively represent the relative distances, the interactions among entities can be represented better. in addition, by eliminating sentences that only have one or no entity, mimic-big misses some language phenomena. mimic-small, despite its smaller size, thus may encounter more diverse language. this could be the explanation of why an entitybert model pretrained on mimic-small gets the best tlink performance (0.651 f1; in the future, we will investigate combining entity-centric masking with deberta @xcite with the goal of developing a strategy for a deep neural model that combines entities and their relative position in an input sequence. we will experiment with more flavors of entitybert with different block sizes for a wider range of clinical applications. further testing entitybert on a wider range of clinical and biomedical tasks would be helpful for understanding its capabilities.