Query-Aware Graph Neural Networks for Enhanced Retrieval-Augmented Generation

TL;DR

This paper introduces a query-aware graph neural network architecture for retrieval-augmented generation that constructs knowledge graphs to improve retrieval accuracy on complex, multi-hop questions, outperforming traditional methods.

Contribution

The paper proposes a novel GNN architecture with query-aware attention and learned scoring for enhanced retrieval in RAG systems, especially for multi-hop reasoning tasks.

Findings

Significant improvement over standard dense retrievers on complex QA tasks

Effective use of knowledge graphs to capture relationships between text chunks

Scalable implementation with PyTorch Geometric

Abstract

We present a novel graph neural network (GNN) architecture for retrieval-augmented generation (RAG) that leverages query-aware attention mechanisms and learned scoring heads to improve retrieval accuracy on complex, multi-hop questions. Unlike traditional dense retrieval methods that treat documents as independent entities, our approach constructs per-episode knowledge graphs that capture both sequential and semantic relationships between text chunks. We introduce an Enhanced Graph Attention Network with query-guided pooling that dynamically focuses on relevant parts of the graph based on user queries. Experimental results demonstrate that our approach significantly outperforms standard dense retrievers on complex question answering tasks, particularly for questions requiring multi-document reasoning. Our implementation leverages PyTorch Geometric for efficient processing of graph-structured data, enabling scalable deployment in production retrieval systems