QAFD-RAG

Official code for ICLR 2026 paper: Query-Aware Flow Diffusion for Graph-Based RAG with Retrieval Guarantees

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

LightRAG

QAFD-RAG

QAFD-RAG reweights edges by query relevance, suppressing irrelevant neighborhoods (e.g., Amazon, Apple fruit).

Quick Start

# 1. Create conda environment and install dependencies
conda create -n qafd-rag python=3.10 -y
conda activate qafd-rag
pip install -r requirements.txt

# 2. Set your OpenAI API key (required for LLM and answer generation)
export OPENAI_API_KEY="sk-..."

# 3. Download pre-built KGs (recommended — saves hours of build time + API costs)
huggingface-cli download tarzanagh/QAFD-RAG --include "kg/multihop/*" --local-dir .
huggingface-cli download tarzanagh/QAFD-RAG --include "kg/ultradomain/*" --local-dir .

# 4. Run a benchmark
#    Note: First run for multihop will auto-download the nvidia/NV-Embed-v2
#    embedding model (~8GB, one-time, requires GPU with 16GB+ VRAM).
#    No GPU? See "Rebuild with openai-small" below.
python benchmarks/run.py --task multihop --dataset musique --questions 10
python 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.

All benchmarks use the unified runner:

python benchmarks/run.py --task <task> --dataset <dataset> [options]

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.

Multi-hop QA

Datasets	MuSiQue, HotpotQA, 2WikiMultiHopQA
Graph type	passage-entity (default)
Metrics	F1, Exact Match
Data	Included in data/multihop/ (from osunlp)
Embedding	nvidia-nv-embed-v2 (GPU, 16GB+ VRAM)

# Download pre-built KGs (recommended — saves hours + API costs)
huggingface-cli download tarzanagh/QAFD-RAG --include "kg/multihop/*" --local-dir .

# Run benchmark
python benchmarks/run.py --task multihop --dataset musique --questions 100
python benchmarks/run.py --task multihop --dataset hotpotqa --questions 100
python benchmarks/run.py --task multihop --dataset 2wikimultihopqa --questions 100

# No GPU? Rebuild KGs with openai-small instead
python benchmarks/run.py --task multihop --dataset musique --force_build --embedding openai-small

# Retrieval only (skip answer generation)
python benchmarks/run.py --task multihop --dataset musique --skip_qa

UltraDomain

Datasets	agriculture, biology, cs, finance, legal, math, medicine, mix, music, philosophy, physics
Graph type	entity (default)
Metrics	Quality scores (comprehensiveness, diversity, relevance, logicality, coherence)
Data	Auto-downloaded from TommyChien/UltraDomain at runtime
Embedding	openai-small (no GPU needed)

# Download pre-built KGs (recommended)
huggingface-cli download tarzanagh/QAFD-RAG --include "kg/ultradomain/*" --local-dir .

# Run benchmark
python benchmarks/run.py --task ultradomain --dataset mix --questions 10

# Use passage-entity graph instead of entity graph
python benchmarks/run.py --task ultradomain --dataset mix --graph_type passage-entity --questions 10

Text-to-SQL

Datasets	Spider2-lite, Bird
Graph type	entity
Metrics	Schema retrieval precision / recall
Data	Example databases included in data/text2sql/ (see Data section)
Embedding	openai-small (no GPU needed)

Two example databases are included with pre-generated DB summaries:

• Pagila (Spider2-lite) -- DVD rental store, 16 tables, SQLite
• superhero (Bird) -- superhero database, 10 tables, SQLite

For the full Spider2-lite benchmark, clone Spider2 and copy databases into data/text2sql/spider2-lite/sqlite/.

# Run on included example databases
./run.sh text2sql --questions 5 --db Pagila                          # Spider2-lite (auto-detected)
./run.sh text2sql --questions 5 --db superhero                       # Bird (auto-detected)
./run.sh text2sql --benchmark bird --questions 5 --db superhero      # explicit benchmark selection

# Build KG only (no benchmark)
python 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.

Summarization

Dataset	SQuALITY
Graph type	entity
Metrics	BLEU, ROUGE, METEOR, quality scores
Data	Auto-downloaded from pszemraj/SQuALITY-v1.3 at runtime

