Query Decomposition for RAG: Balancing Exploration-Exploitation

TL;DR

The paper reframes query decomposition for retrieval-augmented generation as an exploitation–exploration bandit problem, where each sub-query is an arm and retrieval is performed under a fixed budget. It introduces rank-aware Bernoulli rewards, diversity considerations, and Thompson sampling, and extends to hierarchical, correlated MAB to exploit informative sub-queries. Across NeuCLIR and ResearchyQuestions, the approach yields consistent gains in document-level precision and alpha-nDCG, and improves downstream long-form generation with Auto-ARGUE, especially at small budgets. The work demonstrates that principled budgeted retrieval with sub-query hierarchy enhances grounding and reduces noise, offering practical guidance for efficient RAG systems. Future work could explore contextual and deeper hierarchical correlations and joint optimization of retrieval and generation objectives.

Abstract

Retrieval-augmented generation (RAG) systems address complex user requests by decomposing them into subqueries, retrieving potentially relevant documents for each, and then aggregating them to generate an answer. Efficiently selecting informative documents requires balancing a key trade-off: (i) retrieving broadly enough to capture all the relevant material, and (ii) limiting retrieval to avoid excessive noise and computational cost. We formulate query decomposition and document retrieval in an exploitation-exploration setting, where retrieving one document at a time builds a belief about the utility of a given sub-query and informs the decision to continue exploiting or exploring an alternative. We experiment with a variety of bandit learning methods and demonstrate their effectiveness in dynamically selecting the most informative sub-queries. Our main finding is that estimating document relevance using rank information and human judgments yields a 35% gain in document-level precision, 15% increase in α-nDCG, and better performance on the downstream task of long-form generation.

Figures (13)

• Figure 1: A user request is first decomposed into sub-queries. Bandit learning iteratively selects a sub-query (arm), retrieves a document down its ranked list, observes its relevance, and updates the sub-query posterior belief over time. The selected documents across iterations are then used as evidence in a RAG setting to generate a grounded answer, balancing exploration (more sub-queries) and exploitation (more documents per sub-query).
• Figure 2: Exploration–exploitation across $n$ subqueries and $m$ documents; colors indicate document relevance, and arrows show how exploitation, exploration, and bandit-based policies allocate effort across sub-queries.
• Figure 3: Hierarchical query decomposition with correlated bandits, where informative sub-queries are decomposed into arms with inherited posterior beliefs.
• Figure 4: Precision scores on NeuCLIR24 and ResearchyQuestions for baselines and reward policies, evaluated across varying budgets of observed documents. The budgets are expressed as percentages of the total available corpus. Modelling the queries using Bernoulli distribution considering the top ranked k documents performs the best across budgets, with all policies reaching the same performance as the budget covers the whole space of actions.
• Figure 5: Fitting linear regression functions $y = \alpha + \beta x +\epsilon$ between the mean reciprocal ranks $\bar{r}_i = \frac{1}{K}\sum_{k=1}^K r_{i,k}$ for each user request in NeuCLIR24 (represented by a point).