Think-on-Graph 2.0: Deep and Interpretable Large Language Model Reasoning with Knowledge Graph-guided Retrieval

RAG aims to offer real-time knowledge updates and effective utilization of external knowledge sources with high interpretability. An important factor is the granularity of the retrieved data. Coarse-grained retrieval units theoretically can provide more relevant information for the problem, but they may also contain redundant content, which could distract the retriever and language models in downstream tasks. On the other hand, fine-grained retrieval unit granularity increases the burden of retrieval and does not guarantee semantic integrity and meeting the required knowledge semantic integrity and meeting the required knowledge(Gao et al., 2024). This low information density and low utility are due to the inherent limitations of semantic retrieval.

KGs offer dynamic, explicit, and structured representations of knowledge. This structured knowledge representation is particularly beneficial for LLMs because it introduces a level of interpretability and precision in the knowledge that LLMs can access. Early approaches (Sun et al. (2020), Peters et al. (2019), Huang et al. (2024), Liu et al. (2020)) focused on embedding knowledge from KGs directly into the neural networks of LLMs. This embedding could occur either during the pretraining phase or the fine-tuning process. The intent was to infuse the models with rich, structured knowledge right from the foundational stages of model training. Despite the initial promise, integrating KGs directly into LLMs introduced challenges. As noted by Hu et al. (2023), this integration leads to a reduction in the natural explainability of the models. Additionally, it makes updating the knowledge base more complex and less efficient, as any change in the KG requires retraining or significant adjustments to the model. More recent studies(Jiang et al., 2024) have shifted towards using KGs to augment LLMs externally rather than embedding the knowledge directly into the models. Those approaches involve translating relevant structured knowledge from KGs into textual prompts that are then fed to LLMs. The process typically follows a fixed pipeline where extra information from KGs is retrieved to enhance the LLM prompts. The integration of Knowledge Graphs (KGs) with Large Language Models (LLMs) offers numerous advantages but also meets several distinct challenges and limitations, such as incompleteness and ambiguity as discussed in the last section.

In this work, we aim to integrate both KG and unstructured data with LLM, leveraging the strengths of both to mitigate their respective weaknesses.

Selecting appropriate starting points for specific queries can facilitate much more streamlined reasoning. For example: ”Among the founders of Tencent company, who has been a member of the National People’s Congress?”. In this case, a broad or poorly chosen point such as ”Member of the National People’s Congress” could lead to pitfalls of sifting through large amounts of irrelevant data and cause time-consuming and less focused exploration. An effective starting point would be to focus on the entity ’Founders of Tencent’. This principle is essential in reasoning tasks, especially in open-domain question answering, where the question is highly varied. Therefore, given a question $q$, Tog_2.0 first performs Named Entity Recognition (NER) and Topic Prune (TP), which prompts the LLM to evaluate $q$ and appearing entities, selecting topic entities $E_{topic}^{0}=\{e_{1},e_{2},\ldots,e_{N}\}$ that are appropriate to serve as the starting point for the question.

In complex reasoning, the implicit correlation between the question and effective intermediate clue sentences often goes unrecognized by both sparse retrieval models and dense pre-trained retrieval models. To solve this limitation, we prompt the LLM to formulate clue queries $q_{j}^{0}$ based on the current context for every topic entity $e_{j}$, which orients the next-step direction of exploring the relations and contexts. Using the question about Tencent mentioned above as the example again, based on the entity ”National People’s Congress?”, the LLM may generate a clue-query that suggests gathering information about their political roles or affiliations.

Next, we will introduce how Tog_2.0 iteratively utilizes structured and unstructured knowledge for reasoning. Formally, in the ${(i+1)}_{th}$ iteration, given the original question $q$, the clue queries from the ${i}_{th}$ iteration $Q_{c}^{i}=\{q_{1}^{i},q_{2}^{i},\ldots,q_{N}^{i}\}$, the topic entities $E_{topic}^{i}=\{e_{1},e_{2},\ldots,e_{N}\}$ and their preceding triple paths $\mathbf{P}^{i}=\{P_{1}^{i},P_{2}^{i},\ldots,P_{N}^{i}\}$, $P_{j}^{i}=\{p_{j}^{0},p_{j}^{1},\ldots,p_{j}^{i}\}$, each iteration includes three steps: relation prune (RP), entity prune (EP), and examine and reasoning (ER). Note that $i=0$ indicates the initialization phase and the $P^{0}$ is empty.

Relation Prune (RP): Based on $q$ and $Q_{c}^{i}$, we prompt the LLM to select the relations that are most likely to find entities containing helpful context information for solving $q$ and that match the description of $Q_{c}^{i}$. Unlike selecting relations for a single topic entity at a time, we provide GPT-3.5 with all topic entities within a single prompt. This approach not only reduces the number of API calls, thereby accelerating inference time, but also enables the LLM to simultaneously consider the interconnections between multiple reasoning paths, allowing it to make selections from a more global perspective. The selected relations for topic entity $e_{j}$ are denoted as $R_{j}=\{r_{j1},r_{j2},\ldots,r_{jW}\}$, where $W$ denotes the hyper-parameter width.

