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

Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 4672–4681 November 7–11, 2021. c⃝2021 Association for Computational Linguistics 4672 Mixture-of-Partitions: Infusing Large Biomedical Knowledge Graphs into BERT Zaiqiao Meng♠, Fangyu Liu♠, Thomas Hikaru Clark♠ Ehsan Shareghi♣♠, Nigel Collier♠ ♠Language Technology Lab, University of Cambridge ♣Department of Data Science and AI, Monash University ♠{zm324, fl399, thc44, nhc30}@cam.ac.uk ♣ehsan.shareghi@monash.edu Abstract Infusing factual knowledge into pretrained models is fundamental for many knowledge- intensive tasks. In this paper, we propose Mixture-of-Partitions (MoP), an infusion ap- proach 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, Clas- sification), and the results show that our MoP consistently enhances the underlying BERTs in task performance, and achieves new SOTA performances on five evaluated datasets. 1 1 Introduction Leveraging factual knowledge to augment pre- trained language models is of paramount impor- tance for knowledge-intensive tasks, such as ques- tion answering and fact checking (Petroni et al., 2021). Especially in the biomedical domain where public training corpora are limited and noisy, trusted biomedical KGs are crucial for deriving ac- curate inferences (Li et al., 2020; Liu et al., 2021). However, the infusion of knowledge from real- world biomedical KGs, where entity sets are very large (e.g. UMLS, Bodenreider 2004, 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 func- tion to minimize the risk of catastrophic forget- ting (Xiong et al., 2019; Zhang et al., 2019; Wang 1Our code, models and other related resources can be found in https://github.com/cambridgeltl/mop. et al., 2021, 2020; Peters et al., 2019; Yuan et al., 2021). Alternatively, entity masking (or entity prediction) has emerged as one of the most popu- lar self-supervised training objectives for infusing entity-level knowledge into pretrained models (Sun et al., 2019; Zhang et al., 2019; Yu et al., 2020; He et al., 2020). However, due to the large number of entities in biomedical KGs, computing an exact softmax over all entities is very expensive for train- ing and predicting (De Cao et al., 2021). Although negative sampling techniques could alleviate the computational issue (Sun et al., 2020), tuning an appropriately hard set of negative instances can be challenging and predicting a very large number of labels may generalize poorly (Hinton et al., 2015). 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, Lee et al. 2020; SciB- ERT, Beltagy et al. 2019; and PubMedBERT, Gu et al. …

Content truncated to 500 words