Query-Aware Graph Neural Networks for Enhanced Retrieval-Augmented Generation
This work addresses the challenge of enhancing retrieval for complex question answering in AI systems, representing an incremental advancement by building on existing graph neural network and retrieval-augmented generation methods.
The paper tackles the problem of improving retrieval accuracy for complex, multi-hop questions in retrieval-augmented generation by proposing a query-aware graph neural network architecture that constructs knowledge graphs and uses query-guided attention. The result is a significant performance improvement over standard dense retrievers, particularly in multi-document reasoning tasks.
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