A uto RAG - HP : Automatic Online Hyper-Parameter Tuning for Retrieval-Augmented Generation (2024)

Abstract: Recent advancements in Large Language Models have transformed ML/AI development, necessitating a reevaluation of AutoML principles for the Retrieval-Augmented Generation (RAG) systems. To address the challenges of hyper-parameter optimization and online adaptation in RAG, we propose the AutoRAG-HP framework, which formulates the hyper-parameter tuning as an online multi-armed bandit (MAB) problem and introduces a novel two-level Hierarchical MAB (Hier-MAB) method for efficient exploration of large search spaces. We conduct extensive experiments on tuning hyper-parameters, such as top-k retrieved documents, prompt compression ratio, and embedding methods, using the ALCE-ASQA and Natural Questions datasets. Our evaluation from jointly optimization all three hyper-parameters demonstrate that MAB-based online learning methods can achieve Recall@5 \approx 0.8 for scenarios with prominent gradients in search space, using only ~20% of the LLM API calls required by the Grid Search approach. Additionally, the proposed Hier-MAB approach outperforms other baselines in more challenging optimization scenarios. The code will be made available at https://aka.ms/autorag.

## introduction recent advancements in large language models (llms) @xcite @xcite represent a significant shift in the development of ml/ai solutions. traditionally, scenario-specific models were trained for most ml/ai applications. however, in the llm era, foundational models serve as the base, with supplementary modules added for practical applications. this transformation significantly affects the automation of ml/ai solution development, previously known as automl @xcite , necessitating a reevaluation of automl concepts in the context of llms. retrieval-augmented generation (rag) has emerged as a prominent framework for building ml/ai solutions with llms @xcite . while the standard rag framework includes an information retrieval component to ground llm's output in relevant data, numerous variants now integrate additional modules such as query rewriting @xcite , prompt compression @xcite , and query routing @xcite to enhance performance. the increased complexity of rag systems presents two main challenges. first, the multitude of modules and hyper-parameters within the modules complicates the identification of optimal settings. second, as we often receive online feedback from users (e.g., via thumb up/down feature), effectively utilizing those feedback to continuously tune the system is also crucial. to address these challenges, we propose the development of an autonomous and self-optimizing system for rag, termed autorag, in line with the principles of automl. as the first step, this study focuses on hyper-parameter tuning in rag (autorag-hp). while there exist prior works discussing hyper-parameter tuning in rag, they tend to either focus on tunable hyper-parameters in llm api calls @xcite or assess the performance and impacts of rag hyper-parameters through manual tuning @xcite . in this work we focus on the optimization methods that can be applied in the online fashion. specifically, we frame hyper-parameter selection as an online multi-armed bandit (mab) problem (lai and @xcite @xcite and explore several variants in mab. moreover, to efficiently explore large search space when tuning multiple hyper-parameters simultaneously, we introduce a novel two-level hierarchical mab (hier-mab) method, wherein a high-level mab guides the optimization of modules, while several low-level mabs search for optimal settings within each module. our evaluation demonstrates that the mab-based online learning methods are effective for scenarios with prominent gradients in search space, and the proposed hier-mab approach outperforms other baselines in more challenging optimization scenarios. our contributions can be summarized as follows: ## discussion in this work, we address online optimization in rag systems by framing hyper-parameter tuning as an online multi-armed bandit problem. our proposed hier-mab approach can be extended to offline hyper-parameter tuning, demonstrating greater efficiency than traditional grid search methods, especially in scenarios with large search spaces with steep gradients. hier-mab can also serve as an initial step to filter the search space, followed by a more extensive parameter scan. our approach can be applied to a broader range of tunable hyper-parameters. while we focus on hyper-parameters in retrieval and prompt compression modules, it is extendable to other rag modules such as document chunk size in indexing module, llm api settings, or other llm-based solutions (e.g., agent frameworks). due to computational constraints, exhaustive searches for optimal configurations as the ground-truth are challenging, limiting the feasibility of experiments across broader hyper-parameter combinations. the reward settings in our work can also be expanded. currently, we assume the reward is a linear combination of accuracy and input token length, with a user-defined weight parameter allowing users to adjust the weight parameter w. however, determining the right parameter to balance accuracy and llm api cost is challenging in reality. an alternative is to set llm api cost constraints and optimize accuracy within these constraints, incorporating cost constraints as penalty terms in the reward function. this converts a multi-objective optimization problem into a single-objective one, though exploring pareto optimization within rag could also yield valuable insights. our current reward framework only considers feedback from the correctness of the final response. in practice, feedback might also come from intermediate steps (e.g., document relevance evaluation in retrieval modules) or multiple sources in multiturn dialogues. thus, automatically and efficiently integrating these additional feedback sources into the reward definition is also worth exploration. beyond hyper-parameter tuning, developing a comprehensive automl framework for rag involves identifying the optimal combination of available rag modules, automated prompt tuning and other query-dependent parameters, such as those in a routing module that directs queries to appropriate base llms. additionally, an ideal autorag system should auto-generate evaluation data for tuning as needed, supporting the "bring your data" vision where users provide their data, and the platform autonomously configures the entire pipeline to meet their specific requirements. future work will explore these areas. ## summary inspired by traditional automl practices designed to simplify and automate ml/ai development, we introduce the autorag-hp framework. this framework addresses the critical need for efficient and effortless hyper-parameter tuning within the retrieval-augmented generation (rag) system in the context of llms. to address challenges posed by extensive search spaces and the need for online tuning, we formulate hyper-parameter selection in rag as a multi-armed bandit problem and introduce a novel two-level hierarchical upper confidence bound (hier-ucb) method for efficient parameter space exploration. our experiments on simultaneously tuning three hyper-parameters demonstrate that multi-armed bandit-based online learning methods (hier-ucb, ucb, and ts) can achieve recall@5 ≈ 0.8 for scenarios with prominent gradients in search space, using only ∼ 20% of the llm api calls required by the grid search approach. additionally, the proposed hier-ucb approach outperforms other baselines in more challenging optimization scenarios. these promising results motivate further exploration into automatic tuning of the rag system to achieve the full vision of autorag.