python benchmarks/run.py --task summarization --dataset squality --questions 50

Common Options

Option	Description
--task TASK	multihop, ultradomain, text2sql, summarization
--dataset NAME	Dataset name (e.g., musique, mix, spider2-lite)
--graph_type TYPE	passage-entity or entity (auto-selected by task)
--questions N	Number of questions to evaluate
--build_only	Build KG only, skip benchmark
--force_build	Rebuild KG even if it exists
--skip_qa	Run retrieval only, skip QA (passage-entity)
--embedding MODEL	Embedding model (see Configuration)
--llm MODEL	LLM model (gpt-4o-mini, gpt-4o, gpt-5-nano, gpt-5-mini, gpt-5, gpt-oss-120b)
--alpha FLOAT	QAFD alpha parameter (default: 2.0)
--epsilon FLOAT	QAFD convergence threshold (default: 0.01)
--weight_scheme	Query-aware edge weighting: original, multiply, add

Data

Included Data

Task	Directory	Contents	Size
Multi-hop QA	data/multihop/	MuSiQue, HotpotQA, 2WikiMultiHopQA (full datasets from osunlp)	134 MB
UltraDomain	data/ultradomain/	Auto-downloaded from HuggingFace at runtime	--
Text-to-SQL	data/text2sql/	Example databases + question files (see below)	9.6 MB
Summarization	data/summarization/	Auto-downloaded from HuggingFace at runtime	--

Text-to-SQL Data

Example databases with pre-generated DB summaries are included so you can run the text2sql benchmark out of the box:

Backend	Database	Files	Description
SQLite	Pagila	Pagila.sqlite + Pagila_db_summary.json	DVD rental store (16 tables)
SQLite	superhero	superhero.sqlite + superhero_db_summary.json	Superhero database (10 tables)
BigQuery	san_francisco	san_francisco_bigquery_summary.json	SF city data: bikeshare, crime, film locations, fire dept, street trees (8 tables, 118 columns, schema only)
Snowflake	AUSTIN	AUSTIN_db_summary.json	Austin 311 service requests (schema only)

Question files: spider2-lite.jsonl (546 questions) and bird.jsonl (129 questions).

For the full Spider2-lite benchmark with all databases, clone Spider2 and copy the SQLite files into data/text2sql/spider2-lite/sqlite/.

Generating DB Summaries

To add a new SQLite database to the text2sql benchmark, generate a DB summary using the included CLI tool:

# Generate DB summary for a SQLite database
python -m src.indexing.extract_db_summary \
    --db-path data/text2sql/spider2-lite/sqlite/MyDB/MyDB.sqlite

# Specify custom output directory
python -m src.indexing.extract_db_summary \
    --db-path path/to/database.sqlite \
    --output data/text2sql/spider2-lite/sqlite/MyDB/

This produces a <DB_Name>_db_summary.json file containing:

• Table and column metadata (types, primary/foreign keys, nullable)
• Cardinality and distinct value counts per column
• Sample values for each column
• Inferred foreign key relationships

For BigQuery and Snowflake databases (schema summaries only, no raw data):

# BigQuery
python -m src.indexing.extract_db_summary_bigquery \
    --datasets project_id.dataset_id \
    --output-dir data/text2sql/spider2-lite/bigquery/my_dataset/

# Snowflake
python -m src.indexing.extract_db_summary_snowflake \
    --databases MY_DATABASE \
    --output-dir data/text2sql/spider2-lite/snowflake/MY_DATABASE/

Project Structure

