QAFD-RAG

Paper Code Code + Knowledge Graphs

Query-Aware Flow Diffusion for Graph-Based
Retrieval-Augmented Generation

Graph-based RAG systems leverage interconnected knowledge structures to capture complex relationships, enabling multi-hop reasoning. Yet most existing methods suffer from (i) heuristic designs lacking theoretical guarantees and (ii) static exploration strategies that ignore the query’s holistic meaning. We propose QAFD-RAG, a training-free framework that dynamically adapts graph traversal to each query’s semantics. The central innovation is query-aware traversal: edges are dynamically weighted by how well their endpoints align with the query embedding, guiding flow along semantically relevant paths while suppressing irrelevant regions. This enables the first statistical guarantees for query-aware graph retrieval, with exponential convergence and complexity scaling with the retrieved subgraph rather than the full graph.

Key Contributions

Query-Aware Framework

First principled flow diffusion method for graph-based RAG that incorporates query semantics via alignment-based edge weighting. Adapts flow probabilities online, guiding traversal toward semantically relevant regions.

Provable Guarantees

Exponential convergence to a query-dependent stationary distribution. Recovery guarantees ensuring relevant subgraphs are retrieved with high probability under mild signal-to-noise conditions.

State-of-the-Art Results

Leads in 96% of UltraDomain comparisons. Best F1/EM on HotpotQA (73.4/58.1) and MuSiQue (48.0/33.5). Outperforms GraphRAG, LightRAG, and RAPTOR across all benchmarks.

Approach Comparison

Query: “Introduce Steve Jobs’s products in Apple.”

GraphRAG
GraphRAG
Community Detection
LightRAG
LightRAG
One-Hop Entity-Centric
QAFD-RAG
QAFD-RAG (Ours)
Query-Aware Flow Diffusion

GraphRAG retrieves entire communities, mixing relevant and irrelevant nodes. LightRAG extracts 1-hop neighborhoods without semantic alignment. QAFD-RAG reweights edges by query meaning, suppressing irrelevant clusters (Amazon River, Apple fruit) and reinforcing reasoning paths (Apple → Mac → macOS). Edge thickness reflects weight; node color indicates importance.

How It Works

A two-stage pipeline from raw documents to grounded answers

Stage 1 — Knowledge Graph Construction

  1. Document Chunking — Split into context-preserving chunks.
  2. Entity & Relationship Extraction — LLM builds a structured KG with entities as nodes and relationships as edges.
  3. Embedding — All nodes embedded into vector space for similarity matching.

Stage 2 — Query-Aware Retrieval

  1. Seed Selection — Query matched to KG entities via embedding similarity; LLM reranker filters top seeds.
  2. Flow Diffusion — Mass injected at seeds and propagated via push-relabel with query-aware edge weights.
  3. Response Generation — Top-ranked subgraph assembled into context for LLM answer generation.

Benchmarks

Evaluated on UltraDomain QA (11 domains), multi-hop QA, text-to-SQL, and summarization. Full results in the paper.

UltraDomain QA
11 domains · 5 quality metrics
Leads in 96% of comparisons (48/50)
Multi-hop QA
MuSiQue · HotpotQA · 2Wiki
Best F1 on HotpotQA (73.4) & MuSiQue (48.0)
Text-to-SQL
Spider2-lite (SQLite · Snowflake)
26.7% SQLite · 23.7% Snowflake
Summarization
SQuALITY
BLEU · ROUGE · METEOR

Quick Start

Get running in minutes with pre-built knowledge graphs

# Option 1: Clone from HuggingFace (code + pre-built KGs)
git clone https://huggingface.co/tarzanagh/QAFD-RAG
cd QAFD-RAG

# Option 2: Clone from GitHub (code only) + download KGs
git clone https://github.com/Tarzanagh/QAFD-RAG.git
cd QAFD-RAG
huggingface-cli download tarzanagh/QAFD-RAG --include "kg/multihop/*" --local-dir .

# Install and run
pip install -r requirements.txt
export OPENAI_API_KEY="sk-..."
python benchmarks/run.py --task multihop --dataset musique --questions 10

Citation

@inproceedings{zhou2026qafd,
  title={Query-Aware Flow Diffusion for Graph-Based RAG with Retrieval Guarantees},
  author={Zhou, Zhuoping and Ataee Tarzanagh, Davoud and Didari, Sima and Hu, Wenjun
          and Gutow, Baruch and Verkholyak, Oxana and Faraki, Masoud and Hao, Heng
          and Moon, Hankyu and Min, Seungjai},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2026}
}