UMLS-KGI-BERT: Data-Centric Knowledge Integration in Transformers for Biomedical Entity Recognition
Year: 2023

Abstract:
Pre-trained transformer language models (LMs) have in recent years become the dominant paradigm in applied NLP. These models have achieved state-of-the-art performance on tasks such as information extraction, question answering, sentiment analysis, document classification and many others. In the biomedical domain, significant progress has been made in adapting this paradigm to NLP tasks that require the integration of domain-specific knowledge as well as statistical modelling of language. In particular, research in this area has focused on the question of how best to construct LMs that take into account not only the patterns of token distribution in medical text, but also the wealth of structured information contained in terminology resources such as the UMLS. This work contributes a data-centric paradigm for enriching the language representations of biomedical transformer-encoder LMs by extracting text sequences from the UMLS. This allows for graph-based learning objectives to be combined with masked-language pre-training. Preliminary results from experiments in the extension of pre-trained LMs as well as training from scratch show that this framework improves downstream performance on multiple biomedical and clinical Named Entity Recognition (NER) tasks. All pre-trained models, data processing pipelines and evaluation scripts will be made publicly available.

heading: 1 Introduction
text: In recent times, transformer language models (Vaswani et al., 2017) have become the most popular and effective sequence modelling framework in almost all areas of applied Natural Language Processing. Unsupervised pre-training on large quantities of text allows transformers to capture rich semantic and syntactic patterns that can be transferred to many specialised language processing objectives. As such, transformer models that use the transfer learning paradigm whereby the model is trained
in an unsupervised manner on a large text corpus and then fine-tuned on a downstream supervisedlearning task have achieved state-of-the-art results across a wide range of general and domain-specific applications.
The proliferation of textual data in the biomedical domain (Electronic Health Records (EHRs), clinical documents, pharmaceutical specifications, etc) has precipitated the broad adoption of deep learning & NLP techniques for information extraction and processing (Li et al., 2021; Tiwari et al., 2020; Dubois et al., 2017). Moreover, it has been shown that language models are capable of encoding clinical knowledge to a certain extent (Singhal et al., 2022). Biomedical and clinical NLP, however, is widely recognised to present particular challenges that do not apply to the same extent in other domains, in particular the need to incorporate structured domain knowledge into text encodings (Chang et al., 2020). In order for neural language modelling to be reliable in a discipline as highly specialised as medicine, there is a more acute need for models to learn directly from domain-specific terminologies, as opposed to relying solely on corpus-based learning. Thus, a significant amount of research effort in the medical NLP community has been directed towards the question of how best to inject information from knowledge graphs (KGs) into LMs (He et al., 2022; Naseem et al., 2022; Li et al., 2020). However, a generalisable, widely-accepted approach to this technique that can be easily transferred across different problem settings, models and training corpora has yet to emerge. In addition, research into knowledge graph integration in NLP in the biomedical domain has tended to focus on English-language corpora; the utility and transferability of these techniques for other languages, for which less textual resources are available, as well as for multilingual models, remains therefore an under-explored area.
This paper aims to contribute to the resolution of
312
these issues by proposing a general framework for training BERT encoders (Devlin et al., 2019) using the UMLS (Unified Medical Language System, Bodenreider (2004)) alongside free-text corpora.
The main contributions of this work are as follows:
• We propose a data-centric method for formulating the KG-based learning objectives of triple classification and entity/link prediction in the language modelling paradigm, and implement a framework for training transformers using the UMLS knowledge base in parallel with masked-language pre-training.
• Pre-training on the UMLS alongside the European Clinical Case Corpus (Minard et al., 2021; Magnini et al., 2020), we show that this method brings improvements to pre-trained models across a range of biomedical entity recognition tasks in three different languages, as well as functioning as a competitive pre-training strategy that requires much less training data in comparison to state-ofthe-art transformer models. We release the monolingual and multilingual model weights trained in this way, UMLS-KGI-BERT, as open-source resources for the clinical NLP research community.
• Based on this work, we release the Python library bertify_umls, built mainly on the transformers and pandas libraries, which allows researchers to create custom text datasets and effectively use the UMLS knowledge base as a training corpus for BERT-style LMs.

