Venue: NAACL 2025

# Retrieval, Reasoning, Re-ranking: A Context-Enriched Framework for Knowledge Graph Completion

Muzhi Li\({}^{2,1}\) , Cehao Yang\({}^{3,1}\) , Chengjin Xu\({}^{1*}\)+ , Xuhui Jiang\({}^{1}\), Yiyan Qi\({}^{1}\), Jian Guo\({}^{1}\), Ho-fung Leung\({}^{3}\), Irwin King\({}^{2}\)

Equal contribution.Corresponding authors.Independent researcher.

###### 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 KGR\({}^{3}\), a context-enriched framework for KGC. KGR\({}^{3}\) 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 KGR\({}^{3}\) consistently improves various KGC methods. Specifically, the best variant of KGR\({}^{3}\) achieves absolute Hits@1 improvements of \(12.3\%\) and \(5.6\%\) on the FB15k237 and WN18RR datasets.

## 1 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 Chen et al. (2022a, b), question answering Sun et al. (2024), and sentiment analysis Wang and Shu (2023). Nevertheless, mainstream KGs such as Freebase Bollacker et al. (2008) and Wordnet Miller (1995) 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 Bordes et al. (2013); Yang et al. (2015); Sun et al. (2019); Cao et al. (2022) and text-based methods Yao et al. (2019); Wang et al. (2021, 2022). Embedding-based methods implicitly learn rules based on relation patterns observed in triples and make predictions based on the likelihood of these patterns occurring Wu et al. (2023). For example, from triple \((X,\textit{works in},Y)\) and \((Y,\textit{city of},Z)\), it is very likely to deduce that \((X,\textit{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

Figure 1: Limitations of existing embedding-based (top) and LLM-based (bottom) KGC methods.



ever, these methods still cannot surpass the latest embedding-based counterparts [14] 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 [23; 50] 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 [26] or external structured embeddings [58], which inevitably limit the performance and generality of existing approaches.

