GNN-RAG: Graph Neural Retrieval for Large Language Model Reasoning

Published on March 18, 2024 | AI Research

Ever asked your AI a tricky question like, "Which language do Jamaican people speak?" and gotten a half-baked answer? Large language models (LLMs) are great at chatting, but they can fumble when it comes to reasoning over structured data like Knowledge Graphs (KGs). Introducing GNN-RAG, an approach from a paper that teams up Graph Neural Networks (GNNs) with LLMs to deliver razor-sharp answers grounded in facts. It's like giving your AI a detective's brain and a poet's tongue. Let's understand it in a simple language.

Knowledge Graphs: The Fact Web

Knowledge Graphs (KGs) are like a giant web of human knowledge, built from triplets of the form (head, relation, tail). Think (Jamaica, official_language, English) or (Stephen Hawking, studied_at, Oxford). These triplets connect to form a graph, where nodes (like "Jamaica" or "English") are linked by relationships (like "official_language"). Knowledge Graph Question Answering (KGQA) is the task of digging through this graph to find a set of entities {a} that answer a natural-language question q, like "Which language do Jamaican people speak?" The catch? KGs can have millions of nodes and facts, so we need a smart way to navigate them.

The Problem: AI's Got Gaps

LLMs, like ChatGPT or LLaMA, are wizards at understanding natural language but struggle with the structured, interconnected nature of KGs. They might guess or hallucinate answers, especially for complex questions requiring multiple steps (aka multi-hop). Meanwhile, GNNs are champs at reasoning over graphs, think node classification to spot answers but they're not great at spitting out fluent, human like responses. It's like one's a math genius and the other's a storyteller, and they need to team up. That's where GNN-RAG comes in, blending the two in a Retrieval-Augmented Generation (RAG) framework to understand KGQA.

How GNN-RAG Works: A Step-by-Step Breakdown

Let's walk through how GNN-RAG tackles a question like, "Where did the author of A Brief History of Time go to college?" It's a multi-hop question that needs both graph reasoning and natural-language finesse.

Step 1: Subgraph Retrieval

KGs are massive, so GNN-RAG starts by pulling a smaller, question-specific subgraph Gq using techniques like entity linking and neighbor extraction. For our question, this subgraph includes nodes like "Stephen Hawking," "A Brief History of Time," and "Oxford." The goal? Make sure all correct answers {a} are in Gq.

Step 2: GNN Reasoning

The GNN treats KGQA as a node classification problem, labeling nodes as answers or non-answers. It updates node representations using a message-passing mechanism, defined as:

Mathematical Formula:

**h**ᵥ⁽ˡ⁾ = ψ(**h**ᵥ⁽ˡ⁻¹⁾, Σᵥ'∈Nᵥ ω(q, r) · **m**ᵥᵥ'⁽ˡ⁾)

Here, **h**ᵥ⁽ˡ⁾ is the representation of node v at layer l, ψ combines representations across layers, and ω(q, r) measures how relevant a relation r (like "studied_at") is to the question q. Messages **m**ᵥᵥ'⁽ˡ⁾ come from neighboring nodes v'. The GNN uses a pretrained language model (like SBERT or LMSR) to encode the question and relations, ensuring the reasoning is question-aware.

After L layers (e.g., L=3 for deep reasoning), the GNN scores nodes using a softmax over final representations **h**ᵥ⁽ᴸ⁾, picking high-probability nodes as answer candidates (e.g., "Oxford," "Cambridge"). It also extracts the shortest paths from question entities to these answers, like:

A Brief History of Time → authored_by → Stephen Hawking → studied_at → Oxford
Stephen Hawking → studied_at → Cambridge

These paths capture the reasoning logic.

Step 3: Path Verbalization

The GNN's paths are turned into natural-language sentences, like:

"Stephen Hawking is the author of A Brief History of Time. He studied at the University of Oxford."
"He also studied at the University of Cambridge."

This is done using a prompt template, verbalizing paths as "{question entity} → {relation} → {entity} → … → {answer entity}."

Step 4: LLM Reasoning with RAG

The verbalized paths are fed to a fine-tuned LLM (e.g., LLaMA2-Chat-7B) with a prompt like:

"Based on the reasoning paths, please answer the given question.
Reasoning Paths: {paths}
Question: {question}"

The LLM reads the context and outputs a fluent answer: "Stephen Hawking attended both Oxford and Cambridge." To make it robust, the LLM is fine-tuned on question-answer pairs, optimizing for accurate responses given verbalized paths.

Step 5: Retrieval Augmentation (RA)

To boost performance, GNN-RAG+RA combines the GNN's paths with those from an LLM-based retriever (like RoG). RoG uses beam-search decoding to generate diverse relation paths (e.g., <author, studied_at>, <writer, education>), which are mapped to the KG to fetch intermediate entities. The union of GNN and LLM paths increases answer recall, especially for single-hop and multi-hop questions. Alternatively, GNN-RAG+Ensemble combines paths from two GNNs (one using SBERT, another using LMSR) for a cheaper, equally effective boost.

Why GNN-RAG Is a Big Deal

GNN-RAG is like a dream team: GNNs handle the heavy lifting of multi-hop reasoning over complex KGs, while LLMs polish the answers into something you'd hear from a friend. The GNN's message-passing shines for multi-hop questions, with deep GNNs (L=3) achieving 88.5% answer coverage on WebQSP's 2-hop questions, compared to 82.1% for RoG and 79.8% for shallow GNNs (L=1). It's also efficient, using fewer input tokens (357 vs. 435 for RoG).

The paper's results are fire: GNN-RAG matches or beats GPT-4 on WebQSP and CWQ benchmarks using just a 7B LLM. It crushes multi-hop and multi-entity questions, outperforming competitors by 8.9–15.5% in F1 score. For simple 1-hop questions, the RA technique ensures LLM-based retrieval fills any gaps, making GNN-RAG versatile.

The Tech Behind the Magic

Here's the breakdown of GNN-RAG's key components:

GNN Reasoning: Uses message-passing to score nodes and extract shortest paths, leveraging pretrained LMs (SBERT or LMSR) for question-relation matching.
Path Verbalization: Converts graph paths into natural-language context for the LLM.
LLM Fine-Tuning: A LLaMA2-Chat-7B model is tuned with RAG prompt templates to generate accurate, fluent answers.
Retrieval Augmentation: Combines GNN and LLM retrievers (or multiple GNNs) for diverse, high-recall paths.
Efficiency: Deep GNNs retrieve complex multi-hop info with fewer tokens than LLM-based methods.

The paper also includes a theorem showing that the GNN's output hinges on the question-relation matching function ω(q, r), highlighting the importance of the chosen LM (SBERT vs. LMSR).

Limitations and Trade-Offs

GNNs shine for multi-hop questions but can lag on simple 1-hop questions where precise question-relation matching matters more. Here, LLM-based retrievers like RoG edge out, which is why RA is clutch. Also, LLM-based retrieval can be costly due to beam-search decoding, but GNN-RAG+Ensemble offers a leaner alternative by combining GNN outputs.

Implementation

If you want to play with GNN-RAG, the code and results are at https://github.com/cmavro/GNN-RAG. Spin it up, tweak the GNN layers, or mix in your own LLM - it's a playground for AI nerds fr.

The Big Picture

GNN-RAG is a step toward AI that thinks and talks like us, blending structured reasoning with natural language. It's not just about nailing KGQA, it's about building systems that reason logically and communicate clearly. So, next time you're curious about a fact, imagine GNN-RAG weaving through a web of knowledge to drop a crisp, accurate answer.

Paper: Read the full paper here

Pratham Grover