heading: 2 Related Work

heading: 2.1 Pre-trained LMs for Medical Applications
text: In general, the standard methodology for adapting neural text encoders to the biomedical domain has been to take a model that has been pre-trained on general-domain text corpora and continue this unsupervised pre-training on a medical corpus (Alrowili and Shanker, 2021; Lee et al., 2020; Alsentzer et al., 2019). However, recent work has suggested that, given enough training data, it is preferable to pre-train these models on large domain-specific corpora only, without starting from a general-domain checkpoint (Gu et al., 2021; Rasmy et al., 2021). In this work we explore both approaches, extending
existing biomedical and general-domain models as well as training BERT models from scratch on our own generated datasets.

heading: 2.2 Knowledge-enhanced LMs
text: Techniques for the incorporation of knowledge graph structure into BERT models can, broadly speaking, be divided into three categories, each focusing on one of the three fundamental components of a machine learning system, i.e. 1) the training data, 2) the model architecture and 3) the objective function to be optimised. The first type of approach prioritises the augmentation of BERT’s input data with information extracted from a knowledge graph. This extra information can be numerical, e.g. precomputed graph embeddings (Jeong et al., 2019) or textual, e.g. KG triples linked to input sentences (Liu et al., 2019).
The second type of approach focuses on adapting the architecture of BERT so that its language representations become fused with knowledge graph embeddings (KGEs) (Wang et al., 2021; Peters et al., 2019; Zhang et al., 2019). Knowledge graph fusion techniques such as these have been shown to be beneficial on certain English-language medical NLP tasks (Meng et al., 2021; Roy and Pan, 2021).
Thirdly, the self-supervised pre-training objective of BERT models can be augmented using the kind of knowledge graph reasoning tasks used to build KGE models. This approach is more commonly used for knowledge graph completion (Kim et al., 2020; Yao et al., 2019) but has also been shown to be an effective strategy in the biomedical NLP domain (Hao et al., 2020).
As previously mentioned, given that the medical domain is particularly exacting in terms of requirements for the use of structured facts, the exploration of ways in which ontological knowledge can be integrated into automated text processing is a very active area of research (Khosla et al., 2020; Mondal et al., 2019). In particular, there have been multiple successful efforts to integrate the UMLS knowledge graph into BERT models, notably UmlsBERT (Michalopoulos et al., 2021), which proposes a data-augmentation technique allowing for concept and semantic type information to be linked to input text, and SapBERT (Liu et al., 2021b,a), which introduced a self-alignment strategy for learning from UMLS synonym pairs via a multi-similarity (MS) loss function to force related concepts closer to one another in BERT’s repre-
sentation space. Yuan et al. (2022) build on this strategy by applying MS loss to relation triples. In contrast, in this work we show that information from the UMLS can be incorporated into BERT models in a simpler way, using only cross-entropy classification loss, while also balancing this training process with standard masked-language BERT pre-training.
Recent general overviews of the landscape of AI research have highlighted the importance of data-centric approaches to building models (Zha et al., 2023; Hamid, 2022; Jakubik et al., 2022) and in light of these trends this work focuses on types 1) and 3) of knowledge base integration described above, i.e. on improving the performance of standard model architectures by constructing high-quality datasets that can be integrated into the self-supervised language modelling paradigm by modifying the BERT objective function. The motivation for this kind of approach is also to provide a pre-training framework that is more widely transferable and does not rely on any particular transformer-encoder architecture.

heading: 3 Methodology
text: In this work, we experiment with training BERT language models with three knowledge graph reasoning tasks derived from the UMLS, in addition to the standard masked-language modelling objective: entity prediction, link prediction and triple classification.

