Query-Aware Context Selection for Retrieval-Augmented Generation

Abstract

Retrieval augmented generation (RAG) combines language models with external corpora to support knowledge-intensive tasks, such as open-domain question answering. Standard RAG systems typically employ a fixed top-k retrieval strategy, retrieving the same number of passages regardless of query needs. This can lead to either insufficient evidence, or the inclusion of irrelevant contexts that lead to a degradation of generation performance. In this work, we conduct an empirical study of how irrelevant retrieved passages affect downstream generation, analyzing their impact across multiple standard generator models. Building on these insights, we propose a lightweight, context-size classification module that dynamically predicts how much context is required based on query-specific needs. We integrate this approach into a full RAG pipeline and demonstrate improved performance over several baselines.