AutoRAG-HP: Ajuste Automático en Línea de Hiperparámetros para la Generación Aumentada con Recuperación

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

Resumen

Los recientes avances en los Modelos de Lenguaje de Gran Escala han transformado el desarrollo de ML/IA, lo que ha llevado a una reevaluación de los principios de AutoML para los sistemas de Generación Aumentada por Recuperación (RAG, por sus siglas en inglés). Para abordar los desafíos de la optimización de hiperparámetros y la adaptación en línea en RAG, proponemos el marco AutoRAG-HP, que formula el ajuste de hiperparámetros como un problema de bandido multi-brazo (MAB) en línea e introduce un novedoso método de MAB Jerárquico de dos niveles (Hier-MAB) para una exploración eficiente de grandes espacios de búsqueda. Realizamos experimentos exhaustivos en el ajuste de hiperparámetros, como el número de documentos recuperados (top-k), la relación de compresión de prompts y los métodos de incrustación, utilizando los conjuntos de datos ALCE-ASQA y Natural Questions. Nuestra evaluación, basada en la optimización conjunta de los tres hiperparámetros, demuestra que los métodos de aprendizaje en línea basados en MAB pueden alcanzar un Recall@5 de aproximadamente 0.8 en escenarios con gradientes prominentes en el espacio de búsqueda, utilizando solo alrededor del 20% de las llamadas a la API de LLM requeridas por el enfoque de Búsqueda en Cuadrícula. Además, el enfoque Hier-MAB propuesto supera a otras líneas base en escenarios de optimización más desafiantes. El código estará disponible en 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.