heading: 3.1 Dataset Construction
text: Formally, we consider the UMLS KG in the standard fashion, as a directed graph G = (C,E,R) where C is the set of all medical concepts in the KG, E the set of all edges or relations that link these concepts to one another, and R the set of possible relation types, i.e. the labels r for each e ∈ E. The training sequences are thus generated from the KG dataset of ordered triples (h, r, t) where (h, r) ∈ C × C and r ∈ R. As a compendium of multiple different sources of taxonomic biomedical information, the UMLS metathesaurus contains multiple levels of granularity at which meaning representation can be analysed. We consider three such levels of granularity in our work:
• Terms - string descriptors for conceptual enti-
ties
• CUIs (Concept Unique Identifiers) - the basic unit of meaning representation for the nodes in the knowledge graph, i.e. the elements of the set C.
• Semantic Groups - these are groupings of concepts that can be considered to define the type of entity a concept represents; e.g. anatomical structure, chemical, disorder etc.
Each concept (CUI) can be associated with multiple terms and multiple semantic groups. Thus, given that the entities h and t that make up the knowledge graph triples are represented as CUIs, in order to represent them as input text sequences for BERT models, we use the “preferred term” strings associated with the concepts h and t, except in the case of synonym relations where we randomly select another of the terms associated with the concept in question to associate with t. We also introduce a set of special tokens to represent the relation types R, of which there are seven (parent, child, synonymy, allowed qualifier, qualified by, broader, narrower). Concretely, the tokenization function for BERT models forms text classification sequences from triples in the following way;
Tokenize(h, r, t) = [CLS]wh1 · · ·whm [REL]wt1 · · ·wtn[SEP] (1)
where the wi represent the token sequences corresponding to the strings h and t, [CLS] and [SEP] are BERT’s standard classification and sequence-separation tokens as defined by Devlin et al. (2019), and [REL] is one of the relation tokens. For link prediction, we construct a dataset of variable-length paths through the KG by iteratively selecting a list of triples (h1, r1, t1), . . . , (hn, rn, tn) where hi+1 = ti to form a path p = (h1, r1, h2, . . . , rn, tn).
Entity Prediction The entity classification task can be trivially integrated into the masked-language
objective of BERT, by masking the tokens associated with the concept t.
Link Prediction We formulate link prediction as a narrow masked-language task by masking the relation tokens in the path dataset with another hidden relation token, for which the model is trained to fill in one of six relation types - as the triple classification and entity prediction tasks already have the partial goal of improving the model’s capability to associate synonymous terms with each other, we exclude synonym relations from the path dataset.
Triple Classification Following the work of Hao et al. (2020), the triple classification objective is formulated as a binary classification problem where the model is tasked with classifying triples as true or false. In order to generate training examples of false triples, we use two different negative sampling strategies. Firstly, to provide directly contrastive examples for existing relations, we sample triples (h, r, t) where h and t belong to different semantic groups and construct corresponding false triples with the same relation type and semantic group categories, i.e. (ĥ, r, t̂) /∈ G where ĥ and t̂ are of the same semantic group as h and t respectively. Secondly, to provide contrastive examples for relation types, we sample triples for which h and t are of the same semantic group, and form the negative training example by changing the relation type r. To ensure balance, the triple classification datasets used in this work are made up of 50% positive examples (real triples from the KG), 25% examples generated by the first negative sampling method and the rest by the second.
We perform stratified sampling on the base knowledge graph according to semantic groups, i.e. we ensure that the proportional representation of each semantic group in the knowledge-base triples for each language is maintained in the training datasets.
Mixed Objective Function In order to train BERT models using the UMLS-based reason-
ing tasks described above alongside the maskedlanguage objective, each training example is augmented with an indicator label that tells the model which loss function to apply to the sequence in question. The overall loss function is then calculated as
L = LMLM + α1LEP + α2LLP + α3LTC (2)
where the αi are scalar task-weighting coefficients and LMLM, LEP, LLP, and LTC correspond to the loss values for masked language modelling, entity prediction, link prediction and triple classification respectively. We use the standard cross-entropy classification loss for all tasks.

heading: 4 Experiments
text: For the evaluation of the approach described in the previous section, we restrict our attention in this paper to NER tasks. Where possible, we use the datasets and training-evaluation-test splits that are publicly available via the Huggingface datasets library1.

