MaskSearch: Un Marco de Preentrenamiento Universal para Mejorar la Capacidad de Búsqueda Agente

MaskSearch: A Universal Pre-Training Framework to Enhance Agentic Search Capability

May 26, 2025

Autores: Weiqi Wu, Xin Guan, Shen Huang, Yong Jiang, Pengjun Xie, Fei Huang, Jiuxin Cao, Hai Zhao, Jingren Zhou

cs.AI

Resumen

Los Modelos de Lenguaje Aumentados con Recuperación (RALMs, por sus siglas en inglés) representan un paradigma clásico en el que los modelos mejoran sus capacidades generativas utilizando conocimiento externo recuperado a través de un módulo especializado. Los avances recientes en técnicas de Agentes permiten que los Modelos de Lenguaje de Gran Escala (LLMs, por sus siglas en inglés) utilicen herramientas de manera autónoma para la recuperación, planificación y razonamiento. Si bien los métodos basados en entrenamiento existentes muestran potencial, sus habilidades agentivas están limitadas por las características inherentes de los datos específicos de la tarea utilizados durante el entrenamiento. Para mejorar aún más la capacidad de búsqueda universal de los agentes, proponemos un novedoso marco de pre-entrenamiento, denominado MaskSearch. En la etapa de pre-entrenamiento, introducimos la tarea de Predicción de Mascaras Aumentada con Recuperación (RAMP, por sus siglas en inglés), donde el modelo aprende a aprovechar herramientas de búsqueda para completar segmentos enmascarados en una gran cantidad de datos de pre-entrenamiento, adquiriendo así capacidades universales de recuperación y razonamiento para los LLMs. Posteriormente, el modelo se entrena en tareas posteriores para lograr una mejora adicional. Aplicamos tanto el Ajuste Supervisado (SFT, por sus siglas en inglés) como el Aprendizaje por Refuerzo (RL, por sus siglas en inglés) para el entrenamiento. Para el SFT, combinamos métodos basados en agentes y en destilación para generar datos de entrenamiento, comenzando con un sistema multiagente compuesto por un planificador, reescritor, observador y seguido por un modelo docente auto-evolutivo. Mientras que para el RL, empleamos DAPO como marco de entrenamiento y adoptamos un sistema de recompensas híbrido que incluye recompensas por respuestas y recompensas por formato. Además, introducimos un enfoque de aprendizaje curricular que permite al modelo aprender progresivamente desde instancias más fáciles hasta más desafiantes, basado en el número de segmentos enmascarados. Evaluamos la efectividad de nuestro marco en el escenario de respuesta a preguntas de múltiples saltos en dominios abiertos. A través de experimentos exhaustivos, demostramos que MaskSearch mejora significativamente el rendimiento de los agentes de búsqueda basados en LLMs tanto en tareas posteriores dentro del dominio como fuera de él.

English

Retrieval-Augmented Language Models (RALMs) represent a classic paradigm where models enhance generative capabilities using external knowledge retrieved via a specialized module. Recent advancements in Agent techniques enable Large Language Models (LLMs) to autonomously utilize tools for retrieval, planning, and reasoning. While existing training-based methods show promise, their agentic abilities are limited by inherent characteristics of the task-specific data used during training. To further enhance the universal search capability of agents, we propose a novel pre-training framework, MaskSearch. In the pre-training stage, we introduce the Retrieval Augmented Mask Prediction (RAMP) task, where the model learns to leverage search tools to fill masked spans on a large number of pre-training data, thus acquiring universal retrieval and reasoning capabilities for LLMs. After that, the model is trained on downstream tasks to achieve further improvement. We apply both Supervised Fine-tuning (SFT) and Reinforcement Learning (RL) for training. For SFT, we combine agent-based and distillation-based methods to generate training data, starting with a multi-agent system consisting of a planner, rewriter, observer, and followed by a self-evolving teacher model. While for RL, we employ DAPO as the training framework and adopt a hybrid reward system consisting of answer rewards and format rewards. Additionally, we introduce a curriculum learning approach that allows the model to learn progressively from easier to more challenging instances based on the number of masked spans. We evaluate the effectiveness of our framework in the scenario of open-domain multi-hop question answering. Through extensive experiments, we demonstrate that MaskSearch significantly enhances the performance of LLM-based search agents on both in-domain and out-of-domain downstream tasks.