KG-RAG

Abstract

Ensuring factual accuracy while maintaining the creative capabilities of Large Language Model Agents (LMAs) poses significant challenges in the development of intelligent agent systems. LMAs face prevalent issues such as information hallucinations, catastrophic forgetting, and limitations in processing long contexts when dealing with knowledge-intensive tasks. This paper introduces a KG-RAG (Knowledge Graph-Retrieval Augmented Generation) pipeline, a novel framework designed to enhance the knowledge capabilities of LMAs by integrating structured Knowledge Graphs (KGs) with the functionalities of LLMs, thereby significantly reducing the reliance on the latent knowledge of LLMs. The KG-RAG pipeline constructs a KG from unstructured text and then performs information retrieval over the newly created graph to perform KGQA (Knowledge Graph Question Answering). The retrieval methodology leverages a novel algorithm called Chain of Explorations (CoE) which benefits from LLMs reasoning to explore nodes and relationships within the KG sequentially. Preliminary experiments on the ComplexWebQuestions dataset demonstrate notable improvements in the reduction of hallucinated content and suggest a promising path toward developing intelligent systems adept at handling knowledge-intensive tasks.

The Future of Intelligent Agents

The integration of KGs and LLMs through frameworks like KG-RAG suggests a promising path toward developing intelligent systems capable of tackling knowledge-intensive tasks with high reliability. This approach mitigates some of the current limitations in LLM-based Agents and encourages further research and practical applications of this method. There is a significant opportunity to improve the current capabilities of KG Construction by developing curated high-quality datasets specifically tailored for this field.

KG-RAG holds promise as a potential solution to equip LLMs with external, updatable knowledge that is both reliable and efficient. Furthermore, with the continuous advancements in AI hardware accelerators, the reduction in response times (with KG-RAG and reasoning processes happening behind) will make the interactions with these agents feel more natural to humans.

KG-RAG Framework

KG-RAG consists of three main stages:

KG Construction: Converts unstructured text into a structured KG. This step is crucial to maintain the quality of information and can affect later stages if it is not done correctly.
Retrieval: Carried out through a Chain of Explorations (CoE), a novel retrieval algorithm that employs LLM reasoning to explore nodes and relationships within the KG. It ensures that the retrieval process is both relevant and accurate leveraging the benefits of dense vector representations to locate relevant nodes and relationships within large KGs.
Response Generation: Generates coherent and contextually appropriate responses.

Diagram illustrating the KG-RAG framework.

Chain-of-Explorations (CoE)

The retrieval stage constitutes a critical component of the KGQA process, where the goal is to systematically extract pertinent information from a KG in response to a user-specified query. To this end, we perform a Chain of Explorations (CoE) over the KG, a novel approach designed to facilitate the exploration of a KG through a sequential traversal of its nodes and relationships to find relevant information.

CoE, illustrated in the diagram, strategically traverses the KG by navigating through nodes and edges that are directly relevant to the query. It performs a step-by-step exploration to ensure a thorough and directed search, increasing the likelihood of retrieving accurate and relevant information.

The algorithm consists of three main components: planning, KG lookups, and evaluation. The planning phase involves devising a sequence of steps that guide the exploration through the KG. KG lookups are executed using a combination of vector similarity search and cipher queries to identify relevant nodes or relationships. The evaluation phase checks the current state of the traversal against the initial query, deciding whether to continue exploring, refine the exploration strategy, or synthesize an answer based on the gathered context.

The initial step of CoE uses a few-shot learning prompt along with the user query q to guide the planner in formulating a strategic exploration plan throughout the KG. In the first step of the plan, CoE finds top-k starting nodes using vector similarity search over a specified keyword and selects relevant nodes to perform further exploration. The process then progresses into a cyclical lookup phase within the KG, outlined as follows:

Execution of Cypher queries to retrieve connected nodes and relationships in the KG.
Ranking of nodes or relationships by relevance using dense vector embeddings to measure relevance for the current step's task.
Utilizing an LLM to filter and select the most relevant nodes or relationships for continuing the exploration hops over the KG.

This is followed by an evaluation phase, where the LLM assesses the alignment of the current traversal with the initial plan. Based on this assessment, the algorithm decides whether to continue the exploration, adjust the exploratory steps, or synthesize a response based on the context found thus far. Suppose the evaluation phase identifies a conclusive answer within the traversed path. In that case, it compiles the data into a coherent chain of node triples and sends it to the next stage for generating an answer.

Diagram illustrating the components to perform KGQA over a KG using Chain of Explorations (CoE).

BibTeX

@misc{sanmartin2024kgrag,
        title={KG-RAG: Bridging the Gap Between Knowledge and Creativity}, 
        author={Diego Sanmartin},
        year={2024},
        eprint={2405.12035},
        archivePrefix={arXiv},
        primaryClass={cs.AI}
  }

KG-RAG: Bridging the Gap between Knowledge and Creativity

Diagram of the three core components of an AI agent: perception, brain, and action. The brain component integrates LLMs for dynamic reasoning and decision-making, alongside KGs for structured knowledge and memory storage.