AutoRAG-HP: Automatic Online Hyper-Parameter Tuning for Retrieval-Augmented Generation
Jia Fu, Xiaoting Qin, Fangkai Yang, Lu Wang, Jue Zhang, Qingwei Lin, Yubo Chen, Dongmei Zhang, Saravan Rajmohan, Qi Zhang
2024-07-01
Summary
This paper talks about AutoRAG-HP, a new system designed to automatically adjust the settings (hyperparameters) of Retrieval-Augmented Generation (RAG) models during their training. This helps improve the models' performance on various tasks without needing manual tuning.
What's the problem?
As AI technology has advanced, especially with large language models, there's a growing need to optimize how these models are trained. Hyperparameter tuning is crucial because it involves setting important parameters that control how a model learns. However, finding the best settings can be time-consuming and difficult, especially when using traditional methods that require many trial-and-error attempts.
What's the solution?
To tackle this problem, the authors developed AutoRAG-HP, which uses a method called the multi-armed bandit (MAB) approach to make hyperparameter tuning more efficient. This method allows the system to explore different settings dynamically while the model is training, rather than having to set everything beforehand. They introduced a two-level Hierarchical MAB method to effectively search through a wide range of hyperparameter options. Their experiments showed that this approach could achieve high performance with significantly fewer API calls compared to older methods.
Why it matters?
This research is important because it simplifies and speeds up the process of optimizing AI models. By automating hyperparameter tuning, AutoRAG-HP makes it easier for researchers and developers to get better results without spending excessive time on manual adjustments. This advancement can lead to more effective AI applications in various fields, improving how machines understand and generate information.
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 sim20% 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.