heading: 4.1 KG-integrated pre-training
text: Pre-training corpora As a resource for maskedlanguage pre-training, we utilise the European Clinical Case Corpus (E3C) version 2.0.02, a freelyavailable multilingual corpus of clinical narratives. We evaluate our method in three different languages; English, French and Spanish. These languages were chosen as they are the three most well-represented languages in the metathesaurus for which we have access to pre-trained clinical BERT models for comparison. The sizes of the combined UMLS-E3C datasets used are shown in Table 1.
For each language, we compare the performance of 1) a transformer model trained from scratch on each monolingual dataset (KGI-BERTEN,FR,ES) against 2) a multilingual version of the same model trained on all three datasets (KGI-BERTm), 3) a pre-trained monolingual biomedical model and 4) the same pre-trained model with supplementary training on the corresponding monolingual UMLSE3C dataset.
The UMLS-KGI models were trained for 64 epochs on each dataset, using the PyTorch implementation of the weighted ADAM optimizer
1https://huggingface.co/datasets 2https://live.european-language-grid.eu/
catalogue/corpus/7618
(Loshchilov and Hutter, 2019) with default parameters. We use a maximal sequence length of 256 for the masked-language modelling sequences, an effective batch size of 1500 and a triangular learning rate schedule peaking at 7.5 × 10−4. To take into account the varying sizes of the components of the pre-training dataset we set the values of the coefficients of the loss function such that they are inversely proportional to the number of documents available:
αi =
∑3 j=0,j ̸=i nj
2 ∑3
k=0 nk
where the nk correspond to the number of documents in the training set for each UMLS-based task. In this way, the E3C masked-language loss has the same weighting as the UMLS-based task losses.
Pre-trained models For supplementary training, we make use of what are, to the best of our knowledge, the overall best-performing biomedical BERT models of their size (pre-trained using masked-language tasks only) for each language, according to baseline experiments on the NER tasks.
For French, we use DrBERT (Labrak et al., 2023), for Spanish the RoBERTa-based biomedical model released by Carrino et al. (2021), which we refer to as BioRoBERTa-ES, and for English PubMedBERT (Gu et al., 2021). For training from scratch, we use the DistilBERT model configuration (Sanh et al., 2019) with 12 encoder layers and 12 attention heads.

heading: 4.2 Evaluation corpora
text: We evaluate these models on nine different clinical entity recognition tasks; four in French, two in Spanish and three in English. In order to ensure a fair comparison between models and evaluate more directly the knowledge transfer capabilities of the pre-trained models, we restrict ourselves to a oneshot setting for all tasks, i.e. the model is given a single pass over the training data before being evaluated on the test set. For all fine-tuning runs, we use an effective batch size of 4 (we found that very frequent optimizer updates give better results in for few-shot learning), learning rate 2 × 10−5 and weight decay of 0.01.
CAS/ESSAIS CAS (Grabar et al., 2018) and ESSAIS (Dalloux et al., 2021) are corpora of clinical cases in French for which a subset is annotated with part-of-speech tags as well as semantic biomedical annotations (UMLS concepts, negation, and
uncertainty). We evaluate our models on the two corresponding medical POS-tagging tasks, CASPOS and ESSAI-POS, as well as formulating a semantic-group token classification task using the CAS corpus annotations (CAS-SG).
QUAERO The QUAERO French Medical Corpus (Névéol et al., 2014) is a corpus of biomedical documents from EMEA and Medline annotated with UMLS concepts to facilitate entity recognition and document classification tasks. The NER evaluation task we make use of here, QUAEROMEDLINE, involves semantic group identification in the Medline documents.
PharmaCoNER (Gonzalez-Agirre et al., 2019) Designed for the automated recognition of pharmacological substances, compounds and proteins in Spanish-language clinical documents, this is a manually annotated subset of the Spanish Clinical Case Corpus (SPACCC (Intxaurrondo, 2018)).
MEDDOCAN Similarly to PharmaCoNER, the MEDDOCAN corpus (Marimon et al., 2019) is an annotated subset of SPACCC, in this case with semantic entity types relevant to clinical document anonymisation, i.e. words and expressions constituting Personal Health Information (PHI).
NCBI-Disease (Doğan et al., 2014) The NCBI disease corpus is made up of PubMed abstracts with annotated disease mentions. In this work, we restrict our attention to token classification at the mention level.
BioRED (Luo et al., 2022) This corpus is designed for biomedical relation extraction and en-
tity recognition; we focus on the latter in this work. This task can be considered a more semantically general version of the NCBI disease recognition task, in that the BioRED corpus consists of PubMed abstracts annotated with a diverse range of entity types including genes, proteins and chemicals.
JNLPBA04 NER Dataset (Collier and Kim, 2004) Developed in the context of a biomedical entity recognition shared task, this corpus consists of Medline documents annotated with mentions of DNA, RNA, proteins, cell types and cell lines.
We report the macro-averaged precision, recall and F1-score for each task. Results for the French, English and Spanish tasks can be seen in Tables 2, 3, and 4 respectively. We find that the bestperforming models are in general the pre-trained checkpoints for which training has been extended via knowledge graph integration. This is unsurprising given that these are the models that have undergone the most domain-specific pre-training among all variants. It is important to highlight, moreover, the fact that the KGI-BERT variants are competitive with the pre-trained baselines for many tasks, despite being trained on less data. The largest improvements brought about by the UMLS-KGI training strategy can be seen in the French and Spanish tasks, suggesting that this technique will be more beneficial for lower-resource languages for which there is more room for improvement with respect to existing models.
The number of documents and target label classes for each evaluation task is show in Table 5.