Considering the aforementioned challenges, we propose a context-enriched KGC framework named \({\bf KGR^{3}}\), which consists of three modules: \({\mathbf{R}}\textit{etrieval}\), \({\mathbf{Reasoning}}\), and \({\mathbf{Re}}\textit{-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. \({\bf 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 \({\rm 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 \({\rm KGR^{3}}\) achieves state-of-the-art performance with absolute Hits@\(1\) improvements of \(12.3\%\) and \(5.6\%\) on the two datasets. Our contributions are summarized as follows:

* We propose a novel \({\rm KGR^{3}}\) framework for the KGC task, which systematically retrieves relevant supporting contexts, conducts semantic reasoning, and re-ranks candidate answers.
* We notice the semantic gap between KG triples and natural language sentences, and seamlessly bridge this gap with entity contexts.
* We conduct extensive experiments and ablation studies to evaluate the effectiveness of the \({\rm KGR^{3}}\) framework, and discuss the importance of incorporating entity contexts and LLMs.

## 2 Related Works

Embedding-based methodsEmbedding-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 [9] assumes the translated head embedding of a triple is close to the embedding of the tail. DistMult [28] aims to maximize the Hadamard product of the head, relation, and tail embeddings of each triple. To model symmetric and anti-symmetric relations, ComplEx [15] generalizes DistMult by introducing complex embeddings with Hermitian dot product. RotatE [29] interprets relation as a rotation operation in complex space, which can effectively infer inversion and composition patterns. ATTH [1] and GIE [16] 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 [25, 26], with various applications [27, 28, 29] . Considering the heterogeneity of KGs, RGCN [12], WGN [26], and CompGCN [24] 

2020) adapted to consider relation types in their message-passing functions. KBGAT Nathani et al. (2019) proposes a two-layer attentional network to encode each triple and to measure its importance to the tail entity. In addition, NBF-Net Zhu et al. (2021), RED-GNN Zhang and Yao (2022) and A*Net Zhu et al. (2023b) 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 methodsApart 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 Yao et al. (2019) tokenizes triples with textual descriptions of entities and relations, and utilizes BERT Devlin et al. (2019) to assess their plausibility. KEPLER Wang et al. (2021) jointly finetunes the pre-trained BERT with the KG embedding and MLM objectives, showcasing improved KGC accuracy. SimKGC Wang et al. (2022) introduces a contrastive learning strategy, which reduces the computational complexity by re-using in-batch and pre-batch entities as negative samples. CoLE Liu et al. (2022) extends KG-BERT with co-distillation learning. Inspired by the semantic understanding and reasoning capability of LLMs, KICGPT Wei et al. (2023) 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 Ge et al. (2023); Shang et al. (2024). Recently, DIFT Liu et al. (2024) devises a supervised fine-tuning (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.

## 3 Problem Specifications

A Knowledge Graph (denoted as \({\cal G}=\{{\cal E},{\cal R},{\cal T}\}\)) can be represented as a set of triples in the form of \((h,r,t)\in{\cal T}\), where \(h,t\in{\cal E}\), \(r\in{\cal R}\). The notations \(h\) and \(t\) denote the head and the tail entity of a triple. \({\cal E},{\cal R},{\cal T}\) are the set of entities, relations, and triples, respectively. Besides, KGs are usually associated with knowledge bases (KBs), such as YAGO Suchanek et al. (2007), Wikidata Vrandecic and Krotzsch (2014), and DBPedia Lehmann et al. (2014). 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.

## 4 Methods

In this section, we introduce our proposed context-enriched KGC framework KGR\({}^{3}\), which consists of three components: (1) _Retrieval_, (2) _Reasoning_, and (3) _Re-ranking_.

### Stage 1: Retrieval

The retrieval module focuses on gathering structural and semantic knowledge that may contribute to the completion of certain incomplete triples.

#### 4.1.1 Supporting Triple Retrieval

In KGs, the attributes of an entity are represented in structural triples. Different entities connected by the same relation often share common salient properties. The internal knowledge inherent in the graph structure provides the most direct support to the validity of a triple. Given an incomplete query triple in the form of \((h,r,?)\) or \((?,r,t)\), we aim to retrieve \(k\) supporting triples that are the most semantically similar to the incomplete query triple. Intuitively, we prioritize triples with the same entity and relation from the training set. If the number of available triples is less than \(k\), we broaden our choices to triples with the same relation, and with entities that are semantically similar to the known one in the query triple.

#### 4.1.2 Textual Context Retrieval

We note that there is a significant semantic gap between structural triples and natural language sentences. For example, in Figure 2, entity "Kasper Schmeichel" is originally represented by an unique entity id "/m/07Th1h5" while relation "plays for sports teams" is originally represented 

as _"/sports/pro_athlete/teams./sports/sports_team roster/team"_. Such a structured format is difficult for LLMs to process. To fully leverage the semantic understanding capabilities of LLMs, we extract relevant contexts related to entities in the query triple and supporting triples from Wikidata knowledge base Vrandecic and Krotzsch (2014).

In mainstream KGs, entities are represented in numerical or textual IDs. Each entity ID acts as an index to the data frame in its corresponding knowledge base. Apart from triples, the data-frame of an entity contains significant contextual information such as its entity label. Since Google Freebase Bollacker et al. (2008) is deprecated and migrated to Wikidata Vrandecic and Krotzsch (2014), we map the entity IDs in the FB15k237 dataset to corresponding Wikidata QIDs with official data dumps. We then collect the textual entity label, the short description, and aliases from Wikidata URIs. As for entities in the WN18RR dataset, we adopt the same set of entity labels and descriptions used in KGBERT Yao et al. (2019) and SimKGC Wang et al. (2022).

#### 4.1.3 Candidate Answer Retrieval

The widely adopted ranking-based evaluation for the KGC task requires the model to score the plausibility of each entity in the KG as a potential replacement for the missing entity in the query triple. However, given the vast number of entities in the KG, employing LLMs to score and rank each entity is computationally expensive and impractical. Inspired by Lovelace et al. (2021); Wei et al. (2023); Li et al. (2024), we employ a base KGC model to initialize the scoring and ranking of entities within the KG. Formally, we denote the ranked entity list \({\cal A}_{\text{KGC}}\) as following:

\[{\cal A}_{\text{KGC}}=[e^{(k)}_{1},e^{(k)}_{2},...,e^{(k)}_{| \mathcal{E}|}],\] (1) where \[e^{(k)}_{i}=\operatorname*{arg\,max}_{e\in\left\{\xi\backslash\{ e^{(k)}_{<i}\}\right\}}f_{r}(h,e)\text{ or }f_{r}(e,t),\]

\(f_{r}(h,t)\) is the scoring function the KGC model evaluating the plausibility of a triple \((h,r,t)\). Then, we extract the top-\(n\) entities with the highest scorings from \({\cal A}_{\text{KGC}}\) as candidate answers, and extract their labels and descriptions based on their entity IDs. The plausibility of these candidate answers will be re-evaluated in the re-ranking module.

### Stage 2: Reasoning

In the second stage, we first exploit the supporting triples to direct the LLM in performing the KGC task. In addition, we employ the LLM to generate several possible answers drawing upon the description of the known entity in the query triple.

#### 4.2.1 Supporting Triple Demonstrations

In this sub-section, we present the design of our prompt tailored for the demonstration, which is crucial in facilitating in-context learning Wei et al. (2023). For each supporting triple, we first provide LLM with the description of the head entity. The entity description serves three objectives: (1) _disambiguate entities with the same label_, (2) _rejuvenate the LLM's memory about known entities_, and (3) _provide essential information for entities that are not included in the LLM pre-training corpus_.

Figure 2: The end-to-end architecture of the proposed KGR\({}^{3}\) framework, which consists of three modules: 1) Reasoning (left), 2) Reasoning (middle), and 3) Re-ranking (right). (Please see Appendix A for detailed prompts.)



Following the convention of Zhu et al. (2023); Wei et al. (2023), we corrupt the neighboring entity with a "[MASK]" token. In order to narrow the semantic gap between structural triples and natural language sentences, we translate the masked triple into a natural language question. Subsequently, we ask the LLM to generate answers based on its semantic understanding of the short description of the known entity and the question. Different from KICGPT Wei et al. (2023), all supporting triples included in our demonstration prompts are sampled from the training set, which avoids potential information leakage. Finally, we provide the label of the corrupted entity and its entity description as our explanation. Considering multiple supporting triples, the LLM can also discern salient properties of adjacent neighbors connected by the same relation and similar entities. These salient properties play a vital role in helping the LLM to find out proper answers for the query triple.

#### 4.2.2 Context-aware Reasoning

We anticipate that LLMs can harness their information extraction and semantic understanding capabilities by utilizing comprehensive contextual information about the known entity, thereby generating potential answers. Similarly, we pass the description of the known entity and the question translated from the query triple to the LLM. The LLM is then instructed to output a list of answers in its response. It should be noted that generative LLMs do not guarantee that output answers will conform to entities in the KG. Therefore, we postprocess the LLM output by replacing entity aliases with entity labels and filtering out invalid and unreliable answers that do not appear within the top-\(\delta\) positions of \({\mathcal{A}_{\text{KGC}}}\). Finally, we obtain a list of \(m\) answers, which are formally denoted as:

\[{\mathcal{A}_{\text{LLM}}} =[e^{(l)}_{1},e^{(l)}_{2},...,e^{(l)}_{m}]\] \[=f_{\text{LLM}}(q,c^{(q)}_{e},\mathcal{D}(q))\cap{\mathcal{A}_{ \text{KGC}}}[0:\delta], \tag{2}\]

where \(q,c^{(q)}_{e}\) and \(\mathcal{D}(q)\) denote the question, description of the known entity, and supporting triple demonstrations. Entities in \({\mathcal{A}_{\text{LLM}}}\) are ensured to be simultaneously supported by the LLM and the base KGC model.

### Stage 3: Re-ranking

Motivated by the complementary nature of semantic and structural knowledge, we aim to exploit the candidate answer list generated by the LLM and the base KGC model to compose our final rankings.

To better enable the LLM in utilizing entity descriptions and structured neighbor facts for ranking candidate answers to query triples, we introduce supervised fine-tuning (SFT) with LoRA adaptation Chao et al. (2024). Inspired by DIFT Liu et al. (2024), the training objective of SFT is to find out the missing entity of an incomplete triple from a set of candidate answers. Specifically, we construct training samples by corrupting the tail (or head) entity of each triple in the training set. For each corrupted triple, we randomly sample \(n-1\) negative samples from the entity set, where half of them are connected by the same relation as the corrupted ground truth entity. Incorporating these hard negative samples helps the LLM to distinguish between different entities with the same property, which is crucially important since candidate entities suggested by base KGC models usually yield similar characteristics. After that, we translate the masked triple to a natural language question, and gather the label and description for each candidate entity. Finally, we provide the question \(q\), retrieved neighbor facts \(\mathcal{N}(q)\), and the description of the known entity \(c^{(q)}_{e}\), along with candidate answers \({\mathcal{A}}\) and their descriptions \(c({\mathcal{A}})\) to the LLM, and fine-tune the LLM to output the label \(y\) of the ground truth entity. Formally, we have the SFT loss:

\[{\mathcal{L}_{\text{SFT}}}=-\sum_{i=1}^{|\mathcal{T}|}\log(y|q,\mathcal{N}( q),c^{(q)}_{e},{\mathcal{A}},c({\mathcal{A}})). \tag{3}\]

It is important to emphasize that the SFT process does not rely on the prior inference results from existing KGC approaches. This ensures that the KGR\({}^{3}\) framework can be implemented as a plug-and-play solution.

During the inference stage, we construct a candidate answer set \({\mathcal{A}}_{c}\), which is composed of top-\(p\) entities from \({\mathcal{A}_{\text{KGC}}}\) and top-\((n-p)\) entities in \({\mathcal{A}_{\text{LLM}}}\) that are not previously encountered. Formally, we have:

\[{\mathcal{A}}_{c}={\mathcal{A}_{\text{KGC}}}[0:p]\cup\{{\mathcal{A}_{\text{ LLM}}}\backslash{\mathcal{A}_{\text{KGC}}}[0:p]\}[0:n-p]. \tag{4}\]

If \({\mathcal{A}_{\text{LLM}}}\) contains fewer than \((n-p)\) entities (e.g. \(m<n-p\)), we supplement the candidate answer set with additional entities from \({\mathcal{A}_{\text{KGC}}}\) to reach a total of \(n\) entities. Similarly, we instruct the fine-tuned LLM to select the most appropriate candidate entity to complete the query triple, and output its 

entity label. The selected entity \(e_{\text{ans}}\) is then prioritized and moved to the front of the candidate answer list \(\mathcal{A}_{c}\). Finally, we construct the re-ordered entity list \(\mathcal{A}_{\text{RR}}\) for performance evaluation. Here, we have

\[\mathcal{A}_{\text{RR}}=[e_{\text{ans}}||\mathcal{A}_{c}\backslash\{e_{\text{ ans}}\}||\mathcal{A}_{\text{KGC}}[n:|\mathcal{E}|]]\,, \tag{5}\]

where \([\cdot||\cdot]\) denotes the concatenation operation.

## 5 Experiments

In this section, we assess the effectiveness of the KGR3 framework in the KGC task.

### Datasets

We utilize two widely-used benchmark datasets, namely FB15k237 and WN18RR to evaluate the proposed method. FB15k237 is derived from Freebase Bollacker et al. (2008), an encyclopedic KG containing general knowledge about topics such as celebrities, organizations, movies and sports. WN18RR is a subset of WordNet Miller (1995), a lexical KG with knowledge about English morphology. To prevent potential data leakage, FB15k237 and WN18RR excludes reversible relations from their backend KGs. Detailed statistics of the two datasets are shown in Table 1.

Footnote *: KGR3 obtains better performance with Qwen2-7B LLM.

### Baselines and Evaluation Metrics

We compare the KGR3 framework with four types of baseline methods: (1) traditional KG Embedding methods TransE Bordes et al. (2013), ComplEx Trouillon et al. (2016), RotatE Sun et al. (2019), TuckER Balazevic et al. (2019); (2) GNN-based embedding methods CompGCN Vashishth et al. (2020), NBF-Net Zhu et al. (2021); (3) text-based methods KG-BERT Yao et al. (2019), MEMKGC Choi et al. (2021), SimKGC Wang et al. (2022), CoLE Liu et al. (2022); and (4) LLM-based methods KICGPT Wei et al. (2023), DIFT Liu et al. (2024). Among these baselines, we select TransE, RotatE, GIE, SimKGC, CoLE, and NBF-Net as our base KGC models because these methods are highly representative and exhibits strong performance.

Footnote *: KGR3 obtains better performance with Qwen2-7B LLM.

We utilize the widely adopted evaluation metrics, namely Hits@\(k\) (\(k=1,3,10\)) and MRR to evaluate our proposed method. Hits@\(k\) measures the proportion of query triples which the ground truth entities are ranked within the top-\(k\) position. MRR measures the mean reciprocal rank for each ground truth entities. Higher results indicates a better performance. Implementation details of baseline models and our models are described in Appendix B.

### Main Results

Table 2 summarizes the performance of the KGR3 framework on six different base KGC methods. The experiment results show that the best variant of KGR3 significantly outperforms all baseline methods among all evaluation metrics. Compared to the previous state-of-the-art baseline, KGR3, with Llama3-8B as the backbone LLM, achieves absolute Hits@\(1\) improvements of \(11.1\%\) and \(4.0\%\) on the FB15k237 and the WN18RR datasets, respectively. * It demonstrates that KGR3 is highly effective for the KGC task.

Notably, the improvement in Hits@1 is more substantial than that in Hits@3 and Hits@10. This indicates that the KGR3 framework is particularly effective at identifying the most accurate answers. Since our framework primarily focuses on re-ordering top-\(n\) entities from the initial ranked entity list, the upper bound of Hits@1, Hits@3, and Hits@\(10\) are implicitly constrained by the Hits@\(n\) performance of the base KGC model. Given that Hits@1 is typically further from this upper bound, the potential for improvement will be greater. Additionally, by leveraging semantic knowledge from entity contexts, the LLM gains a more comprehensive understanding of the entities, thereby enabling more precise inferences, particularly for top-ranked candidate answers.

Compared to the selected base KGC models, the corresponding variants of KGR3 consistently boost the performance on both datasets to a large margin. It shows that KGR3 is compatible with various types of KGC models, confirming its strong generality and plug-and-play capability. In contrast, KICGPT fails to outperform the state-of-the-art GNN-based or text-based methods, which underutilizes the power of the LLM. In addition, variants of KGR3 consistently outperform counterparts of the LLM-based baseline DIFT Liu et al. (2024) with the same base KGC models. The performance

\begin{table}
\begin{tabular}{l r c} \hline \hline
**Dataset** & **FB15k237** & **WN18RR** \\ \hline \#Entities & 14,541 & 40,943 \\ \#Relations & 237 & 11 \\ \#Train & 272,115 & 86,835 \\ \#Valid & 17,535 & 3,034 \\ \#Test & 20,466 & 3,134 \\ \hline \hline \end{tabular}
\end{table}
Table 1: Statistics of Datasets 

improvements can be primarily attributed to the incorporation of the reasoning module and the inclusion of entity descriptions for candidate answers. It should also be noted that DIFT necessitates an expensive re-training process and pre-trained KG embeddings to adapt to various KGC models Liu et al. (2024). Without such a process, DIFT cannot guarantee the experimental results as claimed.

### Ablation Studies

We verify the effectiveness of each component in the KGR\({}^{3}\) framework by answering the following research questions (RQs). Table 3 and 4 shows the experimental results for ablation studies.

**RQ1: Does the reasoning module improve the final inference performance?** To address this question, we simply apply the re-ranking module to candidate answers retrieved from the base KGC models. The performance drop in the "KGR\({}^{3}\) w/o Reasoning" variant demonstrates that the pretrained LLM can provide plausible answers that are not initially ranked at the top positions by the base models. This capability effectively breaks through the limitations of base KGC models and increases the performance ceiling of our method.

**RQ2: Whether entity descriptions contribute to enhancing KG completion?** In the "KGR\({}^{3}\) w/o Entity Descriptions" variant, we remove the descriptions for the known entity of a query triple and each of its candidate answers, resulting in sub-optimal experiment results. The performance decline re-confirms our hypothesis that LLMs may lack a fundamental understanding of certain entities within the KG, showing the importance of retrieving and leveraging textual contexts.

**RQ3: Can LLMs generate desirable KGC results without the help of KG triples?** We observe a significant performance decline with the "KGR\({}^{3}\) w/o Neighbor Facts" variant, particularly for Hits@1. In this case, the LLM can only rely on limited semantic knowledge derived from entity descriptions and its inherent knowledge base, which

\begin{table}
\begin{tabular}{l|c c c c|c c c c} \hline \hline \multirow{2}{*}{**Methods**} & \multicolumn{5}{c|}{**FB15k237**} & \multicolumn{5}{c}{**WN18RR**} \\  & **MRR** & **Hits@1** & **Hits@3** & **Hits@10** & **MRR** & **Hits@1** & **Hits@3** & **Hits@10** \\ \hline \multicolumn{9}{c}{_KG embedding methods_} \\ TransE & 0.279 & 0.198 & 0.376 & 0.441 & 0.243 & 0.043 & 0.441 & 0.532 \\ ComplEx & 0.247 & 0.158 & 0.275 & 0.428 & 0.440 & 0.410 & 0.460 & 0.510 \\ RotatE & 0.338 & 0.241 & 0.375 & 0.533 & 0.476 & 0.428 & 0.492 & 0.571 \\ TuckER & 0.358 & 0.266 & 0.394 & 0.544 & 0.470 & 0.443 & 0.482 & 0.526 \\ GIE & 0.362 & 0.271 & 0.401 & 0.552 & 0.491 & 0.452 & 0.505 & 0.575 \\ HitTER & 0.373 & 0.279 & 0.409 & 0.558 & 0.503 & 0.462 & 0.516 & 0.584 \\ \hline \multicolumn{9}{c}{_Graph neural network-based methods_} \\ CompGCN & 0.355 & 0.264 & 0.390 & 0.535 & 0.479 & 0.443 & 0.494 & 0.546 \\ NBF-Net & 0.415 & 0.321 & 0.450 & 0.599 & 0.551 & 0.497 & 0.573 & 0.666 \\ \hline \multicolumn{9}{c}{_Text-based methods_} \\ KG-BERT & - & - & - & 0.420 & 0.216 & 0.041 & 0.302 & 0.524 \\ MEM-KGC & 0.346 & 0.253 & 0.381 & 0.531 & 0.557 & 0.475 & 0.604 & 0.704 \\ SimKGC & 0.338 & 0.252 & 0.364 & 0.511 & 0.671 & 0.595 & 0.719 & 0.802 \\ CoLE & 0.389 & 0.294 & 0.429 & 0.572 & 0.593 & 0.538 & 0.616 & 0.701 \\ \hline \multicolumn{9}{c}{_Large language model-based methods_} \\ ChatGPT (0-shot) & - & 0.237 & - & - & - & 0.190 & - & - \\ ChatGPT (1-shot) & - & 0.267 & - & - & - & 0.212 & - & - \\ KICGPT & 0.412 & 0.327 & 0.448 & 0.581 & 0.564 & 0.478 & 0.612 & 0.677 \\ DIFT + TransE & 0.389 & 0.322 & 0.408 & 0.525 & 0.491 & 0.462 & 0.496 & 0.560 \\ DIFT + SimKGC & 0.402 & 0.338 & 0.418 & 0.528 & 0.686 & 0.616 & 0.730 & 0.806 \\ DIFT + CoLE & 0.439 & 0.364 & 0.468 & 0.586 & 0.617 & 0.569 & 0.638 & 0.708 \\ \hline \multicolumn{9}{c}{**KGR\({}^{3}\) + TransE**} \\ \multicolumn{9}{c}{**KGR\({}^{3}\) + HotatE**} & 0.456 & 0.414 & 0.474 & 0.550 & 0.506 & 0.487 & 0.515 & 0.556 \\ \multicolumn{9}{c}{**KGR\({}^{3}\) + GotE**} \\ \multicolumn{9}{c}{**KGR\({}^{3}\) + GlE**} \\ \multicolumn{9}{c}{**KGR\({}^{3}\) + SlimKGC**} \\ \multicolumn{9}{c}{**KGR\({}^{3}\) + SoLE**} \\ \multicolumn{9}{c}{**KGR\({}^{3}\) + NBF-Net**} \\ \hline \hline \end{tabular}
\end{table}
Table 2: Experiment results of the KGC task on FB15k-237 and WN18RR datasets. The best results are in **bold** and the second-best ones are underlined. All results of baseline methods are referred from corresponding original papers. For KGR\({}^{3}\), we adopt LLama3-8B as the backbone LLM of the Reasoning and Re-ranking module.

 

proves insufficient for generating precise predictions. This underscores that KG triples provide accurate and irreplaceable structural knowledge that is not inherently present in the LLM. In general, experimental results in Table 3 reconfirm the complementary relationship between textual and structural contexts.

**Case study.**   In addition, we conduct a case study on the hallucination case that may happen during the reasoning stage. For example, when being instructed to complete the incomplete triple (<missing-entity>, _/olympics/olympic_games/participating_countries_, _/m/04vjh_[Mauritania]) which queries the Olympic Games in which Mauritania has participated, the pre-trained LLM (GPT) incorrectly responds "_Mauritania has never participated in the Olympic Games_". This error highlights how hallucinations can arise from the model's lack of critical factual knowledge.

Nevertheless, triple _(Un/04yih_[Mauritania], _/olympics/olympic_participating_country/athletes - _/olympics/olympic_athlete_affiliation/Olympics_, /m/06ksk96 [2012 Summer Olympics]) shows that Mauritania does attend 2012 Summer Olympic games, which can help the LLM to figure out Mauritania also participated in [2008 Summer Olympics] during the re-ranking stage. This case study also shows that access to accurate factual knowledge is essential for LLMs to perform reasoning tasks like KGC.

**RQ4: Can KGR\({}^{3}\) ensure desirable performance with different LLMs?**   Experimental results in Table 4 show that KGR\({}^{3}\) consistently outperforms all baseline methods across all \(8\) LLM combinations. Notably, KGR\({}^{3}\) achieves the state-of-the-art performance when using Qwen2-7B and GPT3.5 as the LLMs of the re-ranking and the reasoning module. Compared to variants using pre-trained open-source LLMs, those employing GPT3.5 in reasoning stage produce better predictions, showing its stronger semantic understanding and instruction following capabilities.

It should be noted that KGR\({}^{3}\), when using the same LLM, Llama2-7B, still outperforms DIFT by a large margin. Hence, we cannot simply attribute the performance improvements to the power of LLMs. In addition, KGR\({}^{3}\) can also produce plausible predictions with an 1.5B model, which reduces the average SFT time from \(28.3\text{h}\) to \(9.05\text{h}\), demonstrating its strong robustness and effectiveness in low-resource settings.

\begin{table}
\begin{tabular}{l l|c c c c|c c c} \hline \hline \multicolumn{2}{c|}{**LLMs**} & \multicolumn{5}{c|}{**FB15k237**} & \multicolumn{5}{c}{**WN18RR**} \\
**Re-ranking** & **Reasoning** & **MRR** & **Hits@1** & **Hits@3** & **Hits@10** & **MRR** & **Hits@1** & **Hits@3** & **Hits@10** \\ \hline Llama2-7B & Llama2-7B & 0.524 & 0.462 & 0.555 & 0.627 & 0.709 & 0.645 & 0.754 & 0.803 \\ Llama2-7B & GPT3.5 & 0.530 & 0.466 & 0.562 & 0.635 & 0.710 & 0.644 & 0.757 & 0.808 \\ Llama3-8B & Llama3-8B & 0.535 & 0.475 & 0.564 & 0.634 & 0.717 & 0.656 & **0.759** & **0.809** \\ Llama3-8B & GPT3.5 & 0.536 & 0.477 & 0.565 & 0.636 & 0.717 & 0.655 & **0.759** & 0.807 \\ \hline Qwen2-1.5B & Qwen2-1.5B & 0.526 & 0.465 & 0.555 & 0.627 & 0.706 & 0.641 & 0.751 & 0.803 \\ Qwen2-1.5B & GPT3.5 & 0.531 & 0.469 & 0.560 & **0.648** & 0.706 & 0.637 & 0.753 & 0.806 \\ Qwen2-7B & Qwen2-7B & 0.539 & 0.482 & 0.566 & 0.634 & **0.724** & **0.672** & 0.754 & 0.805 \\ Qwen2-7B & GPT3.5 & **0.543** & **0.487** & **0.570** & 0.637 & **0.724** & **0.672** & 0.754 & 0.807 \\ \hline \hline \end{tabular}
\end{table}
Table 4: Results for ablation studies with different combinations of LLMs in the Reasoning and Re-ranking stages. We adopt the best base KGC model, namely NBF-Net and SimKGC for the FB15k237 and the WN18RR datasets.

\begin{table}
\begin{tabular}{l|c c c c c c} \hline \hline \multirow{2}{*}{**Settings**} & \multicolumn{3}{c|}{**FB15k237**} & \multicolumn{3}{c}{**WN18RR**} \\
**Settings** & **MRR** & **Hits@1** & **Hits@3** & **Hits@10** & **MRR** & **Hits@1** & **Hits@3** & **Hits@10** \\ \hline Reordering (RotatE) & 0.382 & 0.293 & 0.417 & 0.559 \\ KGR\({}^{3}\) (RotatE) & 0.456 & 0.400 & 0.476 & 0.569 \\ \hline Reordering (GIE) & 0.391 & 0.301 & 0.426 & 0.573 \\ KGR\({}^{3}\) (GIE) & 0.463 & 0.400 & 0.485 & 0.581 \\ \hline \hline \end{tabular}
\end{table}
Table 5: Ablation Experiments on FB15k-237 dataset with different re-ranking strategies.

 

RQ5: How does the ranking strategy affect the performance of KGR\({}^{3}\)?Rather than instructing the LLM to "sort" or "reorder" the entire candidate entity list, KGR\({}^{3}\) guides the LLM to identify and place the most likely candidate at the top, while preserving the relative order of the remaining candidates. As shown in Table 5, we conduct additional experiments to compare these two ranking strategies. Under the reorder setting, the LLM is fine-tuned to follow the original ranking produced by the corresponding base KGC models. However, the results remain suboptimal, suggesting that LLMs are not well-suited for sorting entire entity lists. This limitation can be attributed to the nature of KGs, which do not inherently provide ground truth rankings of entities for each query triple. As a result, it is not intuitive to claim that ground truth tail entity A is more relevant than ground truth tail entity B for a given query \((h,r,?)\). Additionally, the experimental results demonstrate that constraining LLMs to mimic the ranking patterns of prior KGC methods does not alleviate their inherent limitations.

Furthermore, in some cases, LLMs may select an incorrect, or a false negative candidate answer, which reduces the rankings of the ground truth entity. For example, during the re-ranking stage, the Llama3-8B model suggests "_Solihull_" as the answer of query triple (_England_, _location contains_, <missing-entity>), where the ground truth is "_Pontefract_". In fact, both "_Solihull_" and "_Pontefract_" are towns in England. Due to the inherent defect of the KG, the former is judged as a wrong answer. The re-ranking strategy adopted in KGR\({}^{3}\) ensures that when an LLM makes an incorrect judgment, the ranking of the ground truth answer drops by no more than 1 position, and hence, results in desirable performance improvements.

RQ6: How does the volume of SFT data affect the performance of KGR\({}^{3}\)?From Figure 3, we can conclude that increasing SFT data generally improves the model performance. We attribute the boundary effects observed in the WN18RR dataset to the characteristics of its entities, which consist of common English words that can be readily interpreted by LLMs. It is also noteworthy that, even with a limited number of training samples, KGR\({}^{3}\) still achieves desirable results. Given the substantial performance improvements, taking additional computational costs is deemed justifiable.

## 6 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.

## Acknowledgement

The work described in this paper was partially supported by the Research Grants Council of the Hong Kong Special Administrative Region, China (CUHK 14222922, RGC GRF 2151185). We would also like to express our sincere gratitude to the reviewers and area chairs for their constructive comments and suggestions.

Figure 3: Hits@1 performance over the amount of SFT data on the two datasets with different LLMs.

 



### Limitations

Although the proposed framework achieves a significant breakthrough in the KGC task, it still has some remaining issues to be resolved in the future. Firstly, the proposed KGR\({}^{3}\) framework is not capable of handling the KGC task under an "inductive setting". The KGC task discussed in this paper, along with most related works, operates under a "transductive setting", where entities in test triples also exist in the training set. We plan to tackle unseen entities that are not present in the KG in the future. Secondly, the commonly adopted evaluation metrics Hits@\(k\) necessitate ranking the plausibility of all entities within the KG for each query triple. Due to the limitations imposed by the maximum sequence length of LLMs, it is impractical to rank tens of thousands of entities within a KG. To _meet the evaluation criteria_, one may have to select a proper subset of entity candidates using existing KGC models. Hence, the base KGC model may inevitably constrain the final performance of the proposed framework.

## Ethics Statement

To the best of our knowledge, this work does not involve any discrimination, social bias, or private data. All the datasets are constructed from open-source KGs such as Wikidata. Therefore, we believe that our study complies with the ACL Ethics Policy.

## References

* J. Balazevic, C. Allen, and T. Hospedales (2019)TuckER: tensor factorization for knowledge graph completion. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), Hong Kong, China, pp. 5185-5194. External Links: Link Cited by: SS1.
* K. Bollacker, C. Evans, P. Paritosh, T. Sturge, and J. Taylor (2008)Freebase: a collaboratively created graph database for structuring human knowledge. In Proceedings of the 2008 ACM SIGMOD International Conference on Management of Data, SIGMOD '08, New York, NY, USA, pp. 1247-1250. External Links: ISBN 978-1-4673-387X, Document, ISSN 1053-6901, Link, [https://doi.org/10.1147/S00217719781](https://doi.org/10.1147/S00217719781)

Yang Liu, Zequn Sun, Guangyao Li, and Wei Hu. 2022. I know what you do not know: Knowledge graph embedding via co-distillation learning. In _Proceedings of the 31st ACM International Conference on Information & Knowledge Management_, CIKM '22, page 1329-1338, New York, NY, USA. Association for Computing Machinery.
* Liu et al. (2024) Yang Liu, Xiaobin Tian, Zequn Sun, and Wei Hu. 2024. Finetuning generative large language models with discrimination instructions for knowledge graph completion. _Preprint_, arXiv:2407.16127.
* Loshchilov and Hutter (2017) Ilya Loshchilov and Frank Hutter. 2017. Decoupled weight decay regularization. _arXiv preprint arXiv:1711.05101_.
* Lovelace et al. (2021) Justin Lovelace, Denis Newman-Griffis, Shikhar Vashishth, Jill Fain Lehman, and Carolyn Rose. 2021. Robust knowledge graph completion with stacked convolutions and a student re-ranking network. In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_, pages 1016-1029, Online. Association for Computational Linguistics.
* Ma et al. (2023) Yueen Ma, Zixing Song, Xuming Hu, Jingjing Li, Yifei Zhang, and Irwin King. 2023. Graph component contrastive learning for concept relatedness estimation. In _AAAI_, pages 13362-13370. AAAI Press.
* Miller (1995) George A. Miller. 1995. Wordnet: a lexical database for english. _Commun. ACM_, 38(11):39-41.
* Nathani et al. (2019) Deepak Nathani, Jatin Chauhan, Charu Sharma, and Manohar Kaul. 2019. Learning attention-based embeddings for relation prediction in knowledge graphs. In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_, pages 4710-4723, Florence, Italy. Association for Computational Linguistics.
* Schlichtkrull et al. (2018) Michael Schlichtkrull, Thomas Kipf, Peter Bloem, Rianne Berg, Ivan Titov, and Max Welling. 2018. _Modeling Relational Data with Graph Convolutional Networks_, pages 593-607.
* Shang et al. (2024) Bin Shang, Yinliang Zhao, Jun Liu, and Di Wang. 2024. Mixed geometry message and trainable convolutional attention network for knowledge graph completion. _Proceedings of the AAAI Conference on Artificial Intelligence_, 38(8):8966-8974.
* Song et al. (2022) Zixing Song, Yueen Ma, and Irwin King. 2022. Individual fairness in dynamic financial networks. In _NeurIPS 2022 Workshop: New Frontiers in Graph Learning_.
* Song et al. (2024) Zixing Song, Ziqiao Meng, and Irwin King. 2024. A diffusion-based pre-training framework for crystal property prediction. In _AAAI_, pages 8993-9001. AAAI Press.
* Song et al. (2023a) Zixing Song, Yifei Zhang, and Irwin King. 2023a. No change, no gain: Empowering graph neural networks with expected model change maximization for active learning. In _NeurIPS_.
* Song et al. (2023b) Zixing Song, Yifei Zhang, and Irwin King. 2023b. Optimal block-wise asymmetric graph construction for graph-based semi-supervised learning. In _NeurIPS_.
* Suchanek et al. (2007) Fabian M. Suchanek, Gjergji Kasneci, and Gerhard Weikum. 2007. Yago: A core of semantic knowledge. In _Proceedings of the 16th International Conference on World Wide Web_, WWW '07, page 697-706, New York, NY, USA. Association for Computing Machinery.
* Sun et al. (2024) Jiashuo Sun, Chengjin Xu, Lumingyuan Tang, Saizhuo Wang, Chen Lin, Yeyun Gong, Lionel M. Ni, Heung-Yeung Shum, and Jian Guo. 2024. Think-on-graph: Deep and responsible reasoning of large language model on knowledge graph. _Preprint_, arXiv:2307.07697.
* Sun et al. (2019) Zhiqing Sun, Zhi-Hong Deng, Jian-Yun Nie, and Jian Tang. 2019. Rotate: Knowledge graph embedding by relational rotation in complex space. In _International Conference on Learning Representations_.
* Trouillon et al. (2016) Theo Trouillon, Johannes Welbl, Sebastian Riedel, Eric Gaussier, and Guillaume Bouchard. 2016. Complex embeddings for simple link prediction. In _Proceedings of The 33rd International Conference on Machine Learning_, volume 48 of _Proceedings of Machine Learning Research_, pages 2071-2080, New York, New York, USA. PMLR.
* Vashishth et al. (2020) Shikhar Vashishth, Soumya Sanyal, Vikram Nitin, and Partha Talukar. 2020. Composition-based multi-relational graph convolutional networks. In _International Conference on Learning Representations_.
* Vrandecic and Krotzsch (2014) Denny Vrandecic and Markus Krotzsch. 2014. Wikidata: A free collaborative knowledgebase. _Commun. ACM_, 57(10):78-85.
* Wang and Shu (2023) Haoran Wang and Kai Shu. 2023. Explainable claim verification via knowledge-grounded reasoning with large language models. In _Findings of the Association for Computational Linguistics: EMNLP 2023_, pages 6288-6304, Singapore. Association for Computational Linguistics.
* Wang et al. (2022) Liang Wang, Wei Zhao, Zhuoyu Wei, and Jingming Liu. 2022. SimKGC: Simple contrastive knowledge graph completion with pre-trained language models. In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_, pages 4281-4294, Dublin, Ireland. Association for Computational Linguistics.
* Wang et al. (2021) Xiaozih Wang, Tianyu Gao, Zhaocheng Zhu, Zhengyan Zhang, Zhiyuan Liu, Juanzi Li, and Jian Tang. 2021. KEPLER: A Unified Model for Knowledge Embedding and Pre-trained Language Representation. _Transactions of the Association for Computational Linguistics_, 9:176-194 

Yanbin Wei, Qiushi Huang, Yu Zhang, and James Kwok. 2023. KICGPT: Large language model with knowledge in context for knowledge graph completion. In _Findings of the Association for Computational Linguistics: EMNLP 2023_, pages 8667-8683, Singapore. Association for Computational Linguistics.
* Wu et al. (2023) Hong Wu, Zhe Wang, Kewen Wang, Pouya Ghiasnezhad Omran, and Jiangmeng Li. 2023. Rule Learning over Knowledge Graphs: A Review. _Transactions on Graph Data and Knowledge_, 1(1):7:1-7:23.
* Yang et al. (2015) Bishan Yang, Wen-tau Yih, Xiaodong He, Jianfeng Gao, and Li Deng. 2015. Embedding entities and relations for learning and inference in knowledge bases. In _3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings_.
* Yao et al. (2019) Liang Yao, Chengsheng Mao, and Yuan Luo. 2019. Kgbert: Bert for knowledge graph completion.
* Zhang and Yao (2022) Yongqi Zhang and Quamming Yao. 2022. Knowledge graph reasoning with relational digraph. In _Proceedings of the ACM Web Conference 2022_, WWW '22, page 912-924, New York, NY, USA. Association for Computing Machinery.
* Zhao et al. (2021) Yunxiang Zhao, Jianzhong Qi, Qingwei Liu, and Rui Zhang. 2021. Wgcn: Graph convolutional networks with weighted structural features. In _Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval_, SIGIR '21, page 624-633, New York, NY, USA. Association for Computing Machinery.
* Zheng et al. (2024) Yaowei Zheng, Richong Zhang, Junhao Zhang, Yanhan Ye, Zheyan Luo, and Yongqiang Ma. 2024. Lamafactory: Unified efficient fine-tuning of 100+ language models. _arXiv preprint arXiv:2403.13372_.
* Zhou et al. (2020) Jie Zhou, Ganqu Cui, Shengding Hu, Zhengyan Zhang, Cheng Yang, Zhiyuan Liu, Lifeng Wang, Changcheng Li, and Maosong Sun. 2020. Graph neural networks: A review of methods and applications. _AI Open_, 1:57-81.
* Zhu et al. (2023a) Yuqi Zhu, Xiaohan Wang, Jing Chen, Shuofei Qiao, Yixin Ou, Yunzhi Yao, Shumin Deng, Huajun Chen, and Ningyu Zhang. 2023a. LLMs for knowledge graph construction and reasoning: Recent capabilities and future opportunities. _arXiv preprint arXiv:2305.13168_.
* Zhu et al. (2023b) Zhaocheng Zhu, Xinyu Yuan, Michael Galkin, Louis-Pascal Khonneux, Ming Zhang, Maxime Gazeau, and Jian Tang. 2023b. AST net: A scalable path-based reasoning approach for knowledge graphs. In _Advances in Neural Information Processing Systems_, volume 36, pages 59323-59336. Curran Associates, Inc.
* Zhu et al. (2021) Zhaocheng Zhu, Zuboai Zhang, Louis-Pascal Khonneux, and Jian Tang. 2021. Neural bellman-ford networks: A general graph neural network framework for link prediction. In _Advances in Neural Information Processing Systems_, volume 34, pages 29476-29490. Curran Associates, Inc.

 

Appendix

### Prompt templates of Reasoning stage

Table 6 shows the prompt templates of the **Reasoning** stage and provides an example.

```
##In-contextDemonstrations:User:WashingtonCounty:countyinPennsylvania,U.S.Thequestionistopredicttheheadentity[MASK]fromthegiven([MASK],locationadjoining_relationship adjoins,WashingtonCounty)bycompletingentence'WashingtonCountyistheadjoinsofwhatlocation?Theansweris'.Assistant:TheanswerivestwistemorelandCounty,sothe[MASK]isWestmorelandCounty.WestmorelandCountycountrycountryinPennsylvania,UnitedStatesUser:RocklandCountysuburban countyinNew York,UnitedStatesThequestionistopredicttheheadentity[MASK]fromthegiven([MASK],locationadjointagrelationship adjoins,RocklandCounty)bycompletingestence'RocklandCountyistheadjoinsoftwhatlocation?Theansweris'.Assistant:TheanswerisBergenCounty,sothe[MASK]isBergenCounty.BergenCounty:countyinNew Jersey,UnitedStates
\begin{tabular}{}
```##Querytriple:([MASK],/location/adjoinning_relationship/adjoins,Champaign) ```
##Reasoningprompts:User:Thequestionistopredictheheadentity[MASK]fromthegiven([MASK],locationadjoing_relationship adjoins,Champaign) bycompletingsentence'Champaignistheadjoinsofwhatlocation?Theansweris'.Outputallsomepossibleanswbasedonyoourownknowledge,usingtheformat'[answer1,answer2,...,answerN]'andpleasestartyourresponsewith'Thepossibleanswers:'.Donotoutputanythingexceptthepossibleanswors.
```##Context-awareReasoning:User:Herearesomematerialsforyoureferto:Champaign:ChampaignisacityintchampaignCountyIllinois,UnitedStates.Thepopulationwas88,302atethe2020census.ItisthetenthmostpopulousmunicipalityinillinoisandthefourthmostpopulouscityinthestateoutsidetheChicagometropolitanarea.Itisapricipalcityofthemampaign-Urbana metropolitanarea,whichhad236,000 residentsin2020. Thequestionistopredictheheadentity[MASK]fromthegiven([MASK],locationadjoning_relationship adjoins,Champaign)bycompletingstencence'Champaignistheadjoinsofwhatlocation?Theansweris'.Outputallthepossibleanswersyoanfindinthematerialsusingtheformat'[answer1,answer2,...,answerN]'andpleasestartyourresponsewith'Thepossibleanswers:'.Donotoutputanythingexceptthepossibleanswers.Ifyoucannotfindanyanswer,pleaseoutputsomepossibleansweredonyoourownknowledge. ```
##Reasoningresults:Assistant:Thepossibleanswers:[Urbana,ChampaignCounty,IllinoisSiliconPrairie,ParklandCollege]
```

\end{table}
Table 6: PromptTemplate of context-aware reasoning.

 



### Prompt templates of Re-Ranking stage

Table 7 shows the prompt templates of the **Re-ranking** stage and gives an example. It is noteworthy that this case empirically shows the effectiveness of the **Reasoning** and **Re-ranking** processes. The ground truth answer "Urbana" does not rank in a leading position by the KGC model. However, the LLM provides plausible candidates including the ground truth answer "Urbana", by analyzing the context of the known entity "Champaign" in the incomplete triple during the **Reasoning** process. In addition, the fine-tuned LLM succeeds in selecting the correct answer from the candidate list based on relevant textual contexts and neighbor facts during the **Re-ranking** process.

\begin{table}
\begin{tabular}{l l} \hline \hline
**## Query triple**: (< ??? >, _/location/adjointing_relationship/adjoins_, Champaign) \\ \hline
**## Re-Ranking prompts:** & \\
**User**: Here is an incomplete triple with missing head entity !missing-entity¿: (!missing-entity¿, location adjoin_s. location adjoining_relationship adjoins, Champaign). \\ \end{tabular}
\end{table}
Table 7: Prompt Template of context-aware ranking.



Implementation Details

We conduct all our experiments on a Linux server equipped with \(2\) Intel Xeon Platinum 8358 processors and \(8\) A100-SXM4-40GB GPUs. + For TransE and CoLE, we adopt the candidate entities provided in Liu et al. (2024). For other base KGC methods, candidate entities are retrieved using their publicly available implementations. We tested our proposed method using \(4\) different LLMs, specifically Llama2-7B, Llama3-8B, Qwen2-1.5B, and Qwen2-7B. The time costs of SFT and re-ranking are reported in Table 9. It is noteworthy that the duration of the experiments is subject to specific configurations, which can be further reduced through parallel inference or the application of vLLM service. Additionally, we also evaluated OpenAI's gpt-3.5-turbo-0125 model in the reasoning stage. During the re-ranking stage, we employ the LLaMA-Factory framework Zheng et al. (2024) to fine-tune the mentioned LLMs. For LoRA adjustments, we set the rank to 16 and the alpha value to 32. We utilize the AdamW optimizer Loshchilov and Hutter (2017) to train our model, with an initial learning rate of \(1.0e{-}4\), a per-device batch size of \(2\) and a gradient accumulation step of \(4\) iterations. We adopted BF\(16\) precision to reduce the GPU memory usage. For each query triple, we retrieve \(3\) supporting triples for in-context demonstration and provide \(n=20\) candidate answers (during training, we provide \(1\) ground truth entity and \(19\) negative samples) to the LLM. Detailed hyperparameter settings are listed in Table 8.

Footnote †: Only 2 GPUs are used in our experiments.

## Appendix C

\begin{table}
\begin{tabular}{c c} \hline \hline
**Hyperparameters** & **Settings** \\ \hline SFT learning rate & 1e-4 \\ Per-device batch size & 2 \\ \# GPUs used & 2 \\ Gradient accumulation step & 4 \\ LoR rank & 16 \\ LoRA \(\alpha\) value & \(32\) \\ \hline Re-ranking scope \(n\) & \(20\) \\ \# Easy negative samples & \(9\) \\ \# Hard negative samples & \(10\) \\ \(\delta\) & \(50\) \\ \(p\) & \(10\) \\ \hline \hline \end{tabular}
\end{table}
Table 8: Hyperparameter settings.

\begin{table}
\begin{tabular}{l l l l l l} \hline \hline
**Dataset** & **Owen2-1.5B** & **Qwen2-7B** & **Llama2-7B** & **Llama3-8B** & **Test (Llama3-8B)** \\ \hline FB15k237 & 12h 12min & 36h 50min & 43h 42min & 38h 27min & 0.180s / query \\ WN18RR & 5h 54min & 22h 33min & 19h 30min & 18h 12min & 0.152s / query \\ \hline \hline \end{tabular}
\end{table}
Table 9: Time costs of supervised fine-tuning and re-ranking.