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

Zhuoping Zhou, Davoud Ataee Tarzanagh, Sima Didari, Wenjun Hu, Baruch Gutow, Oxana Verkholyak, Masoud Faraki, Heng Hao, Hankyu Moon, Seungjai Min

Keywords: Graph-Based RAG, Training-Free Retrieval, Multi-Hop Reasoning, Query-Aware Graph Traversal, Subgraph Recovery Guarantees

TL;DR: QAFD-RAG is a training-free graph-based RAG that dynamically adapts traversal to query semantics via flow diffusion, yielding principled subgraph recovery guarantees and consistent gains on reasoning tasks.

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 \textit{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 \emph{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.

Primary Area: learning on graphs and other geometries & topologies

Submission Number: 23598