Can Knowledge-Graph-based Retrieval Augmented Generation Really Retrieve What You Need?

Retrieval-Augmented Generation (RAG) based on knowledge graphs (KGs) enhances large language models (LLMs) with structural and textual external knowledge. Yet, existing KG-based RAG methods struggle to retrieve accurate and diverse information when handling complex queries. By modeling KG-based retrieval as a multi-step decision process, Process Reward Models (PRMs) offer a promising solution to align the retrieval behavior with the query-specific knowledge requirements. However, PRMs heavily rely on process-level supervision signals that are expensive and hard to obtain on KGs. To address this challenge, we propose GraphFlow, a framework that efficiently retrieves accurate and diverse knowledge required for complex queries from text-rich KGs. GraphFlow employs a detailed balance objective with local exploration to jointly optimize a retrieval policy and a flow estimator. The flow estimator factorizes the outcome reward of the retrieval results into the intermediate retrieval steps. Such reward factorization guides the retrieval policy to retrieve candidates from KGs in proportion to their outcome reward. This allows GraphFlow to explore relevant regions of KGs that yield diverse and accurate results. We evaluate GraphFlow on STaRK benchmark, which includes real-world queries from multiple domains over text-rich KGs. GraphFlow outperforms strong KG-based RAG baselines including GPT-4o by 10% performance gain on both retrieval accuracy and diversity metrics. GraphFlow also shows strong generalization by effectively retrieving information from unseen KGs to support new-domain queries, highlighting its effectiveness and robustness 2 .