heading: 4.3 Ablation Experiments
text: In order to measure the relative effect of the three KG-derived pre-training tasks on downstream performance, we perform ablation experiments with the continually pre-trained models. This involved comparing the downstream performance on the NER tasks of different versions of the UMLSextended models, each with one of the three KGbased pre-training tasks excluded from the pretraining process. For ablation, we use identical experimental settings to those described previously, except with 32 pre-training epochs rather than 64.
In general, the ablation results, for which the macro F1 scores are shown in Table 6, suggest that the majority of the benefits in terms of NER performance are brought about by the link prediction task, although there are not enough statistically significant differences among the results to fully justify this conclusion.
It is clear also that certain tasks tend to add unhelpful noise to the model with respect to some tasks, in particular the ESSAI-POS task in French and the MEDDOCAN task in Spanish. This may be due to the nature of these entity recognition tasks being more linked to general semantic patterns (i.e. parts-of-speech and identifying information) such that the addition of biomedical knowledge to the models does not improve their representation of the
relevant concepts.

heading: 5 Conclusions and Future Work
text: This paper introduces UMLS-KGI, a framework for training BERT models using knowledge graphs requiring highly minimal adjustments to the standard language modelling paradigm. We show the potential of this method to increase the performance of BERT models on various NER tasks. The results presented in this paper suggest that for clinical NER tasks, high-quality small-scale datasets derived from structured information, alongside alongside relatively small clinical text corpora, can be as effective as large-scale corpora for pre-training BERT models. We make our models and dataprocessing pipelines freely available online.
Future work in this direction will involve the incorporation of more diverse graph-based reasoning tasks in the pre-training strategy with more finegrained representation of relation types, as well as intrinsic evaluation of the UMLS-KGI-BERT language representations via embedding visualisation and interpretability studies.
Limitations
The work presented in this paper is subject to a number of limitations which will be addressed in future work. Firstly, we evaluate UMLS-KGI-BERT on a very narrow range of tasks limited to token classification - a broader range of information extraction and reasoning tasks would be necessary for a more complete picture of the utility of our pretraining methods. In addition, we only train models for mid-to-high-resource languages; to properly validate the applicability of this approach, in particular the lessening of the need to rely on large training corpora, it will be necessary to train and evaluate such models in more low-resource settings.
Table 6: Macro-F1 scores for the ablation experiments.
Dataset Base Model KG Tasks CAS-POS CAS-SG QUAERO-MEDLINE ESSAI-POS DrBERT-4GB - 90.84 62.20 66.66 94.69
EP+LP 91.59 64.85 66.08 94.62 EP+TC 90.86 62.11 66.75 94.88 TC+LP 92.01 65.98 66.89 94.41 all 92.04 66.22 67.15 94.50
NCBI-Disease BioRED-NER JNLPBA04 PubMedBERT - 93.53 83.35 81.13
EP+LP 93.24 82.40 81.25 EP+TC 93.37 83.09 82.66 TC+LP 94.13 83.38 84.30 all 94.11 83.45 84.36
PharmaCoNER MEDDOCAN BioRoberta-ES - 81.11 91.84
EP+LP 81.12 91.86 EP+TC 82.40 91.80 TC+LP 83.22 91.71 all 83.46 91.77

