LLMs and GNNs are a killer combo for RAG applications

GNN-RAG brings together the knowledge graph–processing abilities of graph neural networks and the language abilities of LLMs to unlock new applications.

Knowledge graphs (KG) are rich sources of information for knowledge-intensive tasks such as question answering. However, retrieving relevant information from KGs is challenging, especially when dealing with large-scale knowledge bases that contain millions of facts.

To solve this problem, researchers at the University of Minnesota have developed GNN-RAG, a novel method that combines the language understanding abilities of large language models (LLM) with the reasoning abilities of graph neural networks (GNN) for knowledge graph question answering (KGQA).

The technique can have important uses for enterprises that want to integrate their knowledge graphs into their generative AI applications.

Combining GNNs and LLMs

KGs represent human-crafted factual knowledge in the form of triplets (head, relation, tail), which collectively form a graph and are stored in a format that can be easily updated. As KGs capture complex interactions between the stored entities, they can help to answer questions that require multiple hops across different entities. 

On the other hand, LLMs are very good at combining contextual information with their inside knowledge and presenting to parse and answer natural language queries. Therefore knowledge graphs can be very good sources of information for retrieval-augmented generation (RAG) applications, where an LLM is provided with additional context to generate answers. But matching the right KG elements to a user prompt is tricky, which makes the retrieval mechanism difficult.

"Knowledge Graphs (KGs) encode knowledge in a structured form that can be easily updated and that can be used to reduce LLM hallucination via RAG," Costas Mavromatis, lead author of the GNN-RAG paper, told TechTalks. "However, KGs include complex graph information and thus, graph-based algorithms are needed to process this information. This led us to repurpose Graph Neural Networks (GNNs), which are powerful graph learners, for retrieval over KGs."

GNNs are deep learning architectures that can learn and reason over graph-structured data. GNNs have proven to be very effective at complex reasoning tasks over knowledge graphs. However, contrary to LLMs, GNNs don't have natural language understanding capabilities. 

GNN-RAG

"Both GNNs and LLMs have been applied for KGQA [knowledge graph question answering]," Mavromatis said. "However, the missing element is how we can combine these two model families in the best way possible and leverage their unique strengths. Our GNN-RAG recipe utilizes GNNs for retrieval due to their graph reasoning abilities, while leveraging the natural language abilities of LLMs for ultimate KGQA."

GNN-RAG uses a two-step process to answer questions. First, a GNN reasons over a dense KG subgraph to retrieve answer candidates for a given question. The nodes with the highest probability scores are chosen along with the shortest paths connecting the question entities with the answers. 

Next, the selected paths that connect the entities and the answer candidates are extracted to represent KG reasoning paths. The extracted paths are then verbalized and given as input to an LLM such as ChatGPT or Llama, which uses it for reasoning and answering the question.

"In our GNN-RAG framework, the GNN acts as a dense subgraph reasoner to extract useful graph information, while the LLM leverages its natural language processing ability for ultimate KGQA," the researchers write.

The researchers further augmented GNN-RAG with an LLM-based retriever to boost its performance on simple questions that don't require multi-hop reasoning. 

Evaluating GNN-RAG

The researchers evaluated their method on two widely used KGQA benchmarks: WebQuestionsSP (WebQSP) and Complex WebQuestions 1.1 (CWQ). They compared GNN-RAG with state-of-the-art GNN and LLM methods for KGQA as well as zero-shot and few-shot LLM prompts.

"GNN-RAG is the method that performs overall the best, achieving state-of-the-art results on the two KGQA benchmarks in almost all metrics," the researchers write. "The results show that equipping LLMs with GNN-based retrieval boosts their reasoning ability significantly."

GNN-RAG outperformed competing RAG-based systems for KGQA by up to 15.5 percentage points on complex questions. It also outperformed or matched the performance of a system that uses GPT-4 while using a much smaller LLM and fewer API calls. The researchers estimate that their GPT-4–based baseline has an overall cost above $800 to evaluate the datasets, while GNN-RAG can be deployed on a single 24GB GPU with a fine-tuned open LLM such as Llama-2-Chat-7B.

GNN-RAG also outperformed other state-of-the-art methods on multi-entity and multi-hop questions, showing that it is an effective retrieval method when deep graph search is important.

"GNN-RAG can help in applications where KGs encode useful or critical information for the in-domain task," Mavromatis said. "For example, GNN-RAG can leverage medical KGs for QA in health-care or customer-based KGs for personalized recommendations."

The researchers plan to continue working on GNN-RAG and expand it in several directions, including detecting missing information in knowledge graphs and combining KGs and GNN-RAG with other modalities, such as documents and images.