AdaRAG: Budget-Aware Adaptive Retrieval-Augmented Generation via Hierarchical Reinforcement Learning

Abstract

Multi-turn retrieval-augmented generation (RAG) improves question answering by decomposing evidence seeking into iterative retrieval and reasoning steps. Existing multi-turn RAG methods usually optimize when and how to retrieve while fixing the number of retrieved documents per step. However, we discovered that this fixed-TopK design is suboptimal: single-hop questions tend to benefit from fewer retrieval rounds with larger per-round evidence sets, whereas multi-hop questions require more retrieval rounds with smaller evidence sets to support stepwise reasoning. To bridge this gap, we introduce AdaRAG, a budget-aware adaptive RAG framework that learns how to retrieve under a hard document budget, including how many retrieval rounds to perform, how many documents to retrieve in each round, and which retrieval source to use. AdaRAG implements this idea with a two-level policy architecture. ModeHead, a lightweight retrieval-mode classifier, selects passage retrieval, graph retrieval, or answer generation; TopkHead, a budget-aware document-allocation classifier, selects a legal TopK after query generation according to the remaining budget. These discrete policy heads are decoupled from language-model token generation, enabling direct reinforcement-learning optimization through hierarchical GRPO after supervised action-format learning. Our experiments across five QA benchmarks demonstrate AdaRAG's good generalization performance under constrained document budgets. In detailed comparisons on HotpotQA, it surpasses the strongest baselines by an average of 10.8 percentage points in Exact Match (EM) and F1 score.