QAFD-RAG/
├── benchmarks/
│   ├── run.py                       # Unified benchmark runner (--graph_type, --task)
│   ├── multihop/                    # Multi-hop reasoning
│   ├── ultradomain/                 # General domain QA
│   ├── text2sql/                    # Natural language to SQL
│   └── summarization/               # Document summarization
├── src/
│   ├── QAFD_RAG.py                  # Entity graph pipeline (main class)
│   ├── base.py                      # Storage interfaces, QueryParam
│   ├── llm.py                       # LLM and embedding functions
│   ├── storage.py                   # KV, vector, and graph storage
│   ├── operate.py                   # Query operations
│   ├── evaluation.py                # Answer evaluation
│   ├── answering/                   # Query processing (entity graph)
│   │   ├── handler.py               # kg_query entry point
│   │   ├── context.py               # Context building (local/global/hybrid)
│   │   ├── clusters.py              # Flow diffusion clustering
│   │   └── text_units.py            # Text chunk retrieval
│   ├── retrievers/
│   │   ├── flow_diffusion.py        # QAFD on NetworkX (entity graph)
│   │   └── base.py                  # Retriever interface
│   ├── passage_entity/           # Passage-entity graph pipeline
│   │   ├── kg_builder.py            # OpenIE → igraph (entities + passages + facts)
│   │   ├── graph_adapter.py         # QAFD on igraph (passage-entity graph)
│   │   ├── retriever.py             # Fact reranking → seed selection → QAFD
│   │   ├── config.py                # Pipeline configuration
│   │   ├── openie.py                # NER + triple extraction
│   │   ├── embedding_store.py       # Parquet-backed vector store
│   │   ├── reranker.py              # LLM-based fact reranker
│   │   ├── prompts.py               # Prompt templates
│   │   ├── benchmark_runner.py      # Standalone runner
│   │   └── utils.py                 # Helpers
│   ├── indexing/                     # KG construction (entity graph)
│   │   ├── chunkers.py              # Token-based chunking
│   │   ├── extractors.py            # Entity/relationship extraction
│   │   ├── schema_builder.py        # Database schema KG builder
│   │   ├── extract_db_summary.py    # SQLite DB summary generator (CLI)
│   │   └── build_kg.py              # CLI KG builder
│   ├── embedding_models/            # Embedding model implementations
│   ├── prompts/                     # LLM prompt templates
│   ├── text2sql/                    # Text-to-SQL support
│   └── utils/                       # Helpers
├── data/                            # Datasets (included or auto-downloaded)
│   ├── multihop/                    # MuSiQue, HotpotQA, 2WikiMultiHopQA
│   ├── ultradomain/                 # Auto-downloaded at runtime
│   ├── text2sql/                    # Example DBs (Pagila, superhero) + questions
│   └── summarization/               # Auto-downloaded at runtime
├── kg/                              # Knowledge graphs (auto-generated or downloaded)
├── docs/figs/                       # Figures for README
└── results/                         # Benchmark results (JSON)

Configuration

Embedding Models

Key	Model	Dimensions
openai-small	text-embedding-3-small	1536
openai-large	text-embedding-3-large	3072
jina-v3	jinaai/jina-embeddings-v3	1024
gritlm	GritLM/GritLM-7B	4096
nvidia-nv-embed-v2	nvidia/NV-Embed-v2	4096

The default embedding is openai-small which requires OPENAI_API_KEY.

Use --embedding <key> to select a model:

./run.sh ultradomain --questions 10 --embedding jina-v3

Local Embedding Models

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):

# Use Jina v3 (1024-dim, lightweight, good quality)
./run.sh ultradomain --questions 10 --embedding jina-v3

# Use GritLM (4096-dim, unified embedding+generation)
./run.sh multihop --dataset musique --questions 10 --embedding gritlm

# Use NVIDIA NV-Embed-v2 (4096-dim, 32K context, high quality)
./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.

LLM Models

gpt-4o-mini (default), gpt-4o, gpt-5-nano, gpt-5-mini, gpt-5, gpt-oss-120b (free, local)

How It Works

QAFD-RAG processes documents through a two-stage pipeline:

Stage 1 -- Knowledge Graph Construction

Entity Graph:

1. Documents are split into token-based chunks.
2. An LLM extracts entities and relationships from each chunk.
3. Entities become nodes, relationships become weighted edges.

Passage-Entity Graph:

1. Documents are split into passages (chunks).
2. OpenIE extracts named entities and (subject, predicate, object) triples.
3. 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

1. Seed Selection: Query is matched to entities/facts via embedding similarity. An LLM reranker filters the most relevant facts.
2. 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.
3. 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.
4. 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

@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}
}