Retrieval, Reasoning, Re-ranking: A Context-Enriched Framework for Knowledge Graph Completion (2025)

Abstract: The Knowledge Graph Completion (KGC) task aims to infer the missing entity from an incomplete triple. Existing embedding-based methods rely solely on triples in the KG, which is vulnerable to specious relation patterns and long-tail entities. On the other hand, text-based methods struggle with the semantic gap between KG triples and natural language. Apart from triples, entity contexts (e.g., labels, descriptions, aliases) also play a significant role in augmenting KGs. To address these limitations, we propose KGR3, a context-enriched framework for KGC. KGR3 is composed of three modules. Firstly, the Retrieval module gathers supporting triples from the KG, collects plausible candidate answers from a base embedding model, and retrieves context for each related entity. Then, the Reasoning module employs a large language model to generate potential answers for each query triple. Finally, the Re-ranking module combines candidate answers from the two modules mentioned above, and fine-tunes an LLM to provide the best answer. Extensive experiments on widely used datasets demonstrate that KGR3 consistently improves various KGC methods. Specifically, the best variant of KGR3 achieves absolute Hits@1 improvements of 12.3% and 5.6% on the FB15k237 and WN18RR datasets.

## introduction knowledge graphs (kgs) are graph-structured knowledge bases (kbs) that organize factual knowledge as triples in the form of (head entity, relation, tail entity). recently, kgs have become a crucial foundation for various downstream applications, such as recommendation systems @xcite , question answering @xcite , and sentiment analysis @xcite . nevertheless, mainstream kgs such as freebase @xcite and wordnet @xcite suffer from serious incomplete issues. this problem highlights the importance of the knowledge graph completion (kgc) task, which aims to predict the missing entity from an incomplete triple. existing kgc methods can be roughly categorized into embedding-based methods @xcite @xcite ) and text-based methods @xcite @xcite . embedding-based methods implicitly learn rules based on relation patterns observed in triples and make predictions based on the likelihood of these patterns occurring @xcite . for example, from triple (x, works in, y ) and (y, city of, z), it is very likely to deduce that (x, citizen of, z). however, these methods ignore the contextual semantics supporting these triples, leading to conclusions that do not align with the facts. text-based methods employ pre-trained language models (plms) to embed entities and relations with their labels and descriptions. how-ever, these methods still cannot surpass the latest embedding-based counterparts @xcite due to the substantial semantic gap between structural kg triples and natural language sentences. large language models (llms), trained by extensive corpora, demonstrate emergent semantic understanding and in-context learning (icl) capabilities. recent studies @xcite have proposed utilizing llms for the kgc task, as these models harbor general knowledge that can be leveraged to mitigate information scarcity for long-tail entities. however, the application of llms in kgc tasks encounters several limitations. firstly, if the pre-training corpora of the llms lack adequate contextual information on specific entities, the llms may produce hallucinated or biased responses. secondly, the structured nature of kg triples limits the ability of llms to effectively capture and leverage contextual information from the graph structure. these shortcomings necessitate a strong reliance on a considerable amount of in-context demonstrations @xcite or external structured embeddings @xcite , which inevitably limit the performance and generality of existing approaches. considering the aforementioned challenges, we propose a context-enriched kgc framework named kgr 3 , which consists of three modules: retrieval, reasoning, and re-ranking. given a query triple, the retrieval module gathers semantically relevant supporting triples with the same relation and similar entities, and extracts plausible candidate answers from a base kgc model. to ensure that the llm attains a fundamental understanding of the retrieved information, this module also collects and augments relevant contextual information to entities in supporting triples and the candidate answer list. then, the reasoning module exploits the semantic understanding capability of pre-trained llm to suggest several potential answers based on in-context demonstrations and the description of the known entity. finally, the re-ranking module fine-tunes the llm to select out the corrupted entity of the training triple from a set of candidate entities, enabling it to process structured knowledge. kgr 3 possesses strong plug-and-play capability, making it compatible with all base kgc methods without costly re-training. during inference, the re-ranking module integrates the candidate answers derived from the base kgc model and the reasoning module, and then instructs the llm to output the entity label that best completes the query triple. we validate the proposed framework on two conventional kgc datasets: fb15k237 and wn18rr. the extensive experiments show that kgr 3 significantly and consistently outperforms all baseline methods with different types of base kgc models and backbone llms, showing its superiority. notably, the best variant of kgr 3 achieves stateof-the-art performance with absolute hits@1 improvements of 12.3% and 5.6% on the two datasets. our contributions are summarized as follows: ## related works embedding-based methods embedding-based methods are fundamental in kgc, which focuses on learning a set of low-dimensional embeddings for entities and relations with certain geometric or mathematical constraints. most typically, transe @xcite assumes the translated head embedding of a triple is close to the embedding of the tail. distmult @xcite aims to maximize the hadamard product of the head, relation, and tail embeddings of each triple. to model symmetric and anti-symmetric relations, complex @xcite generalizes dist-mult by introducing complex embeddings with hermitian dot product. rotate @xcite interprets relation as a rotation operation in complex space, which can effectively infer inversion and composition patterns. atth @xcite and @xcite further leverage hyperbolic embeddings and operations to capture the intrinsic hierarchical structure in kgs. recently, graph neural networks (gnns) have emerged as powerful methods for graph embedding @xcite , with various applications @xcite @xcite . considering the heterogeneity of kgs, rgcn @xcite , wgcn @xcite , and compgcn @xcite adapted to consider relation types in their message-passing functions. kbgat @xcite proposes a two-layer attentional network to encode each triple and to measure its importance to the tail entity. in addition, nbf-net @xcite , red-gnn @xcite and a*net @xcite ) integrate structural information from paths between the two entities, which also supports inductive kgc. despite the simplicity and high scalability, embedding methods suffer from the long-tail entity distribution in kgs, and hence, cannot generate semantic expressive embeddings for boundary entities, which limits the performance of kgc. text-based methods. apart from the graph structure, textual information in kgs also entails rich semantic knowledge. dkrl proposes to initialize entity embeddings with a convolutional neural network. kg-bert @xcite tokenizes triples with textual descriptions of entities and relations, and utilizes bert @xcite to assess their plausibility. kepler @xcite jointly finetunes the pre-trained bert with the kg embedding and mlm objectives, showcasing improved kgc accuracy. simkgc @xcite introduces a contrastive learning strategy, which reduces the computational complexity by re-using in-batch and pre-batch entities as negative samples. cole @xcite extends kg-bert with co-distillation learning. inspired by the semantic understanding and reasoning capability of llms, kicgpt @xcite proposes a gpt-based in-context learning (icl) paradigm for the kgc task. however, due to the large semantic gap between kg triples and natural language sentences, all approaches mentioned above cannot outperform the latest embedding-based methods @xcite . recently, dift @xcite devises a supervised finetuning (sft) solution to guide the llm in completing triples. despite achieving state-of-the-art performance, dift necessitates costly re-training to adapt to different base kgc models, thereby limiting its compatibility and generality. ## problem specifications a knowledge graph (denoted as @xmath0, r, t }) can be represented as a set of triples in the form of (h, r, t) ∈ t , where h, @xmath1, @xmath2. the notations h and t denote the head and the tail entity of a triple. e, r, t are the set of entities, relations, and triples, respectively. besides, kgs are usually associated with knowledge bases (kbs), such as yago @xcite , wikidata @xcite , and dbpedia @xcite . kbs are renowned for representing general knowledge about real-world objects, including people, organizations, places, products, and among others. apart from triples, kbs also store rich contextual information (or "contexts") for entities in kgs, which includes entity labels, entity descriptions, aliases, etc. we argue that these entity contexts contain valuable semantic knowledge for the kgc task. in this paper, we exploit llms to reveal the missing entity of an incomplete triple based on relevant triples and the contexts of entities involved. ## methods in this section, we introduce our proposed contextenriched kgc framework kgr 3 , which consists of three components: (1) retrieval, (2) reasoning, and (3) re-ranking. ## experiments in this section, we assess the effectiveness of the kgr 3 framework in the kgc task. ## conclusion in this paper, we propose kgr 3 , an llm-based context-enriched kgc framework with three modules: retrieval, reasoning, and re-ranking. by leveraging contextual information, kgr 3 effectively bridges the semantic gap between structural kg triples and natural language. experimental results show that incorporating supporting triples and entity descriptions with llm in-context learning and sft significantly improves the kgc performance. future work will focus on adapting llms to other kg reasoning tasks such as inductive kgc and knowledge-based question answering.