Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT (2021)

Abstract: Infusing factual knowledge into pre-trained models is fundamental for many knowledge-intensive tasks. In this paper, we proposed Mixture-of-Partitions (MoP), an infusion approach that can handle a very large knowledge graph (KG) by partitioning it into smaller sub-graphs and infusing their specific knowledge into various BERT models using lightweight adapters. To leverage the overall factual knowledge for a target task, these sub-graph adapters are further fine-tuned along with the underlying BERT through a mixture layer. We evaluate our MoP with three biomedical BERTs (SciBERT, BioBERT, PubmedBERT) on six downstream tasks (inc. NLI, QA, Classification), and the results show that our MoP consistently enhances the underlying BERTs in task performance, and achieves new SOTA performances on five evaluated datasets.

## introduction leveraging factual knowledge to augment pretrained language models is of paramount importance for knowledge-intensive tasks, such as question answering and fact checking @xcite . especially in the biomedical domain where public training corpora are limited and noisy, trusted biomedical kgs are crucial for deriving accurate inferences @xcite . however, the infusion of knowledge from realworld biomedical kgs, where entity sets are very large (e.g. umls, @xcite , contains ∼4m entities) demands highly scalable solutions. although many general knowledge-enhanced language models have been proposed, most of them rely on a computationally expensive joint training of an underlying masked language model (mlm) along with a knowledge-infusion objective function to minimize the risk of catastrophic forgetting @xcite @xcite @xcite . alternatively, entity masking (or entity prediction) has emerged as one of the most popular self-supervised training objectives for infusing entity-level knowledge into pretrained models @xcite @xcite . however, due to the large number of entities in biomedical kgs, computing an exact softmax over all entities is very expensive for training and predicting @xcite . although negative sampling techniques could alleviate the computational issue @xcite , tuning an appropriately hard set of negative instances can be challenging and predicting a very large number of labels may generalize poorly @xcite . to address the aforementioned challenges, we propose a novel knowledge infusion approach, named mixture-of-partitions (mop), to infuse factual knowledge based on partitioned kgs into pretrained models (biobert, @xcite @xcite . more concretely, we first partition a kg into several sub-graphs each containing a disjoint subset of its entities by using the metis algorithm @xcite , and then the transformer adapter module @xcite is applied to learn portable knowledge parameters from each subgraph. in particular, using adapter module to infuse knowledge does not require fine-tuning the parameters of the underlying berts, which is more flexible and efficient while avoiding the catastrophic forgetting issue. to utilise the independently learned knowledge from sub-graph adapters, we introduce mixture layers to automatically route useful knowledge from these adapters to downstream tasks. figure 1 illustrates our approach. our results and analyses indicate that our "divide and conquer" partitioning strategy effectively preserves the rich information presented in two biomedical kgs from umls while enabling us to sub-graph (a) graph partitioning add & norm ff up ff down add & norm multi-head attention feed forward add & norm + group 1 adapter group 2 adapter group k adapter (b) infuse knowledge with adapter (d) mixture of adapters add & norm add & norm multi-head attention feed forward add & norm group 1 adapter mixture layer (c) collect pretrained adapters we denote a kg as a collection of ordered triples @xmath0, r, t)|h, @xmath1, @xmath2, where e and r are the sets of entities and relations, respectively. all the entities and relations are associated with their textual surface forms, which can be a single word (e.g. fever), a compound (e.g. sars-cov-2), or a short phrase (e.g. has finding site). given a pretrained model θ 0 , our task is to learn φ g based on an input knowledge graph g, such that it encapsulates the knowledge from g. the training objective, l g , can be implemented in many ways such as relation classification @xcite , entity linking @xcite , next sentence prediction (goodwin and demner-fushman, 2020), or entity prediction @xcite . in this paper, we focus on entity prediction, one of the most widely used objectives, and leave exploration of other objectives for future work. as mentioned earlier, exact softmax over all entities is extremely expensive @xcite for large-scale kgs, hence we resort to the principle of "divide and conquer", and propose a novel approach called mixture-of-partition (mop). specifically, our mop first partitions a large kg into smaller sub-graphs (i.e., @xmath3, g 2 , . . . , g k }, §2.1), and learns subgraph specific parameters on each sub-graph separately (i.e., {φ g 1 , φ g 2 , . . . , φ g k }, §2.2). then these sub-graph parameters are fine-tuned through mixture layers to route the sub-graph specific knowledge into a target task ( §2.3). ## conclusion and future work in this paper, we proposed mop, a novel approach for infusing knowledge by partitioning knowledge graphs into smaller sub-graphs. we show that while the knowledge-encapsulated adapters perform very differently over different sub-graphs, our proposed mop can automatically leverage and balance the useful knowledge across those adapters to enhance various downstream tasks. in the future, we will evaluate our approach using some general domain kgs based on some general domain tasks.