A Framework for Adapting Pre-Trained Language Models to Knowledge Graph Completion (2022)

Abstract: Recent work has demonstrated that entity representations can be extracted from pre-trained language models to develop knowledge graph completion models that are more robust to the naturally occurring sparsity found in knowledge graphs. In this work, we conduct a comprehensive exploration of how to best extract and incorporate those embeddings into knowledge graph completion models. We explore the suitability of the extracted embeddings for direct use in entity ranking and introduce both unsupervised and supervised processing methods that can lead to improved downstream performance. We then introduce supervised embedding extraction methods that can extract more informative representations. We then synthesize our findings and develop a knowledge graph completion model that significantly outperforms recent neural models.

## introduction knowledge graphs (kg) are structured representations of knowledge that contain a collection of factual relations between entities. kgs are valuable resources with applications in different areas such as representation learning @xcite , question answering @xcite @xcite , and entity linking @xcite . however, the difficulty of curating knowledge at scale means that existing kgs are highly incomplete. this has led to the widespread study of knowledge graph completion (kgc) which aims to develop automated solutions that can suggest new facts to add to the kg @xcite @xcite . kgc is typically formulated as ranking problem where an incomplete fact is used as a query to retrieve entities that complete the fact. recent work has utilized pre-trained language models to develop approaches that are more robust to the naturally occurring sparsity within knowledge graphs. these approaches utilize textual entity descriptions to develop entity representations that are less reliant on graph connectivity. such work either fine-tunes the language model directly during training to encode the entities (e.g. @xcite ) or extracts a set of entity embeddings prior to training which can then be used to train a kgc model using standard training procedures (e.g. @xcite ). while fine-tuning language models often improves downstream performance @xcite , it increases the computational overhead of computing entity representations. as a result, standard kgc training procedures that involve evaluating a large number of negative candidates for each positive instance are typically infeasible. sampling only a small set of negative candidates enables training, but can harm performance. approaches that extract entity embeddings prior to training @xcite do not introduce any overhead for computing entity representations and are able to take advantage of standard training protocols. however, such approaches do not utilize any supervision to adapt the pre-trained language model to kgc. while both lines of previous work have demonstrated their approaches effectiveness at retrieving sparsely connected entities, they still lag behind kgc models that do not incorporate any textual information on standard benchmark datasets. in this work, we develop a framework for adapting pre-trained language models to kgc that takes advantage of the strengths of both prior lines of work. we accomplish this by decoupling the entity representations used for computing the query representation and the entity representations used for retrieval (see figure 1 ). for candidate ranking, we extract and cache en-tity representations from a pre-trained language model prior to training. we then introduce lightweight unsupervised and supervised embedding processing techniques that improve the suitability of the space for candidate retrieval without sacrificing the scalability necessary to use standard kgc training procedures. the embedding processing techniques introduced in this work lead to significant performance improvements across datasets from diverse domains. this decoupling also enables us to scalably finetune pre-trained language models to extract more informative entity representations for the query. however, naively fine-tuning the language model overfits the knowledge graph and actually degrades performance. we find that parameter-efficient finetuning methods such as prompt-tuning mitigate this and improve downstream performance. we synthesize our findings and utilize the most effective candidate representation processing and entity extraction techniques with a recently proposed neural ranking architecture. although we do not make any modifications to the ranking architecture, our representation extraction and processing techniques lead to significant improvements across four diverse datasets. the findings and analysis from this work provide useful guidelines for developing and utilizing effective textual entity representations for kgc. the rest of our paper is organized as follows. we discuss related work in section 2, present a formal description of our task in section 3, and describe the datasets used in this work in section 4. we introduce unsupervised and supervised techniques to improve the suitability of entity embeddings for candidate ranking in section 5. we then introduce supervised methods to extract more informative representations for the query entity in section 6 and explore the effect of language model selection in section 7. finally, we synthesize our findings in section 8 and compare against recent work on our datasets. our contributions are as follows. ## task formulation given a set of entities e and relations r, a kg can be defined as a collection of entity-relation-entity triplets @xmath0, r j , e k )} ⊂ e × r × e where e i , e @xmath1. the aim of kgc is to develop a model that accepts a query consisting of a head entity and a relation, (e i , r j , ?), and ranks all candidate entities e @xmath2. an effective kgc model should rank correct candidates more highly than incorrect candidates. neural kgc models embed the head entity and relation and compute a query vector f θ (e i , r j ) = q where f θ (•) is a neural network and e i , r j , @xmath3. scores for each candidate, e @xmath4, are computed as the inner product between the query vector and the candidate entity embedding y @xmath5. we follow @xcite and use textual descriptors to extract the entity embeddings from pre-trained language models while learning relation embeddings during training. we evaluate the kgc models with standard ranking metrics: mean reciprocal rank (mrr), hits at 1 (h@1), hits at 3 (h@3), and hits at 10 (h@10). we follow standard procedure and consider both forward and reverse relations and use the filtered evaluation setting @xcite . we validate the significance of improvements in the mrr with the paired bootstrap significance testing @xcite and correct for multiple hypothesis testing with the benjamini/hochberg method @xcite . ## datasets we work with kgc datasets that cover diverse domains such as commonsense, biomedical, and encyclopedic knowledge. for the commonsense kg dataset, we work with the cn-82k dataset introduced by @xcite which is derived from conceptnet. for the biomedical kgc dataset, we work with the snomed-ct core dataset introduced by @xcite . for the encyclopedic dataset, we utilize the widely used benchmark kgc dataset, fb15k-237 @xcite . we additionally utilize the widely studied wn18rr @xcite dataset which is derived from wordnet. dataset statistics are reported in the appendix in ## query entity extraction we explore supervised techniques to extract more informative representations from pre-trained language models for the query entity. fine-tuning: we fine-tune the language model during training and extract the entity representation by mean pooling across the intermediate states in each layer and aggregating across layers with a learned linear combination. ## effect of language model selection further performance improvements can often be gained by scaling up the size of the language model @xcite or from using specialized, domain-specific language models @xcite . in this section, we examine the effect of these two aspects on downstream kgc performance. we conduct experiments with both unsupervised and supervised query entity extraction techniques while using our best candidate ranking approach, the residual mlp. we conduct experiments with bert-base-uncased and bert-large-uncased for all three kgs. to evaluate the effect of specialization, we use pubmedbert, which is the same size as bert-base, for snomed-ct core. we report the results of these experiments in the supervised extraction and larger language models do lead to lower training loss, but that improvement does not consistently translate to stonger test performance. thus, the mixed results likely arise from overfitting which could potentially be mitigated with careful regularization. domainspecific pretraining is particularly effective, with pubmedbert consistently outperforming other models. ## comparison against recent work we synthesize our findings to develop a kgc model and compare against recent work. we again simply repurpose the bert-resnet ranking architecture with the default hyperparameters from @xcite to demonstrate the impact of the decisions explored in this work. we report results across the two sparser datasets in our kgc model is very effective and outperforms the models that do not incorporate any additional information. although this seems natural, this was actually not the case with previous work. therefore, our method integrates textual information in a way that leads to competitive performance even for these widely studied benchmark datasets. ## conclusion we present a framework for adapting pre-trained language models for kgc. the key insight driving the development of our framework was that decoupling the entity representations used for computing the query representation and the entity representations used for candidate retrieval enabled us to better integrate the information from pre-trained language models while maintaining the scalability necessary to train performant kgc models. we introduced unsupervised and supervised techniques to improve the suitability of entity embeddings for candidate ranking (section 5), introduced methods to extract entity embeddings from language models (section 6), and explored the effect of language model selection (section 7). by synthesizing the insights from our exploration, we developed a kgc model that significantly outperforms recent work while simply repurposing an existing ranking architecture. while innovations in neural ranking architecture have been valuable, our work demonstrates the importance of developing more informative entity representations. the findings and analysis from this work provide a useful framework for adapting pre-trained language models for knowledge graph completion. the residual mlp, while lightweight, does increase the training time per iteration. for the candidate transformation experiment on the wn18rr dataset (section 5), the baseline completes one epoch in 3m56s while the residual mlp increases this to 5m44s. however, the residual mlp also accelerates convergence, offsetting the per-iteration slowdown. although it takes a similar amount of time to train the baseline for 6 epochs as it does to train the residual mlp model for 4 epochs, the residual mlp actually outperforms the baseline at that time despite training for fewer iterations. therefore, the baseline is only more effective in the earliest stage of training before being surpassed by the residual mlp model. for the wn18rr dataset, this breakeven point occurs within only 29m of training. this trend holds across all datasets, with the worst breakeven point being only 1h43m. therefore the accelerated convergence offsets the increased per-iteration cost for all but the shortest of training times. techniques such as prompt-tuning require the application of a language model, which increases the time per iteration. for the query extraction experiment on the wn18rr dataset (section 6), the baseline completes one epoch in 3m54s, while prompttuning increases this to 8m47s. when controlling for wall clock time, we observe a similar trend where the baseline is more effective early in training before being surpassed by prompt-tuning. however, the breakeven point occurs much later (e.g. at 14h1m for wn18rr). therefore, in settings with limited training budgets, the performance improvement from prompt-tuning may not justify the additional training cost. we note that none of the techniques explored in our work introduce any overhead at inference time. after training, the improved entity representations from the residual mlp or prompt-tuning can be computed and cached for inference, reducing the cost of computing entity embeddings to a simple lookup like the original bert-resnet model.