Diversifying Knowledge Enhancement of Biomedical Language Models using Adapter Modules and Knowledge Graphs

PLMs are trained on enormous corpora of training data, ranging from 3.3 billion tokens in the case of the original BERT Devlin et al., (2019), all the way to 3.5 trillion tokens in the case of the recent Falcon-180B model Almazrouei et al., (2023). The power of the model architecture, combined with transfer learning, has led to these models showing impressive capabilities on most NLP tasks. While the textual data used for the model training is usually completely unstructured in nature, research has shown that models like BERT do encode, to some extent, syntactic structures, hierarchical concepts, and certain semantic conceptual relations Rogers et al., (2021). Still, other studies have shown weakness in modeling tasks dealing with structured knowledge, such as hyponymy relations Ravichander et al., (2020) or preserving the association between text and meaning Di Sciullo, (2018).

In most cases, the knowledge we find and gather, especially scientific knowledge, can be represented in a structured manner. This is the underlying idea of knowledge graphs (KGs), a data structure that models concepts (entities) and relations between them in a graph-like format Ji et al., (2021). KGs have been used in the field of NLP to enhance the performance of NLP models in many downstream NLP tasks Schneider et al., (2022). There are multiple ways to combine KGs with PLMs. The knowledge triples from KGs can be embedded as vector representations such as TransE Wang et al., (2014) or TuckER Balazevic et al., (2019) and then combined with the vectors encoding text. Alternatively, the triples from KGs can be converted to sentences, and, in turn, these textual representations can then be used to fine-tune PLMs in the same way as with any other text. This approach was followed by COMET Bosselut et al., (2019), which utilized the knowledge graph ConceptNet Speer et al., (2017) to enhance the performance on commonsense reasoning tasks. Besides knowledge graphs, lexicons are sometimes used for knowledge enhancement Hoang et al., (2022).

While there are numerous ways to inject structured knowledge into PLMs such as adding it to the input and output of models Wei et al., (2021), an especially promising approach is adding adapters to the architecture of the model Colon-Hernandez et al., (2021). Adapters are small layers that are inserted within a language model and are subsequently fine-tuned to a specific task. The major benefit of adapters is that they add a minimal amount of additional parameters, thus significantly reducing the needed training time. Combined with freezing original model weights, adapters can avoid catastrophic forgetting, where the PLM’s performance deteriorates when all of its weights are fine-tuned with a new knowledge source. Adapters have been used for numerous purposes such as learning hierarchical representation Chronopoulou et al., (2022), transferring models trained on English to low-resource languages Wang et al., (2021), and in the domain of efficient transformers as low-rank adapters (LoRA) Hu et al., (2022). General knowledge-enhanced PLMs utilizing adapters include, for example, KnowBERT Peters et al., (2019) and K-Adapter Wang et al., (2020). A practical tool emerged that combines well-known adapter architectures in one place, called AdapterHub Pfeiffer et al., 2020b .

A major focus of knowledge enhancement in PLMs is in domain adaption to expert domains such as the biomedical domain. So far, most of the advancements have focused on utilizing the knowledge graph UMLS Bodenreider, (2004) for this purpose. Examples include BERT-MK He et al., (2020) and KeBioLM Yuan et al., (2021), which both fine-tune the whole weights of the base language model by using masked language modeling of triples from UMLS. Biomedical PLMs can then be used for various NLP tasks, such as biomedical text summarization Abacha et al., (2021), named entity recognition Sung et al., (2022), medical fact-checking Vladika and Matthes, (2023), information retrieval Luo et al., (2022), or health question answering Vladika et al., (2023).

There are also existing approaches using adapters for biomedical knowledge enhancement. Representative works are DAKI Lu et al., (2021), which fine-tunes the adapters with entity prediction task, and KEBLM Lai et al., (2023), which fine-tunes the adapters on three different knowledge types from UMLS and PubChem Kim et al., (2019), namely entity descriptions, entity-entity relations, and entity synonyms. The most similar approach to ours and a direct inspiration was the Mixture-of-Partitions (MoP) approach Meng et al., (2021), where the adapters were fine-tuned on smaller subgraphs of UMLS.

In essence, our work builds on the present foundations of adapter-based biomedical models and uses the yet unexplored knowledge graph OntoChem, which is rich with chemical knowledge. For our experiments, we use the well-known biomedical PLM PubMedBERT as well as the yet unexplored but powerful BioLinkBERT base model. Following the suggestions of Meng et al., (2021), we use only the triplets corresponding to the 20 most frequent relations of OntoChem for the knowledge injection. An example of an OntoChem triplet can be seen in Figure 1. Finally, we provide a deeper qualitative analysis of learned structured knowledge on a specific dataset. Notably, our work achieves the SOTA (averaged) performance on the question-answering BioASQ-7b dataset.

