AutoRAG-HP: 검색 증강 생성(Retrieval-Augmented Generation)을 위한 자동 온라인 하이퍼파라미터 튜닝

AutoRAG-HP: Automatic Online Hyper-Parameter Tuning for Retrieval-Augmented Generation

June 27, 2024

저자: Jia Fu, Xiaoting Qin, Fangkai Yang, Lu Wang, Jue Zhang, Qingwei Lin, Yubo Chen, Dongmei Zhang, Saravan Rajmohan, Qi Zhang

cs.AI

초록

대형 언어 모델(Large Language Models)의 최근 발전은 머신러닝/인공지능(ML/AI) 개발을 혁신적으로 변화시켰으며, 이로 인해 검색 증강 생성(Retrieval-Augmented Generation, RAG) 시스템을 위한 AutoML 원칙의 재평가가 필요하게 되었습니다. RAG에서의 하이퍼파라미터 최적화와 온라인 적응의 문제를 해결하기 위해, 우리는 AutoRAG-HP 프레임워크를 제안합니다. 이 프레임워크는 하이퍼파라미터 튜닝을 온라인 다중 슬롯 머신(Multi-Armed Bandit, MAB) 문제로 공식화하고, 대규모 탐색 공간을 효율적으로 탐색하기 위한 새로운 2단계 계층적 MAB(Hierarchical MAB, Hier-MAB) 방법을 도입합니다. 우리는 ALCE-ASQA 및 Natural Questions 데이터셋을 사용하여 상위 k개의 검색 문서, 프롬프트 압축 비율, 임베딩 방법과 같은 하이퍼파라미터 튜닝에 대한 광범위한 실험을 수행했습니다. 세 가지 하이퍼파라미터를 동시에 최적화한 평가 결과, MAB 기반 온라인 학습 방법은 탐색 공간에서 뚜렷한 기울기가 있는 시나리오에서 Grid Search 접근 방식에 필요한 LLM API 호출의 약 20%만 사용하여 Recall@5 약 0.8을 달성할 수 있음을 보여주었습니다. 또한, 제안된 Hier-MAB 접근 방식은 더 어려운 최적화 시나리오에서 다른 베이스라인을 능가하는 성능을 보였습니다. 코드는 https://aka.ms/autorag에서 공개될 예정입니다.

English

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.