Entity Prune (EP): Given a topic entity $e_{j}$ and one of the selected relation $r_{j}k$, Tog_2.0 will identify all interconnected candidate entity nodes $\{e_{jkl}\}$ within the Wiki Knowledge Graph (KG) and get their associated Wikipedia page documents $D_{jkl}$ through locally deployed service. The document context of each candidate entity is initially segmented into appropriately sized chunks$\{t_{jklm}\}$. Subsequently, a two-stage search $F_{retr}$ is employed, utilizing pre-trained language models for all candidate entities’ chunks. Formally, $s_{jklm}=F_{retr}([q,q_{j}^{i},p_{jkl}^{i}],t_{jklm})$ denotes the relevance score of the $m_{th}$ paragraph of the $l_{th}$ candidate, where $p_{jkl}^{i}$ is the triples from which the current candidate entity is derived. Then, the ranking score of a candidate entity $e_{jkl}$ is calculated as the exponentially decayed weighted sum of scores of its chunks that rank in top-$K$, and the weight for the $i_{th}$ ranked chunk is calculated as $w=e^{-\alpha\cdot i}$, where $K$ and $\alpha$ are hyper-parameters. Finally, top-W candidate entities are selected as the new topic entities $E_{topic}^{i+1}$ for the next iteration, meanwhile the corresponding preceding triple paths $\mathbf{P}^{i+1}$ will be updated.

Examine and reasoning (ER): Following RP and EP, we give LLM carefully aggregated references, including $q$, $Q_{c}^{i}$, $\mathbf{P}^{i+1}$ and the top $L$ ($L\leq K$) chunks. Then the LLM is prompted to examine the logical coherence and the completeness of factual evidence. If the LLM believes it can answer the question, the iteration ends. If not, based on the question and the collected contextual clues, a new clue-query needs to be generated for the next round.

We evaluated our method on two multi-hop KBQA datasets WebQSP(Yih et al., 2016) and QALD10-en(Usbeck et al., 2023), a multi-hop complex document QA dataset HotpotQA(Yang et al., 2018), and a fact verification dataset FEVER(Thorne et al., 2018). The evaluation metric for FEVER is Accuracy, while the metric for WebQSP & HotpotQA & QALD-10-en is Exact Match (EM).

We compare ToG_2.0 with both widely used baselines and state-of-the-art methods to provide a more comprehensive overview: 1) Standard prompting (Vanilla Prompt) directly answers the question. 2) Chain-of-thought (CoT)(Wei et al., 2022b) generates several intermediate rationales before the final answer to improve the complex reasoning capability of LLMs. 3) Chain-of-Knowledge (CoK)(Li et al., 2024) a heterogeneous source augmented large language model framework. 4) Think-on-Graph (ToG)(Sun et al., 2024) a KG method that searches useful triples for reasoning.

In this study, considering the experimental costs and for ease of comparison with other baselines, we conducted experiments on two LLMs: GPT-3.5-turbo and Llama-2-13b-chat. We used the OpenAI API to access GPT-3.5-turbo, while Llama-2-13B-chat was deployed on 8 A100-40G GPUs without quantization. Consistent with the ToG settings, we set the temperature parameter to 0.4 during exploration and 0 during reasoning. The maximum token length for generation was capped at 256. For context retrieval, we utilized the pre-trained BGE-embedding model without any fine-tuning. We choose Wikidata as the knowledge source for all experiments. During the TP, RC, relation pruning, and reasoning stages, we employed a 2-shot demonstration for all prompts.

As shown in Table1, we analyze the performance of our proposed method, Tog_2.0, in comparison with state-of-the-art baselines, including the Vanilla RAG strategy, the Chain-of-Thought (CoT), and the current SOTA baseline (CoK). The evaluation metrics include Exact Match (EM) for WebQSP, HotpotQA, and QALD-10-en, and Accuracy for FEVER.

We note that Tog_2.0 outperforms other baselines on WebQSP, HotpotQA and QALD-10-en. Notably, on HotpotQA, it significantly surpasses the SOTA baseline CoK by 5.51%. Compared to the original ToG, Tog_2.0 achieved a substantial improvement on HotpotQA (14.6%) and also demonstrated notable enhancements on other datasets (4.93% on WebQSP, 3.85% on QALD-10-en and 5.84% on FEVER). This demonstrates the advantages of our ”KG+context” framework in addressing complex problems. Although the performance on the fact-checking dataset FEVER is slightly inferior to CoK, this may be due to CoK utilizing more knowledge sources and an additional LLM self-verification mechanism. To save computational costs and reduce inference latency, we ultimately decided not to use a self-verification mechanism, which could be applied based on application requirements in the future.

To evaluate the contribution of each component in Tog_2.0, we conducted comprehensive ablation experiments across all datasets.

Compared to the performance on the other three datasets, on WebQSP, the effectiveness of Topic Prune (TP) is more pronounced, possibly due to the higher relative proportion of general entities in WebQSP questions, which tends to introduce more unnecessary noise.

Although Relation Prune (RC) may slightly decrease the performance due to the increased difficulty for the LLM in understanding multiple tasks within a single prompt, the benefit is a significant reduction in the number of inferences and latency. Assuming a reasoning depth of $N$ and a width of $W$, the complexity can theoretically be reduced from $O(W^{N})$ to $O(N)$.

Additionally, clue-query also brought relatively consistent improvements across each dataset, indicating that adaptive query optimization can help the LLM better understand the tasks. We also tested the vanilla RAG process and the basic Tog_2.0 process on Llama-2-13B. It can be seen that on a less capable LLM, Tog_2.0 can bring a greater performance improvement. This suggests that Tog_2.0 might be more adaptable. Weaker LLMs often encounter bottlenecks when handling complex tasks, while Tog_2.0 utilizes knowledge graphs as clues to optimize the reasoning path and reduce task complexity. It then uses entity context to further guide the model to focus on relevant information, thereby improving task understanding and response accuracy. In contrast, GPT-3.5, due to its higher inherent capabilities, may not exhibit as significant a performance improvement because it is already close to its performance ceiling.