Query-Aware Flow Diffusion for Graph-Based RAG with Retrieval Guarantees

TL;DR

This paper introduces QAFD-RAG, a query-aware, training-free graph traversal framework that adaptively guides retrieval based on query semantics, providing theoretical guarantees and improved performance in graph-based RAG tasks.

Contribution

The paper proposes a novel query-aware flow diffusion method for graph traversal in RAG systems, offering the first statistical guarantees and dynamic, semantics-driven exploration.

Findings

QAFD-RAG recovers relevant subgraphs with high probability.

The method converges exponentially fast, scaling with subgraph size.

Experiments show improved results over state-of-the-art methods.

Abstract

Graph-based Retrieval-Augmented Generation (RAG) systems leverage interconnected knowledge structures to capture complex relationships that flat retrieval struggles with, enabling multi-hop reasoning. Yet most existing graph-based methods suffer from (i) heuristic designs lacking theoretical guarantees for subgraph quality or relevance and/or (ii) the use of static exploration strategies that ignore the query's holistic meaning, retrieving neighborhoods or communities regardless of intent. We propose Query-Aware Flow Diffusion RAG (QAFD-RAG), a training-free framework that dynamically adapts graph traversal to each query's holistic semantics. The central innovation is query-aware traversal: during graph exploration, edges are dynamically weighted by how well their endpoints align with the query's embedding, guiding flow along semantically relevant paths while avoiding structurally connected but irrelevant regions. These query-specific reasoning subgraphs enable the first statistical guarantees for query-aware graph retrieval, showing that QAFD-RAG recovers relevant subgraphs with high probability under mild signal-to-noise conditions. The algorithm converges exponentially fast, with complexity scaling with the retrieved subgraph size rather than the full graph. Experiments on question answering and text-to-SQL tasks demonstrate consistent improvements over state-of-the-art graph-based RAG methods.