TLDR: This research introduces a multi-armed bandit framework for Retrieval-Augmented Generation (RAG) systems to efficiently handle complex user queries. By treating sub-queries as “arms” and dynamically balancing exploration (finding new relevant information) and exploitation (deepening retrieval for promising sub-queries), the system significantly improves document precision, α-nDCG, and downstream report generation performance. The approach, especially when incorporating rank information and hierarchical query structures, allows for adaptive document allocation under budget constraints, leading to more accurate and less noisy information retrieval.
In the evolving landscape of artificial intelligence, Retrieval-Augmented Generation (RAG) systems are becoming increasingly vital for handling complex user queries. These systems work by breaking down a user’s request into smaller, more manageable sub-queries, then finding relevant documents for each, and finally combining this information to generate a comprehensive answer. A new research paper titled “Query Decomposition for RAG: Balancing Exploration-Exploitation” delves into a novel approach to optimize this process, ensuring that RAG systems retrieve the most informative documents efficiently.
The core challenge in RAG is a delicate balance: retrieving broadly enough to capture all relevant information (exploration) versus limiting retrieval to avoid overwhelming the system with irrelevant data and computational costs (exploitation). The authors, Roxana Petcu, Kenton Murray, Daniel Khashabi, Evangelos Kanoulas, Maarten de Rijke, Dawn Lawrie, and Kevin Duh, propose framing this problem within an exploitation-exploration setting, specifically using a multi-armed bandit framework.
Imagine each sub-query as an “arm” in a slot machine. When you “pull” an arm (retrieve a document for a sub-query), you get feedback on its utility (relevance). This feedback helps the system decide whether to continue “exploiting” that promising sub-query by retrieving more documents from it, or to “explore” other sub-queries that might yield better results. This adaptive strategy is crucial because traditional methods often retrieve a fixed number of documents for each sub-query, regardless of its actual usefulness, leading to inefficiencies.
The research highlights that this bandit learning approach, where documents are retrieved one at a time, allows the system to build a “belief” about the relevance of each sub-query. This belief then guides subsequent retrieval decisions under a fixed budget. This is particularly beneficial for complex queries that involve exclusions, negations, or missing entities, which often require evidence from multiple documents.
The paper introduces several reward policies for the bandit framework, incorporating factors like rank information (the implicit relevance encoded in a document’s position in a search result list), document diversity (to avoid retrieving redundant information), and an exploration term (to ensure all sub-queries get a chance). Their main finding is that estimating document relevance using rank information and human judgments significantly improves document-level precision by 35% and increases α-nDCG by 15%. This also leads to better performance in generating long-form answers.
Furthermore, the study explores hierarchical query decomposition, where complex requests are broken down into sub-queries, which can then be further decomposed. By modeling correlations between these hierarchical sub-queries, the system can selectively expand the most informative branches of the query tree. This “correlated bandit” approach demonstrated a 30% precision gain over single-level decomposition, especially under small retrieval budgets, by focusing on promising information paths early on.
The researchers conducted experiments on two datasets: NeuCLIR, which features decomposable user requests, and ResearchyQuestions, which contains hierarchically decomposed questions. The results consistently showed that bandit learning outperforms simpler “full exploitation” (always picking the best-known option) and “full exploration” (randomly trying options) strategies. The benefits extend to downstream tasks like report generation, where selecting documents using multi-armed bandits boosted citation support, nugget coverage, and sentence support by 6.0–9.9%, 6.7–8.5%, and 7.3–9.6% respectively.
Also Read:
- Adaptive Search: How Reinforcement Learning Powers Intelligent AI Agents
- GraphFlow: Enhancing LLM Retrieval with Accurate and Diverse Knowledge from Graphs
In conclusion, this research presents a compelling case for using multi-armed bandit frameworks to manage query decomposition and document retrieval in RAG systems. By dynamically balancing exploration and exploitation, these systems can more efficiently identify and retrieve relevant information, leading to more precise and comprehensive answers, even for the most intricate user requests. For more technical details, you can refer to the full paper here.
