Scientific Language Models for Biomedical Knowledge Base Completion: An Empirical Study

1 . Introduction . Understanding complex diseases such as cancer , HIV , and COVID-19 requires rich biological , chemical , and medical knowledge . This knowledge plays a vital role in the process of discovering therapies for these diseases — for example , identifying targets for drugs [ Lindsay , 2003 ] requires knowing what genes or proteins are involved in a disease , and designing drugs requires predicting whether a drug molecule will interact with specific target proteins . In addition , to alleviate the great costs of designing new drugs , drug repositioning [ Luo et al. , 2021 ] involves identification of existing drugs that can be re-purposed for other diseases . Due to the challenging combinatorial nature of these tasks , there is need for automation with machine learning techniques . Given the many links between biomedical entities , recent work [ Bonner et al. , 2021a , b ] has highlighted the potential benefits of knowledge graph ( KG ) data representations , formulating the associated tasks as KG completion problems — predicting missing links between drugs and diseases , diseases and genes , and so forth . The focus of KG completion work — in the general domain , as well as in biomedical applications — is on using graph structure to make predictions , such as with KG embedding 1. https : //github.com/rahuln/lm-bio-kgc ( KGE ) models and graph neural networks [ Zitnik et al. , 2018 , Chang et al. , 2020 ] . In parallel , recent work in the general domain has explored the use of pretrained language models ( LMs ) as “ soft ” knowledge bases , holding factual knowledge latently encoded in their parameters [ Petroni et al. , 2019 , 2020 ] . An emerging direction for using this information for the task of KG completion involves fine-tuning LMs to predict relations between pairs of entities based on their textual descriptions [ Yao et al. , 2019 , Kim et al. , 2020 , Wang et al. , 2021 , Daza et al. , 2021 ] . In the scientific domain , this raises the prospect of using LMs trained on millions of research papers to tap into the scientific knowledge that may be embedded in their parameters . While this text-based approach has been evaluated on general domain benchmarks derived from WordNet [ Miller , 1995 ] and Freebase [ Bollacker et al. , 2008 ] , to our knowledge it has not been applied to the task of scientific KG completion . Our contributions . We perform an extensive study of LM-based KG completion in the biomedical domain , focusing on three datasets centered on drugs and diseases , two of which have not been used to date for the KG completion task . To enable exploration of LM-based models , we collect missing entity descriptions , obtaining them for over 35k entities across all datasets . We evaluate a range of KGE models and domain-specific scientific LMs pretrained on different biomedical corpora [ Beltagy et al. , 2019 , Lee et al. , 2020 , Alsentzer et al. , 2019 , Gu et al. , 2020 ] . We conduct analyses of predictions made by both types of models and find them to have complementary strengths , echoing similar observations made in recent work in the general domain [ Wang et al. , 2021 ] and motivating integration of both text and graph modalities . Unlike previous work , we train a router that selects for each input instance which type of model is likely to do better , finding it to often outperform average-based ensembles . Integration of text and graph modalities provides substantial relative improvements of 13– 36 % in mean reciprocal rank ( MRR ) , and routing across multiple LM-based models further boosts results . Finally , we demonstrate the utility of LM-based models when applied to entities unseen during training , an important scenario in the rapidly evolving scientific domain . Our hope is that this work will encourage further research into using scientific LMs for biomedical KG completion , tapping into knowledge embedded in these models and making relational inferences between complex scientific concepts . 2 . Task and Methods . We begin by presenting the KG completion task and the approaches we employ for predicting missing links in biomedical KGs , including our model integration and inductive KG completion methodologies . An overview of our approaches is illustrated in Figure 1 . 2.1 KG Completion Task . Formally , a KG consists of entities E , relations R , and triples T representing facts . Each triple ( h , r , t ) ∈ T consists of head and tail entities h , t ∈ E and a relation r ∈ R. An entity can be one of many types , with the type of an entity e denoted as T ( e ) . In our setting , each entity is also associated with some text , denoted as text ( e ) for e ∈ E . This text can be an entity name , description , or both ; we use the entity ’ s name concatenated with its description when available , or just the name otherwise . For example , the fact ( aspirin , treats , headache ) might be an ( h , r , t ) triple found in a biomedical KG that relates drugs and diseases , with the head and tail entities having types T ( h ) = drug and T ( t ) = disease . The task of KG completion or link prediction involves receiving a triple ( h , r , ? ) ( where ? can replace either the head or tail entity ) and scoring all candidate triples { ( h , r , t′ ) | t′ ∈ S } such that the correct entity that replaces ? has the highest score . For the example listed above , a well-performing model that receives the incomplete triple ( aspirin , treats , ? ) should rank the tail entity headache higher than an incorrect one such as diabetes . S can be the entire set of entities ( i.e. , S = E ) or some fixed subset . In the transductive setting , the set of facts T is split into a training set Ttrain and a test set Ttest such that all positive triples in the test set contain entities seen during training . In contrast , for inductive KG completion the triples in the test set may contain entities not seen during training ( see Section 2.4 ) . 2.2 Methods . Ranking-based KG completion . Each KG completion model in our experiments learns a function f that computes a ranking score s = f ( x ) for a given triple x = ( h , r , t ) . Models are trained to assign a high ranking score to correct positive triples from the set of known facts T and a low ranking score to triples that are likely to be incorrect . To do so , we use the max-margin loss function Lrank ( x ) = 1N ∑N i=1 max ( 0 , λ − f ( x ) + f ( x′i ) ) , where λ is a margin hyperparameter , x ∈ T is a known positive triple in the KG , and x′i is a negative triple constructed by randomly corrupting either the head or tail entity of x with an entity of the same type . KG embedding ( KGE ) models . For each entity e ∈ E and each relation r ∈ R , KG embedding ( KGE ) models learn a vector representation E ( e ) ∈ Rm and R ( r ) ∈ Rn . For a given triple ( h , r , t ) , each model computes the ranking score f ( h , r , t ) as a simple function of these embeddings ( Figure 1b ) . We include a variety of different KGE models in our experiments , including TransE [ Bordes et al. , 2013 ] , DistMult [ Yang et al. , 2015 ] , ComplEx [ Trouillon et al. , 2016 ] , and RotatE [ Sun et al. , 2019 ] . LM-based models . KGE methods do not capture the rich information available from textual descriptions of nodes . To address this limitation , previous KG completion ap- proaches have incorporated textual representations [ Toutanova et al. , 2015 , Wang and Li , 2016 ] , most recently with approaches such as KG-BERT [ Yao et al. , 2019 ] that fine-tune the BERT language model ( LM ) [ Devlin et al. , 2019 ] for the task of KG completion . Our focus in this work is on LMs pretrained on corpora of biomedical documents ( e.g. , PubMedBERT [ Gu et al. , 2020 ] ; see Appendix B.1.2 for full details ) . To score a triple using an LM , we use a cross-encoder approach [ Yao et al. , 2019 , Kim et al. , 2020 ] ( Fig . 1a ) , where we encode the text of the head and tail entities together along with the appropriate special tokens . Specifically , a triple ( h , r , t ) is encoded as v = LM ( [ CLS ] text ( h ) [ SEP ] text ( t ) [ SEP ] ) , where v is the contextualized representation of the [ CLS ] token at the last layer.2 We then apply an additional linear layer with a single output dimension to v to compute the ranking score for the triple ( f ( x ) = Wrankv ∈ R ) , and train the LM with the same max-margin loss . Recent work on applying BERT for KG completion on general domain benchmarks has shown that multi-task training improves performance [ Wang et al. , 2021 , Kim et al. , 2020 ] . We use the approach of Kim et al . [ 2020 ] and incorporate two additional losses for each LM : a binary triple classification loss to identify if a triple is positive or negative , and a multi-class relation classification loss.3 2.3 Integrating KGE and LM : Model Averaging vs . Routing . Previous work using text for KG completion on general domain benchmarks has demonstrated the benefit of combining KGE and text-based models [ Xie et al. , 2016 , Wang et al. , 2021 ] . We study integration of graph-based and text-based methods ( Figure 1c ) , exploring whether learning to route input instances adaptively to a single model can improve performance over previous approaches that compute a weighted average of ranking scores [ Wang et al. , 2021 ] . We also explore the more general setup of combining more than two models . More formally , for a given triple x = ( h , r , t ) , let φ ( x ) be its feature vector . We can learn a function g ( φ ( x ) ) that outputs a set of weights α = [ α1 , . . . , αk ] , ∑ i αi = 1 , αi > 0 ∀i . These weights can be used to perform a weighted average of the ranking scores { s1 , . . . , sk } for a set of k models we wish to combine , such that the final ranking score is s = ∑ i αisi . We use a variety of graph- , triple- , and text-based features to construct the feature vector φ ( x ) such as node degree , entity and relation type , string edit distance between head and tail entity names , and overlap in graph neighbors of head and tail nodes . We explore these features further in Section 4.1 , and provide a full list in Appendix B.1.3 ( Table 6 ) . For the function g ( · ) , we experiment with an input-dependent weighted average that outputs arbitrary weights α and a router that outputs a constrained α such that αi = 1 for some i and αj = 0 , ∀j 6= i ( i.e. , α is a one-hot vector ) .4 In practice , we implement the router as a classifier which selects a single KG completion model for each example by training it to predict which model will perform better.5 For the input-dependent weighted average we train a multilayer perceptron ( MLP ) using the max-margin ranking loss . We train all models on the validation set and evaluate on the test set for each dataset . 2 . We experiment with encoding the relation text as well , but find that this did not improve performance . 3 . See details in Appendix B. Wang et al . [ 2021 ] omit the relation classification loss and use a bi-encoder ; we find that both of these modifications reduce performance in our setting . 4 . We also try a global weighted average with a single set of weights ; see Appendix B.1.3 for details . 5 . We explore a range of methods for the router ’ s classifier , with the best found to be gradient boosted decision trees ( GBDT ) and multilayer perceptrons ( MLP ) . When performing ranking evaluation , we use the features φ ( x ) of each positive example to compute the weights α , then apply the same weights to all negative examples ranked against that positive example .

Summary: This paper presented a comprehensive empirical study of applying various scientific language models to the biomedical knowledge graph completion task. The authors found that language models and knowledge graph embedding models have complementary strengths, and further proposed a routing method, which is able to integrate these two models. Domain-specific biomedical language models are carefully evaluated.