AutoRAG-HP: Sintonização Automática de Hiperparâmetros Online para Geração Aprimorada por Recuperação

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

June 27, 2024

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

cs.AI

Resumo

Os avanços recentes em Modelos de Linguagem de Grande Escala transformaram o desenvolvimento de ML/AI, tornando necessária uma reavaliação dos princípios de AutoML para os sistemas de Geração com Recuperação Aprimorada (RAG). Para lidar com os desafios de otimização de hiperparâmetros e adaptação online em RAG, propomos o framework AutoRAG-HP, que formula a afinação de hiperparâmetros como um problema de bandit multi-armado (MAB) online e introduz um novo método Hierárquico de MAB em dois níveis (Hier-MAB) para exploração eficiente de grandes espaços de busca. Realizamos experimentos extensivos na afinação de hiperparâmetros, como documentos recuperados top-k, taxa de compressão de prompt e métodos de embedding, utilizando os conjuntos de dados ALCE-ASQA e Natural Questions. Nossa avaliação da otimização conjunta dos três hiperparâmetros demonstra que métodos de aprendizado online baseados em MAB podem alcançar Recall@5 aproximadamente 0.8 para cenários com gradientes proeminentes no espaço de busca, utilizando apenas 20% das chamadas de API do LLM necessárias pela abordagem de Busca em Grade. Além disso, a abordagem Hier-MAB proposta supera outras baselines em cenários de otimização mais desafiadores. O código estará disponível em 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.