QAFD-RAG uses query-aware flow diffusion to retrieve contextually relevant subgraphs from a knowledge graph. Unlike community-based (GraphRAG) or entity-centric (LightRAG) approaches, QAFD-RAG dynamically re-weights edges based on query relevance and propagates flow through the graph to discover multi-hop context with retrieval guarantees.
Query: "Introduce Steve Jobs's products in Apple."
GraphRAG
GraphRAG
LightRAG
LightRAG
QAFD-RAG
QAFD-RAG
QAFD-RAG reweights edges by query relevance, suppressing irrelevant neighborhoods (e.g., Amazon, Apple fruit).
Quick Start
bash
1# 1. Create conda environment and install dependencies2conda create -n qafd-rag python=3.10 -y
3conda activate qafd-rag
4pip install -r requirements.txt
56# 2. Set your OpenAI API key (required for LLM and answer generation)7exportOPENAI_API_KEY="sk-..."89# 3. Download pre-built KGs (recommended — saves hours of build time + API costs)10huggingface-cli download tarzanagh/QAFD-RAG --include "kg/multihop/*" --local-dir .11huggingface-cli download tarzanagh/QAFD-RAG --include "kg/ultradomain/*" --local-dir .1213# 4. Run a benchmark14# Note: First run for multihop will auto-download the nvidia/NV-Embed-v215# embedding model (~8GB, one-time, requires GPU with 16GB+ VRAM).16# No GPU? See "Rebuild with openai-small" below.17python benchmarks/run.py --task multihop --dataset musique --questions 1018python benchmarks/run.py --task ultradomain --dataset mix --questions 10
Embeddings: For best multihop results, download the pre-built KGs which use nvidia-nv-embed-v2 (requires GPU with 16GB+ VRAM, auto-downloaded on first run). If you don't have a GPU, you can rebuild the KGs from scratch with openai-small instead: python benchmarks/run.py --task multihop --dataset musique --force_build --embedding openai-small.
Graph Types: Multihop defaults to passage-entity graph (entities + passages + facts as nodes). UltraDomain/text2sql/summarization default to entity graph (classic KG). Override with --graph_type.
Two Graph Types
QAFD-RAG supports two knowledge graph representations, both using the same Query-Aware Flow Diffusion algorithm:
Entity Graph
Passage-Entity Graph
Nodes
Entities + relationships
Entities + passages + facts (Gutiérrez et al., 2024)
Extraction
LLM entity/relationship extraction
OpenIE (NER + triple extraction)
Edges
Entity-entity relationships
Fact edges + passage edges + synonymy edges
QAFD traversal
Flow reaches entities, passages looked up after
Flow reaches passages directly as graph nodes
Best for
General QA, text2sql, summarization
Multi-hop reasoning
Default tasks
ultradomain, text2sql, summarization
multihop
Both graph types use query-aware flow diffusion with the same parameters (alpha=1.5, epsilon=0.01, step_size=0.2, weight_scheme=multiply, linking_top_k=10).
Benchmarks
QAFD-RAG is evaluated on four tasks. Each subsection below covers data setup, pre-built KGs, and how to run.
The runner automatically selects the appropriate graph type (passage-entity for multihop, entity for others). Override with --graph_type. A legacy CLI (./run.sh) is also available for entity-graph benchmarks.
Pre-built KGs for all tasks are available at huggingface.co/tarzanagh/QAFD-RAG. Downloading is recommended to avoid hours of build time and API costs.
For the full Spider2-lite benchmark, clone Spider2 and copy databases into data/text2sql/spider2-lite/sqlite/.
bash
1# Run on included example databases2./run.sh text2sql --questions 5 --db Pagila # Spider2-lite (auto-detected)3./run.sh text2sql --questions 5 --db superhero # Bird (auto-detected)4./run.sh text2sql --benchmark bird --questions 5 --db superhero # explicit benchmark selection56# Build KG only (no benchmark)7python benchmarks/run.py --task text2sql --dataset spider2-lite --build_only --db Pagila
Adding a new database: Place your .sqlite file in data/text2sql/spider2-lite/sqlite/<DB_Name>/, then generate a DB summary (see Generating DB Summaries). The benchmark will auto-build the KG on first run.
Local embeddings (jina-v3, gritlm, nvidia-nv-embed-v2) run on your GPU and do not require an API key for embeddings. They are downloaded automatically from HuggingFace on first use.
Requirements:
CUDA-capable GPU with sufficient VRAM (8GB+ recommended)
Models are cached in ~/.cache/huggingface/
Usage with local embeddings (no OpenAI API needed for embeddings):
bash
1# Use Jina v3 (1024-dim, lightweight, good quality)2./run.sh ultradomain --questions 10 --embedding jina-v3
34# Use GritLM (4096-dim, unified embedding+generation)5./run.sh multihop --dataset musique --questions 10 --embedding gritlm
67# Use NVIDIA NV-Embed-v2 (4096-dim, 32K context, high quality)8./run.sh ultradomain --questions 10 --embedding nvidia-nv-embed-v2
Note: Even with local embeddings, an LLM API key is still required for entity extraction (KG building) and answer generation. Set OPENAI_API_KEY or use --llm gpt-oss-120b for a free open-source LLM.
QAFD-RAG processes documents through a two-stage pipeline:
Stage 1 -- Knowledge Graph Construction
Entity Graph:
Documents are split into token-based chunks.
An LLM extracts entities and relationships from each chunk.
Entities become nodes, relationships become weighted edges.
Passage-Entity Graph:
Documents are split into passages (chunks).
OpenIE extracts named entities and (subject, predicate, object) triples.
Entities and passages are both graph nodes, connected by fact edges, passage-entity edges, and synonymy edges (entity pairs with cosine similarity > 0.8).
Stage 2 -- Query-Aware Flow Diffusion
Seed Selection: Query is matched to entities/facts via embedding similarity. An LLM reranker filters the most relevant facts.
Flow Diffusion: Mass is injected at seed nodes and propagated through the graph via push-relabel. Edge weights are dynamically adjusted based on each node's similarity to the query.
Passage Ranking: Nodes accumulate importance scores proportional to flow. In the passage-entity graph, passages are ranked directly. In the entity graph, associated text chunks are retrieved from top-ranked entities.
LLM Answering: Top passages are assembled into context and passed to an LLM for answer generation.
The flow diffusion algorithm is the key differentiator: rather than simple graph traversal or vector search alone, it combines graph structure with query relevance to find contextually important information that may be several hops away from the initial match.
Citation
bibtex
1@inproceedings{zhou2026qafd,
2 title={Query-Aware Flow Diffusion for Graph-Based RAG with Retrieval Guarantees},
3 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},
4 booktitle={International Conference on Learning Representations (ICLR)},
5 year={2026}
6}