A central element of our method is the knowledge graph (KG). This KG is a structured representation of information denoted as a collection of ordered triples Ji et al., (2021). We denote these triples as $(s,r,o)$, where s is a subject, r is a relation, and o is an object. Both s and o are entities that come from an entity set E, while relations come from a relation set R. Each entity and relation in the KG is associated with its corresponding textual surface form. This form can take the shape of a single word or a compound term (e.g., for names of chemicals) or even a concise phrase, especially in case of relations. This textual association is critical as it bridges the gap between the structured KG and natural language, allowing for easier injection of KG knowledge into the language models and associated fine-tuning.

The primary objective is to enhance the capabilities of a pre-trained language model, denoted as LM, by integrating the knowledge contained within the KG. To achieve this, we need a training objective that effectively incorporates the KG knowledge into the model. Most encoder-only PLMs based on the original BERT use masked language modeling (MLM) as one of its pre-training objectives. This task consists of masking a certain word in a given sentence and having the model predict which word would fit the best in the place of the masked token. We follow the established approach of using an entity prediction objective, where we mask one of the entities and have the model predict which token would best fit. In this way, the model incorporates the structured knowledge of $(s,r,o)$ triples into its internal weights.

During the prediction of the masked token, the model produces a probability distribution (with a softmax function) over all of the entities from KG’s entity set E. Considering the massive size of the biomedical KGs we use in the paper, computing the softmax function over all its entities carries a lot of computation complexity. This issue can significantly slow down model training and inference. To bridge this challenge, some approaches have been suggested in the literature. We opt for the approach of Meng et al., (2021), which involves partitioning the KG into smaller subgraphs, which are then trained on independently, and later, their knowledge combined to unified knowledge representations.

The process of dividing a KG yields smaller subgraphs that we denote as $G_{1},G_{2},...,G_{k}$. We set $k$ in final experiments to be $20$, following empirical observations and previous literature, which balances efficiency and graph coverage well. Ideally, these $20$ smaller subgraphs should be almost equal-sized components, meaning nodes are balanced across partitions. Additionally, the capacity of edges between different components should be minimized to maximize the retention of factual knowledge. This is a problem known as balanced graph partitioning and is known to be an NP-complete problem Andreev and Räcke, (2004). Several good approximations have been developed to determine the exact solution despite its computational complexity. We opt for the METIS algorithm Karypis and Kumar, (1997), which was used in other works dealing with large-scale KG partitioning Zheng et al., (2020).

Once the KG is appropriately partitioned, the process of fine-tuning the LM can be started. We deploy adapter modules for this purpose. As mentioned previously, adapters are newly initialized feed-forward networks inserted between the transformer model’s layers. Notably, the training of adapter modules does not require fine-tuning the existing parameters of the pre-trained model. Instead, it focuses solely on updating the parameters within the adapters. This strategy ensures that the pre-trained model’s core knowledge remains intact while enabling the model to specialize in the biomedical domain by adapting to the specific knowledge contained in the KG.

There are multiple adapter module configurations, such as Houlsby et al., (2019) and Bapna and Firat, (2019). The adapter module configuration used in the paper is based on the one by Pfeiffer et al., 2020a , the so-called Pfeiffer architecture. In this configuration, only one adapter module is added as a down-projection and up-projection, unlike the Houlsby architecture, where there are two projections. While the Houlsby architecture has more learning capacity, it comes with training and inference speed costs. Previous studies showed no significant difference in performance between the model architectures, making Pfeiffer architecture a very lightweight choice that brings powerful learning capabilities.

As already mentioned, masked language modeling is used to fine-tune the adapter modules. More precisely, it is a task of entity prediction since a missing entity from the graph triple is being predicted. Given a subgraph $G_{k}$ and its triples $(s,r,o)$, each triple has a textual representation. The object entity o is removed from each triple, and the remaining two elements of the triple are transformed into a textual representation like: ”[CLS] s [SEP] r [SEP]”. The adapter module is then trained to predict the missing object entity using the representation of the [CLS] token. The parameters of the adapter module are optimized by minimizing the cross-entropy loss.

Finally, with a set of knowledge-encapsulated adapter modules at hand, we need to fuse their knowledge together into a final representation. For this, we use the so-called AdapterFusion mixture layers Pfeiffer et al., 2020a . These layers serve the purpose of combining knowledge from various adapters to enhance the model’s performance on downstream tasks. It is a relatively recent approach designed to effectively learn how to combine information from a set of task-specific adapters. It does so by employing a softmax attention mechanism that assigns contextual mixture weights over the adapters. These weights are then used to predict task labels in the final layer. The composition of these layers and their interactions ultimately contribute to the model’s ability to generalize and perform well on a range of tasks.

All of our datasets, with the exception of MedNLI, originated from a collection of common biomedical NLP tasks known as BLURB – Biomedical Language Understanding and Reasoning Benchmark.³³3https://microsoft.github.io/BLURB/index.html Inspired by a similar suite of tasks for general-purpose natural language understanding (NLU) known as GLUE Wang et al., (2018), BLURB covers a wide-range of tasks related to biomedical NLU. This means no tasks include text generation and are all essentially classification tasks, which makes them convenient to evaluate with common classification metrics such as precision, recall, accuracy, and F1 score. The four datasets are described in continuation.