SlowGuess/NAACL_2021_OCR huggingface dataset:
George Michalopoulos, Yuanxin Wang, Hussam Kaka, Helen Chen, Alexander Wong
University of Waterloo,
{gmichalo, yuanxin.wang, hussam.kaka, helen.chen, alexander.wong} @uwaterloo.ca
Contextual word embedding models, such as BioBERT and Bio_ClinicalBERT, have achieved state-of-the-art results in biomedical natural language processing tasks by focusing their pre-training process on domain-specific corpora. However, such models do not take into consideration structured expert domain knowledge from a knowledge base.
We introduce UmsBERT, a contextual embedding model that integrates domain knowledge during the pre-training process via a novel knowledge augmentation strategy. More specifically, the augmentation on UmsBERT with the Unified Medical Language System (UMLS) Metathesaurus is performed in two ways: (i) connecting words that have the same underlying 'concept' in UMLS and (ii) leveraging semantic type knowledge in UMLS to create clinically meaningful input embeddings. By applying these two strategies, UmsBERT can encode clinical domain knowledge into word embeddings and outperform existing domain-specific models on common namedentity recognition (NER) and clinical natural language inference tasks.
In recent years, the volume of data being collected in healthcare has grown considerably. A significant proportion of the data is in text form, which requires advanced Natural Language Processing (NLP) models to process. This has led to the creation of high-performing, optimized NLP models focused on the biomedical domain.
Contextual word embedding models, such as ELMo (Peters et al., 2018) and BERT (Devlin et al., 2019) have achieved state-of-the-art results in many NLP tasks. Initially tested in a general domain, these models have also been successfully applied in the biomedical domain by pre-training them on biomedical corpora, leading to the best performances in a variety of biomedical NLP tasks (Lee
et al., 2019), (Alsentzer et al., 2019). However, current biomedical applications of transformer-based Natural Language Understanding (NLU) models do not incorporate structured expert domain knowledge from a knowledge base into their embedding pre-training process.
The Unified Medical Language System (UMLS) (Bodenreider, 2004) Metathesaurus is a compendium of many biomedical terminologies with the associated information, such as synonyms and categorical groupings. It allows for the connection of words that represent the same or similar 'concept'. For example, the words 'lungs' and 'pulmonary' share a similar meaning and thus can be mapped to the same concept unique identifier (CUI) CUI: C0024109. Additionally, UMLS allows the grouping of concepts according to their semantic type (McCray et al., 2001). For example, 'skeleton' and 'skin' have the same 'Body System' semantic type, and 'inflammation' and 'bleed' are in the 'Pathologic Function' semantic type.
In this paper, we present and publicly release<sup>1</sup> a novel architecture for augmenting contextual embeddings with clinical domain knowledge. Specifically: (i) We are the first, to the best of our knowledge, to propose the usage of domain (clinical) knowledge from a clinical Metathesaurus (UMLS Metathesaurus) in the pre-training phase of a BERT-based model (UmlsBERT) in order to build 'semantically enriched' contextual representations that will benefit from both the contextual learning (BERT architecture) and the domain knowledge (UMLS Metathesaurus). (ii) We propose a new multi-label loss function for the pre-training of the Masked Language Modelling (Masked LM) task in the UmlsBERT that incorporates the connections between clinical words using the CUI attribute of UMLS. (iii) We introduce a semantic type embedding that enriches the input embeddings process of the UmlsBERT by forcing the model to take into
$^{1}$ https://github.com/gmichalo/UmlsBERT
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1744-1753 June 6-11, 2021. ©2021 Association for Computational Linguistics
consideration the association between words that are of the same semantic type. (iv) Finally, we demonstrate that UmlsBERT outperforms two popular clinical-based BERT models (BioBERT and Bio_ClinicalBERT) and a general domain BERT model on different clinical named-entity recognition (NER) tasks and on one clinical natural language inference task.
The rest of paper is organized as follows. Related work is presented in Section 2. The data that were used to pre-train and test the new UmlsBERT are described in Section 3. The characteristics of the proposed UmlsBERT architecture for augmenting contextual embeddings with clinical knowledge are detailed in Section 4. Finally, the results of the down-stream tasks and the qualitative analysis are reported in Section 5, and a conclusion and a plan for future work are presented in Section 6.
In (Peters et al., 2018), contextualized word embeddings were introduced in a bidirectional language model (ELMo). This allowed the model to change the embedding of a word based on its imputed meaning, which was derived from the surrounding context. Subsequently, (Devlin et al., 2019) proposed the Bidirectional Encoder Representations from Transformers (BERT) which used bidirectional transformers (Vaswani et al., 2017) to create context-dependent representations. For both models, pre-training is done on massive corpora and the context-sensitive embeddings can be used for downstream tasks.
Other approaches enhance the BERT's performance by injecting external knowledge from a knowledge base. Sense-BERT (Levine et al., 2020) is pre-trained to predict the supersenses (semantic class) of each word by incorporating lexical semantics (from the lexical database WordNet Miller, 1995)) into the model's pre-training objective and by adding supersense information to the input embedding. In addition, GlossBERT (Huang et al., 2019) focuses on improving word sense disambiguation by using context-gloss pairs on the sentence-pair classification task of a BERT model.
Furthermore, there have been multiple attempts to improve the performance of contextual models in the biomedical domain. BioBERT is a BERT-based model which was pre-trained on both general (BooksCorpus and English Wikipedia) and biomedical corpora (PubMed abstracts and PubMed Cen
tral full-text articles) (Lee et al., 2019). The authors demonstrate that incorporating biomedical corpora in the pre-training process improves the performance of the model in downstream biomedical tasks. This is likely because medical corpora contains terms that are not usually found in a general domain corpus (Habibi et al., 2017). Finally, Bio_ClinicalBERT (Alsentzer et al., 2019) further pre-trains BioBERT on clinical text from the MIMIC-III v1.4 database (Johnson et al., 2016). It is shown that the usage of clinical specific contextual embeddings can be beneficial for the performance of a model on different clinical NLP downstream tasks.
We use the Multiparameter Intelligent Monitoring in Intensive Care III (MIMIC-III) dataset (Johnson et al., 2016) to pre-train the UmlsBERT model. MIMIC dataset consists of anonymized electronic medical records in English of over forty-thousand patients who were admitted to the intensive care units of the Beth Israel Deaconess Medical Center (Boston, MA, USA) between 2001 and 2012. In particular, UmlsBERT is trained on the NOTEEVENTS table, which contains 2,083,180 rows of clinical notes and test reports.
Dataset	Train	Dev	Test	C
MedNLi	11232	1395	14 22	3
i2b2 2006	44392	5547	18095	17
i2b2 2010	14504	1809	27624	7
i2b2 2012	6624	820	5664	13
i2b2 2014	45232	5648	32586	43
Table 1: Number of sentences for the train/dev/test set of each dataset. We also include the number of classes (C) for each dataset. We use the same splits that are used in (Alsentzer et al., 2019).
We evaluate the effects of the novel features of the UmlsBERT model on the English MedNLI natural language inference task (Romanov and Shivade, 2018) and on four i2b2 NER tasks (in IOB format (Ramshaw and Marcus, 1995)). More specifically, we experiment on the following English i2b2 tasks: the i2b2 2006 de-identification challenge (Uzuner et al., 2007), the i2b2 2010 concept extraction challenge (Uzuner et al., 2011), the i2b2 2012 entity extraction challenge (Uzuner et al., 2011) and the i2b2 2014 de-identification challenge (Stubbs et al., 2015). These datasets are chosen because of their
use in benchmarking prior biomedical BERT models, thereby allowing for performance comparison. In addition, these publicly available datasets enable the reproducibility of our results and meaningful comparison with future studies. Table 1 lists the statistics of all the datasets. Finally, it should be noted that for the identification of the UMLS terms, we use the UMLS 2020AA version.
4.1 BERT Model
The original BERT model (Devlin et al., 2019) is based on multi-layer bidirectional transformers (Vaswani et al., 2017), which generates contextualized word representations. Incorporating information from bidirectional representations allows the BERT model to capture more accurately the meaning of a word based on its surrounding context, i.e. sentence.
The pre-training phase of the BERT model consists of two self-supervised tasks: Masked Language Modelling (LM), in which a percentage of the input is masked at random and the model is forced to predict the masked tokens, and Next Sentence Prediction, in which the model has to determine whether two segments appear consecutively in the original text. Since our UmlsBERT model is focused on augmenting the Masked LM task with clinical information from the UMLS Metathesaurus, we omit the description of the Next Sentence Prediction task and only describe the details of the Masked LM task herein.
In Masked LM,  $15\%$  of the tokens of each sentence are replaced by a [MASK] token. For the  $j^{th}$  input token in the sentence, an input embedding vector  $u_{input}^{(j)}$  is created by the following equation:
$$
u _ {i n p u t} ^ {(j)} = p ^ {(j)} + S E G s e g _ {i d} ^ {(j)} + E w _ {j} \tag {1}
$$
where  $p^{(j)}\in \mathbb{R}^d$  is the position embedding of the  $j^{th}$  token in the sentence, and  $d$  is the transformer's hidden dimension. Additionally,  $SEG\in \mathbb{R}^{d\times 2}$  is called the segment embedding, and  $seg_{id}\in \mathbb{R}^2$ , a 1-hot vector, is the segment id that indicates the sentence to which the token belongs. In Masked LM, the model uses only one sentence and therefore, the segment id indicates that all the tokens belong to the first sentence.  $E\in \mathbb{R}^{d\times D}$  is the token embedding where  $D$  is the length of the model's vocabulary and  $w_{j}\in \mathbb{R}^{D}$  is a 1-hot vector corresponding to the  $j^{th}$  input token.
The input embedding vectors pass through multiple attention-based transformer layers where each layer produces a contextualized embedding of each token. Finally, for each masked token  $w$ , the model outputs a score vector  $y_w \in \mathbb{R}^D$  with the goal of minimizing the cross-entropy loss between the softmax of  $y_w$  and the 1-hot vector corresponding to the masked token  $(h_w)$ :
$$
l o s s = - l o g \left(\frac {e x p \left(y _ {w} [ w ]\right)}{\sum_ {w ^ {\prime}} e x p \left(y _ {w} [ w ^ {\prime} ]\right)}\right) \tag {2}
$$
4.2 Enhancing Contextual Embeddings with Clinical Knowledge
In the UmlsBERT model, we update the Masked LM procedure to take into consideration the associations between the words specified in the UMLS Metathesaurus.
4.2.1 Semantic type embeddings
We introduce a new embedding matrix called  $ST \in \mathbb{R}^{D_s \times d}$  into the input embedding of the BERT model, where  $d$  is BERT's transformer hidden dimension and  $D_s = 44$  is the number of unique UMLS semantic types that can be identified in the vocabulary of our model. In particular, in this matrix, each row represents the unique semantic type in UMLS that a word can be identified with (for example the word 'heart' is associated with the semantic type T023: 'Body Part, Organ, or Organ Component' in UMLS).
To incorporate the  $ST$  embedding matrix into the input embedding of our model, all words with a clinical meaning defined in UMLS are identified. For each of these words, the corresponding concept unique identifier (CUI) and semantic type are extracted. We use  $s_w \in \mathbb{R}^{D_s}$  as a 1-hot vector corresponding to the semantic type of the medical word  $w$ . The identification of the UMLS terms and their UMLS semantic type is accomplished using the open-source Apache clinical Text Analysis and Knowledge Extraction System (cTakes) (Savova et al., 2010). Thus, by introducing the semantic type embedding, the input vector (equation 1) for each word is updated to:
$$
u _ {i n p u t} ^ {(j) \prime} = u _ {i n p u t} ^ {(j)} + S T ^ {\top} s _ {w} \tag {3}
$$
where the semantic type vector  $ST^{\top}s_{w}$  is set to a zero-filled vector for words that are not identified in UMLS.
We hypothesize that incorporating the clinical information of the semantic types into the input
Figure 1: (a) Original input vector of the BERT model (Devlin et al., 2019). (b) Augmented input vector of the UmlsBERT where the semantic type embeddings is available. For the words 'chest' and 'cavity', their word embeddings are enhanced with the embedding of the semantic type 'Body Location or Region'(  $E_{T029}$ ) and 'Body Space or Junction'(  $E_{T030}$ ) respectively. The rest of the words are not related to a medical term, so a zero-filled tensor  $E_{null}$  is used.
tensor could be beneficial for the performance of the model as the semantic type representation can be used to enrich the input vector of words that are rare in the training corpus and the model do not have the chance to learn meaningful information for their representation. Figure 1 presents an overview of the insertion of the semantic type embeddings into the standard BERT architecture.
4.2.2 Updating the loss function of Masked LM task
Furthermore, we update the loss function of the Masked LM pre-training task to take into consideration the connection between words that share the same CUI. As described in Subsection 4.1, the loss function of the Masked LM pre-training task of a BERT model is a cross-entropy loss between the softmax vector of the masked word and the 1-hot vector that indicates the actual masked word. We proposed to 'soften' the loss function and updated it to a multi-label scenario by using information from the CUIs.
More specifically, instead of using a 1-hot vector  $(h_w)$  that corresponds only to the masked word  $w$ , we use a binary vector indicating the presence of all the words which share the same CUI of the masked word  $(h_w')$ . Finally, in order for the model to properly function in a multi-label scenario, the cross entropy loss (equation 2) is updated to a binary cross entropy loss:
$$
\begin{array}{l} l o s s = \sum_ {i = 0} ^ {D} \left(- h _ {w} ^ {\prime} [ i ] l o g \left(y _ {w} [ i ]\right) \right. \\ + \left(1 - h _ {w} ^ {\prime} [ i ]\right) \log \left(1 - y _ {w} [ i ]\right)) \tag {4} \\ \end{array}
$$
These changes force UmlsBERT to learn the
semantic relations between words, which are associated with the same CUI in a biomedical context.
An example of predicting the masked word 'lungs' with and without the clinical information is presented in Figure 2. As seen in this figure, the UmlsBERT model tries to identify the words 'lung', 'lungs' and 'pulmonary' because all three words are associated with the same CUI: C0024109 in the UMLS Metathesaurus.
(b)
Figure 2: An example of predicting the masked word 'lungs' (a) the BERT model tries to predict only the word lungs (b) whereas the UmlsBERT tries to identify all words that are associated with the same CUI (e.g. lungs, lung, pulmonary).
4.3 UmlsBERT Training
We initialize UmlsBERT with the pre-trained Bio_ClinicalBERT model (Alsentzer et al., 2019), and then we further pre-train it with the updated Masked LM task on MIMIC-III notes. Afterwards, in order to perform the downstream tasks, we add a single linear layer on top of UmlsBERT and 'finetuned' it to the task at hand, using either the associated embedding for each token or the embedding of the [CLS] token. The same fine-tuning method is applied to all other models used for comparison. In order to keep the experiment controlled, we use
Dataset	BERTbased	BioBERT	Bio_ClinicalBERT	UmlsBERT
MedNLI	epochs	4	4	4	3
batch size	16	16	32	16
learning rate	5e-5	3e-5	3e-5	3e-5
i2b2 2006	epochs	20	20	20	20
batch size	32	16	16	32
learning rate	2e-5	2e-5	2e-5	5e-5
i2b2 2010	epochs	20	20	20	20
batch size	16	32	32	16
learning rate	3e-5	3e-5	5e-5	5e-5
i2b2 2012	epochs	20	20	20	20
batch size	16	32	16	16
learning rate	3e-5	3e-5	5e-5	5e-5
i2b2 2014	epochs	20	20	20	20
batch size	16	16	32	16
learning rate	2e-5	2e-5	5e-5	3e-5
Table 2: Hyperparameter selection of all the models for each dataset
the same vocabulary and WordPiece tokenization (Wu et al., 2016) across all the models. WordPiece divides words not in the vocabulary into frequent sub-words.
Since our goal is to demonstrate the beneficial effect of incorporating domain knowledge in this study, we haven't experimented with a more complicated layer on top of UmlsBERT (e.g. the BiLSTM layer in (Si et al., 2019)). This is because our goal is to demonstrate that incorporating domain knowledge was beneficial for the performance of the model by showing that UmlsBert outperformed the other medical-based BERT models on a variety of medical NLP tasks (Section 5). It should be noted that we chose the UMLS Metathesaurus in our process of augmenting the UmlsBERT model for two reasons:
1. We aim to create a clinical contextual embedding model that is capable of integrating domain (medical) knowledge.
2. The UMLS Metathesaurus is a compendium of many popular biomedical vocabularies (e.g. MeSH (Dhammi and Kumar, 2014) and ICD-10 (Organization, 2004)). By choosing to utilize the domain (medical) knowledge of UMLS, we actually incorporate domain knowledge from all major internationally standardized clinical terminologies.
In the pre-training phase, UmlsBERT is trained for 1,000,000 steps with a batch size of 64, maximum sequence length of 128 and learning rate of
$5 \cdot 10^{-5}$ . All other hyper-parameters are kept to their default values. UmlsBERT is trained by using 2 nVidia V100 16GB GPU's with 128 GB of system RAM running Ubuntu 18.04.3 LTS.
In this section, we present the results of an empirical evaluation of the UmlBERT model. In particular, we provide a comparison between different available BERT models to show the efficiency of our proposed model on different clinical NLP tasks. In addition, we provide the results of an ablation test to exam the effect of the semantic type embeddings on the performance of the model. Furthermore, we conduct a qualitative analysis of the embedding of each model in order to illustrate how medical knowledge improves the quality of medical embeddings. Finally, we provide a visualized comparison of the embeddings of the words that are associated with semantic types between UmlBERT and Bio_ClinicalBert.
5.1 Downstream Clinical NLP Tasks
In this section, we report the results of the comparison of our proposed UmlsBERT model with the other BERT-based models on different downstream clinical NLP tasks described in Section 3. All BERT-based models are implemented using the transformers library (Wolf et al., 2019) on PyTorch 0.4.1. All experiments are executed on a Tesla P100 16.3 GB GPU with 32G GB of system RAM on Ubuntu 18.04.3 LTS.
Dataset	BERTbased	BioBERT	Bio_ClinicalBERT	UmlsBERT
MedNLI	Test Ac.	77.9 ± 0.6	82.2 ±0.5	81.2 ± 0.8	83.0 ± 0.1
Val. Ac.	79.0 ± 0.5	83.2 ± 0.8	83.4 ± 0.9	84.5 ± 0.1
Run. time(sec)	308	307	269	305
#parameters	108,312,579	108,312,579	108,312,579	108,346,371
i2b2 2006	Test F1	93.5 ± 1.4	93.3 ± 1.3	93.1 ± 1.3	93.6 ± 0.5
Val. F1	94.2 ± 0.6	93.8 ± 0.3	93.4 ± 0.2	94.4 ± 0.2
Run. time(sec)	12508	12807	12729	13167
#parameters	108,322,576	108,322,576	108,322,576	108,356,368
i2b2 2010	Test F1	85.2 ± 0.2	87.3 ± 0.1	87.7 ± 0.2	88.6 ± 0.1
Val. F1	83.4 ± 0.3	85.2 ± 0.6	86.2 ± 0.2	87.7 ± 0.5
Run. time(sec)	5325	5244	5279	5219
#parameters	108,315,655	108,315,655	108,315,655	108,349,447
i2b2 2012	Test F1	76.5 ± 0.2	77.8 ± 0.2	78.9 ± 0.1	79.4 ± 0.1
Val. F1	76.2 ± 0.7	78.1 ± 0.5	77.1 ± 0.4	78.3 ± 0.4
Run. time(sec)	2413	2387	2403	2432
#parameters	108,320,269	108,320,269	108,320,269	108,354,061
i2b2 2014	Test F1	95.2 ± 0.1	94.6 ± 0.2	94.3 ± 0.2	94.9 ± 0.1
Val. F1	94.5 ± 0.4	93.9 ± 0.5	93.0 ± 0.3	94.3 ± 0.5
Run. time(sec)	16738	17079	16643	16554
#parameters	108,343,339	108,343,339	108,343,339	108,377,131
Table 3: Results of mean  $\pm$  standard deviation of five runs from each model on the test and the validation test; we use the acronym Ac. for accuracy; average running time and number of parameters is also provided for each model. The number of parameters is different between datasets as we included the linear layers that were used on top of the Bert-based model for text and token classification; best values are bolded;
5.1.1 Hyperparameter tuning
In order to address the reproducibility concerns of the NLP community (Dodge et al., 2019), we provide the search strategy and the bound for each hyperparameter as follows: the batch size is set between 32 and 64, and the learning rate is chosen between the values 2e-5, 3e-5 and 5e-5. For the clinical NER tasks, we take a similar approach to (Lee et al., 2019) and set the number of training epochs to 20 to allow for maximal performance, except for MedNLI, for which we train the models on 3 and 4 epochs.
The best values are chosen based on validation set F1 values using the seqevals python framework for sequence labeling evaluation, due to the fact that it can provide an evaluation of a NER task on entity-level² for the i2b2 tasks and validation set accuracy, which is the standard metric for this task³ for the MedNLI dataset. In the interest of providing a fair comparison, we also tune the hyperparameters of each model in order to demonstrate its best
performance. The final hyper-parameters selection of all the models for each dataset can be found in Table 2.
In order to achieve more robust results, we run our model on five different (random) seeds (6809, 36275, 5317, 82958, 25368) and we provide the average scores and standard deviation for the testing and the validation set. It should be noted that  $\mathrm{BERT}_{\text {base }}$ , BioBERT and Bio_ClinicalBERT have the exact same number of parameters as they use the same BERT-based architecture. However, because we introduce the semantic type embeddings into the UmlsBERT model, our model has an additional 33792 [the number of unique UMLS semantic types (44) × transformer's hidden dimension(768)] parameters<sup>4</sup>. In Table 3, we provide the number of parameters for each dataset where we include the linear layer on top of the BERT-based models for the text and token classification.
$^{2}$ https://github.com/chakki-works/seqeval
 $^{3}$ https://tinyurl.com/transformers-metrics
${}^{4}$  UmlsBert also contains an additional zero-filled vector, that we use as the semantic type vector of the words that are not identified in UMLS, which was not included in the calculation of the number of the parameters of the model.
ANATOMY	DISORDER	GENERIC
feet	kidney	mass	bleeding	school	war
BERTbased	ft	liver	masses	bleed	college	battle
foot	lung	massive	sweating	university	conflict
BioBERT	foot	liver	masses	bleed	college	wartime
wrist	lung	weight	strokes	schooling	battle
Bio_ClinicalBERT	foot	liver	masses	bleed	college	warfare
legs	lung	weight	bloody	university	wartime
UmlSBERT	foot	Ren	lump	bleed	college	warfare
pedal	liver	masses	hem	students	military
5.1.2 BERT-based model comparison
The mean and standard deviation (SD) of the scores for all the competing models on different NLP tasks are reported in Table 3. UmlsBERT achieves the best results in 4 out of the 5 tasks. It achieves the best F1 score in three i2b2 tasks (2006, 2010 and 2012)  $(93.6\%, 88.6\%$  and  $79.4\%)$  and the best accuracy on the MedNLI task  $(83.0\%)$ .