Specializing Static and Contextual Embeddings in the Medical Domain Using Knowledge Graphs: Let’s Keep It Simple (2022)

Abstract: Domain adaptation of word embeddings has mainly been explored in the context of retraining general models on large specialized corpora. While this usually yields good results, we argue that knowledge graphs, which are used less frequently, could also be utilized to enhance existing representations with specialized knowledge. In this work, we aim to shed some light on whether such knowledge injection could be achieved using a basic set of tools: graph-level embeddings and concatenation. To that end, we adopt an incremental approach where we first demonstrate that static embeddings can indeed be improved through concatenation with in-domain node2vec representations. Then, we validate this approach on contextual models and generalize it further by proposing a variant of BERT that incorporates knowledge embeddings within its hidden states through the same process of concatenation. We show that this variant outperforms plain retraining on several specialized tasks, then discuss how this simple approach could be improved further. Both our code and pre-trained models are open-sourced for future research. In this work, we conduct experiments that target the medical domain and the English language.

## introduction the popularization of transfer learning, particularly in the context of pre-training language models to serve as encoders in downstream tasks, has led to an ever-expanding set of methods for representing textual data: e.g. elmo @xcite , bert @xcite , xlnet @xcite . while these models continuously push forward the expected level of performance on so-called "general domain" tasks (e.g. glue foot_0 ), they usually lag behind when it comes to more specialized areas like the medical domain (see blue 2 and blurb foot_2 benchmarks). as a result, there is a growing interest in finding ways in which these out-of-the-box representations can be specialized, with most efforts focusing on retraining general models on specialized corpora: e.g. clinicalbert @xcite , biobert @xcite , and biomed-roberta @xcite . however, pre-trained language models have also been shown to benefit from external knowledge injection, with approaches like libert @xcite , knowbert @xcite , and kepler @xcite in the general domain, or @xcite and @xcite in the medical domain. yet, these efforts usually involve complex modifications to the architecture of underlying models and/or their pre-training procedure, which may convey the impression that knowledge injection cannot be achieved in simpler ways. in this work, we propose a simple approach to embedding specialization that relies on knowledge graph embeddings and concatenation. we argue that by building a simple but strong baseline first, we lay the foundation for future improvements that may be easily achieved by upgrading to more sophisticated knowledge embeddings or combination methods. in practice, we show that medical concept embeddings obtained from an in-domain knowledge graph can be combined through concatenation with word representations to effectively construct specialized "meta-embeddings" @xcite . moreover, in the specific case of contextual embeddings, we show that these concept embeddings can be combined either externally, with a general-domain model, or internally, during the pretraining of a specialized model, to achieve varying levels of model specialization. all our models are trained and evaluated in pairs, and in exactly the same conditions, to highlight to the greatest extent the impact of our strategies. our contributions are the following: • we build two sets of knowledge representations by applying node2vec @xcite to concepts from mesh (biomedical) and snomed ct (clinical). ## related work our approach is related to the similar but usually distinct topics of knowledge injection and domain adaptation. in fact, most attempts at domain adaptation do not aim to inject external knowledge directly into models but rather indirectly, through retraining on specialized corpora, as this is known to bring significant improvements when such indomain corpora are available @xcite . on the other hand, research concerned with knowledge injection usually tackles the problem within the same domain. for instance, sembert @xcite , comet @xcite , ernie @xcite , k-bert @xcite , and kepler all inject general knowledge into general-domain models. similar efforts in the medical domain @xcite @xcite directly inject medical knowledge during medical pre-training. in this work, we first set out to determine whether the performance of general-domain models, both static and contextual, can be improved solely using specialized knowledge embeddings, then only do we incorporate this approach into the usual model adaptation via pre-training. methods that utilize knowledge graphs, for instance @xcite @xcite , can be broadly grouped into two categories: those that use the structured data directly and those that encode this data into numerical representations. instances of direct utilization include kg-bert @xcite where triples (concept_1, relation, concept_2) are used to inject bert with medical information through auxiliary tasks like knowledge graph completion and triple classification. entity linking in @xcite or more specialized tasks in @xcite are also used as auxiliary tasks for performing such injection. while these methods can be effective, we argue that an indirect approach is desirable as it presents the specialized knowledge in the same format as the word embeddings, thus reformulating knowledge injection as a meta-embedding problem. meta-embeddings combine two or more underlying sets of embeddings into a single final representation. there are many approaches to meta-embeddings like dynamic meta embeddings (dme, @xcite ) where each embedding is projected down to the same dimension before being used in a linear combination, or word prisms @xcite , which further improve upon dmes by enforcing desirable orthogonality properties during training. in this work, we use a simple but strong baseline for meta-embeddingsconcatenation-which ensures that both word and knowledge information is accessible at all times. more sophisticated approaches, although likely to improve our overall performance, are left for future work. ## knowledge graph representations in order to use concatenation to specialize word embeddings with a knowledge base, we first need to be able to convert this knowledge base into dense numerical representations. there are several methods for embedding knowledge graphs (e.g. rotate @xcite , tucker @xcite ), and these usually produce multifaceted relationdependent concept representations. however, for simplicity, we only consider a single relation which enables us to use a graph-level method instead, namely node2vec @xcite . ## results and discussion for better visibility and given the large number of experiments, we present our results in pairs composed of a baseline and a specialized version of that baseline. we report the performances of each model pair as a set of two consecutive rows with the baseline on top (see bert vs. @xcite when looking at the results for contextual embeddings, we can see several instances where the concatenation with node2vec proves to be beneficial. however, there seems to be a discrepancy where sometimes this concatenation does improve the characterbert baseline on one hand but impairs the bert baseline on the other (see chemprot and ddi). a closer look at these cases shows that plain characterbert performs slightly lower than plain bert in these situations, which may mean that the knowledge representations compensate for any information that may be missing in the baseline characterbert model, relative to the task. ## conclusion and future work in this paper, we focused on exploring the extent to which specialized information from a knowledge graph could be injected into existing word embeddings using a very simple set of tools: graph embeddings and concatenation. while focusing on the medical domain in the english language, we conducted multiple evaluations on tasks ranging from entity recognition to sentence similarity. these evaluations demonstrated that concatenation with in-domain graph representations can be a simple yet effective approach to model specialization, with significant gains on several tasks. moreover, applying the same process of concatenation within transformer-based contextual models proved beneficial as well, with notable improvements using knowledge injection modules (kims) on several downstream tasks. as mentioned in section 3.1, many more types of relations beyond is_a could be used to improve the quality of the generated knowledge representations. an interesting path for future work may be to use recent meta-embedding methods like word prisms to learn multifaceted knowledge representations from multiple underlying representations corresponding to two or more types of relations.