Source-linked AI summary
UmlsBERT: Clinical Domain Knowledge Augmentation of Contextual Embeddings Using the Unified Medical Language System Metathesaurus
George Michalopoulos, Yuanxin Wang, Hussam Kaka, Helen Chen, Alexander Wong
TL;DR
Biomedical contextual models achieve strong results through domain-specific pre-training but generally omit structured expert knowledge. UmlsBERT addresses this gap by integrating UMLS concept and semantic-type information into BERT pre-training, and it outperforms other biomedical BERT models across clinical downstream tasks.
Problem
Biomedical contextual models do not incorporate structured expert domain knowledge from a knowledge base into embedding pre-training.
Method
UmlsBERT augments BERT pre-training with UMLS concept connections, a multi-label masked-language-modeling loss, and semantic type embeddings.
Results
UmlsBERT outperforms other biomedical BERT models across various clinical NER and natural language inference tasks.
Takeaways & Limitations
UmlsBERT can encode associations among clinically related terms and create more meaningful biomedical input embeddings.
Takeaways & Limitations
Pre-training requires substantial computational resources, including 10 days on two NVIDIA V100 16GB GPUs and 224 GB of system RAM.
Abstract
from arXiv · showhide
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 expert domain knowledge. In this work, we introduced UmlsBERT, a contextual embedding model that integrates domain knowledge during the pre-training process via a novel knowledge augmentation strategy. More specifically, the augmentation on UmlsBERT with the Unified Medical Language System (UMLS) Metathesaurus was performed in two ways: i) connecting words that have the same underlying `concept' in UMLS, and ii) leveraging semantic group knowledge in UMLS to create clinically meaningful input embeddings. By applying these two strategies, UmlsBERT can encode clinical domain knowledge into word embeddings and outperform existing domain-specific models on common named-entity recognition (NER) and clinical natural language inference clinical NLP tasks.
1 Introduction
UmlsBERT augments BERT-based contextual embeddings with UMLS clinical knowledge during pre-training. It connects words sharing concepts, incorporates semantic types, and outperforms existing BERT models across clinical NER and inference tasks.
- Existing biomedical contextual models use domain-specific corpora but do not incorporate structured expert knowledge from a knowledge base during embedding pre-training.This gap motivates augmenting contextual representations with UMLS information.
- The model integrates UMLS knowledge into BERT pre-training to produce semantically enriched contextual representations.The approach is designed to combine BERT’s contextual learning with clinical domain knowledge.
- A multi-label masked-language-modeling loss connects clinical words that share a UMLS concept unique identifier.The loss incorporates connections between clinical words through the UMLS CUI attribute.
- Semantic type embeddings associate words belonging to the same UMLS semantic type during input embedding.This mechanism encourages the model to consider clinically meaningful associations between words of the same semantic type.
- UmlsBERT outperforms BioBERT, Bio_ClinicalBERT, and general-domain BERT on multiple clinical NER and natural language inference tasks.The paper presents this as evidence that its knowledge-augmentation modifications improve downstream performance.
2 Related Work
Related work progresses from general contextual embeddings to biomedical models that use domain-specific corpora, while some approaches inject external lexical knowledge into BERT.
- ELMo and BERT established contextualized representations whose word embeddings depend on surrounding context.ELMo uses a bidirectional language model, while BERT uses bidirectional transformers.
- Knowledge-enhanced BERT approaches incorporate external lexical semantics into pre-training objectives or input embeddings.Sense-BERT uses WordNet supersenses, while GlossBERT targets word sense disambiguation.
- Biomedical models such as BioBERT and Bio_ClinicalBERT improve contextual representations by pre-training on biomedical or clinical corpora.BioBERT combines general and biomedical text, while Bio_ClinicalBERT further pre-trains BioBERT on MIMIC-III clinical text.
3 Data
UmlsBERT is pre-trained on clinical notes from MIMIC-III and evaluated on MedNLI and four i2b2 NER tasks using established dataset splits and UMLS terminology identification.
- UmlsBERT is pre-trained on 2,083,180 MIMIC-III clinical notes and test reports from the NO-TEEVENTS table.The notes come from anonymized English electronic medical records covering over forty-thousand intensive-care patients admitted between 2001 and 2012.
- The datasets use the same train, development, and test splits as Alsentzer et al. (2019), with dataset statistics summarized in Table 1.Table 1 reports sentence counts and the number of classes for each dataset.
- The evaluation covers English MedNLI natural language inference and four i2b2 named-entity recognition tasks.The selected benchmarks support comparison with prior biomedical BERT models and future studies.
- UMLS term identification uses the UMLS 2020AA version.
4 Methods
UmlsBERT augments BERT pre-training with UMLS concept relationships and semantic-type information to encode clinical knowledge in contextual embeddings. It adds semantic-type embeddings to inputs and replaces single-label masked-word prediction with CUI-based multi-label learning.
- 4.2.1 Semantic type embeddings: UmlsBERT adds semantic-type embeddings to BERT input vectors, enriching clinically identified words while assigning non-UMLS words a zero-filled vector.The semantic-type matrix represents 44 UMLS types, and the added vector is derived from a one-hot semantic-type indicator.
- 4.2.1 Semantic type embeddings: The semantic-type augmentation is intended to enrich representations of rare words whose meanings are difficult to learn from corpus context alone.The model hypothesizes that semantic-type information can supply meaningful representation information for rare training-corpus words.
- 4.2.2 Updating the loss function of Masked LM task: UmlsBERT updates masked language modeling to predict all words sharing the masked token’s UMLS CUI rather than only the masked word.For example, masking “lungs” targets “lung,” “lungs,” and “pulmonary,” which share CUI C0024109.
- 4.3 UmlsBERT Training: The model applies the augmented masked-language-modeling procedure to a Bio_ClinicalBERT initialization further pre-trained on MIMIC-III notes before downstream fine-tuning.The same fine-tuning method, vocabulary, and WordPiece tokenization are used across comparison models.
- 4.2 Updating BERT with UMLS knowledge: The study uses UMLS because it combines major biomedical vocabularies and standardized clinical terminologies, supplying concepts and semantic groups for augmentation.UMLS connects synonymous or similar terms through CUIs and groups concepts by semantic type.
- 4.3 UmlsBERT Training: UmlsBERT is trained for 1,000,000 steps with batch size 64, maximum sequence length 128, and learning rate 5 · 10^-5.Training used two NVIDIA V100 16GB GPUs and otherwise default hyperparameters.
5 Results
UmlsBERT is evaluated against BERT-based clinical models across downstream tasks, with ablations and qualitative analyses examining its semantic type and UMLS-based knowledge augmentation.
- 5.1 Downstream Clinical NLP Tasks: UmlsBERT achieves the best results in 4 of 5 tasks, including F1 scores of 93.6%, 88.6%, and 79.4% on i2b2 2006, 2010, and 2012, plus 83.0% MedNLI accuracy.The comparison reports mean and standard deviation across five runs for competing models.
- 5.1.2 BERT-based model comparison: On i2b2 2014, BERTbase reaches 95.2%, while UmlsBERT does not surpass it but performs better than the other biomedical BERT models.The authors attribute this boundary partly to differences between MIMIC-III’s PHI marker and the synthetic masks used in the de-identification task.
- 5.1.3 Effect of semantic type embeddings: Adding semantic type embeddings improves UmlsBERT performance on every evaluated dataset compared with the variation without them.The ablation compares UmlsBERT variants with and without semantic type embeddings.
- 5.2 Qualitative Embedding Comparisons: UmlsBERT uniquely discovers clinically meaningful associations such as kidney–ren, mass–lump, bleeding–hem, and feet–pedal through UMLS-linked representations.The associated word pairs share common UMLS CUIs, which are incorporated through the multi-label Masked LM training objective.
- 5.2 Qualitative Embedding Comparisons: The generic nearest-neighbor analysis finds no meaningful difference among the four models, indicating that medical-domain training did not eliminate general-domain associations.The qualitative comparison uses anatomy, disorder, and generic word categories.
- 5.3 Semantic Type Embedding Visualization: UMAP shows semantic-group clustering in UmlsBERT input embeddings but not in Bio_ClinicalBERT embeddings.The authors connect this clustering to adding semantic type vectors, which make words of the same semantic type more similar.
6 Conclusion and Future Work
UmlsBERT incorporates UMLS biomedical knowledge into contextual embeddings and outperforms other biomedical BERT models on downstream tasks. Future work will examine more complex downstream layers, UMLS hierarchical associations, and additional datasets and tasks.
- UmlsBERT incorporates biomedical knowledge during contextual-embedding pre-training and outperforms other biomedical BERT models on various downstream tasks.
- Future work will test more complex layers atop UmlsBERT embeddings, explore UMLS hierarchical associations, and evaluate additional datasets and biomedical tasks.
Ethical Considerations
Medical contextual embeddings may support processing and understanding health data, but adopting UmlsBERT requires substantial computational resources. Public release is intended to help groups use the model without repeating its costly pre-training.
- Medical contextual embeddings may help process and understand the rapidly growing volume of health and biomedical data.
- UmlsBERT pre-training required 10 days, two NVIDIA V100 16GB GPUs, and 224 GB of system RAM, resources many groups may not afford.
- Publicly releasing UmlsBERT allows clinical NLP groups to use its embeddings or fine-tune it for downstream tasks without repeating pre-training.
- Automatic encoding of diseases and procedures in electronic medical records could reduce data-extraction and reporting time and cost.