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

Vibhor Agrawal, Fay Wang, Rishi Puri

Keywords: Graph Neural Networks, Retrieval-Augmented Generation, Query- Aware Attention, Information Retrieval

TL;DR: We propose a GNN-based RAG system that builds knowledge graphs from audio transcripts and uses query-aware attention to outperform dense retrieval on complex multi-hop questions.

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.

Submission Number: 13