heading: A Dataset Statistics
text: A.1 UMLS Knowledge Graph We use the 2022AB release of the UMLS knowledge graph, which contains 8,751,471 concepts defined by 3,711,072 unique identifiers (CUIs), and 25,369,590 relations. Restricting our attention to semantic types related to human biology and medicine, we end up with the base dataset outlined in Table 7.

heading: B Supplementary Experimental Details
text: Pre-trained Checkpoints We use the following pre-trained model weights downloaded from the HuggingFace model hub as baseline models;
• DrBERT: Dr-BERT/DrBERT-4GB
• PubMedBERT: microsoft/BiomedNLP-PubMedBERT-baseuncased-abstract-fulltext
• BioRoBERTa-ES: PlanTL-GOB-ES/roberta-basebiomedical-clinical-es
Model Hyperparameters The hyperparameter settings used for the pre-training on the UMLSbased dataset are shown in Table 8. The pretraining process used a linear learning rate schedule with warmup, where the learning rate increases from zero over the warmup period until it reaches the specified before decaying linearly over the rest of the training steps. In the interest of minimising the energy consumption of our experiments, we carried out very minimal hyperparameter search, leaving most parameters at their default values. The experiments were run using Python 3.8.15, with PyTorch version 2.0.0 and CUDA 11.8, along with the transformers library version 4.27.4.
Hardware specifications The pre-training experiments were run on four Nvidia Tesla V100 GPUs with 32GB of RAM, while the fine-tuning experiments were run on an RTX 2080 Ti with 11GB of RAM.

heading: Future Work
text: Future work in this direction will involve the incorporation of more diverse graph-based reasoning tasks in the pre-training strategy with more finegrained representation of relation types, as well as intrinsic evaluation of the UMLS-KGI-BERT language representations via embedding visualisation and interpretability studies. Limitations
The work presented in this paper is subject to a number of limitations which will be addressed in future work. Firstly, we evaluate UMLS-KGI-BERT on a very narrow range of tasks limited to token classification - a broader range of information extraction and reasoning tasks would be necessary for a more complete picture of the utility of our pretraining methods. In addition, we only train models for mid-to-high-resource languages; to properly validate the applicability of this approach, in particular the lessening of the need to rely on large training corpora, it will be necessary to train and evaluate such models in more low-resource settings. Table 6: Macro-F1 scores for the ablation experiments. Dataset Base Model KG Tasks CAS-POS CAS-SG QUAERO-MEDLINE ESSAI-POS DrBERT-4GB - 90. 84 62. 20 66. 66 94. 69
EP+LP 91. 59 64. 85 66. 08 94. 62 EP+TC 90. 86 62. 11 66. 75 94. 88 TC+LP 92. 01 65. 98 66. 89 94. 41 all 92. 04 66. 22 67. 15 94. 50
NCBI-Disease BioRED-NER JNLPBA04 PubMedBERT - 93. 53 83. 35 81. 13
EP+LP 93. 24 82. 40 81. 25 EP+TC 93. 37 83. 09 82. 66 TC+LP 94. 13 83. 38 84. 30 all 94. 11 83. 45 84. 36
PharmaCoNER MEDDOCAN BioRoberta-ES - 81. 11 91. 84
EP+LP 81. 12 91. 86 EP+TC 82. 40 91. 80 TC+LP 83. 22 91. 71 all 83. 46 91. 77