MaskSearch: Um Framework Universal de Pré-Treinamento para Aprimorar a Capacidade de Busca 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

Resumo

Os Modelos de Linguagem Aumentados por Recuperação (RALMs, na sigla em inglês) representam um paradigma clássico em que os modelos aprimoram suas capacidades generativas utilizando conhecimento externo recuperado por meio de um módulo especializado. Avanços recentes em técnicas de Agentes permitem que Modelos de Linguagem de Grande Escala (LLMs, na sigla em inglês) utilizem ferramentas de forma autônoma para recuperação, planejamento e raciocínio. Embora os métodos baseados em treinamento existentes mostrem potencial, suas habilidades agentivas são limitadas pelas características inerentes dos dados específicos da tarefa utilizados durante o treinamento. Para aprimorar ainda mais a capacidade de busca universal dos agentes, propomos um novo framework de pré-treinamento, o MaskSearch. Na etapa de pré-treinamento, introduzimos a tarefa de Predição de Máscara Aumentada por Recuperação (RAMP, na sigla em inglês), na qual o modelo aprende a utilizar ferramentas de busca para preencher trechos mascarados em um grande volume de dados de pré-treinamento, adquirindo assim capacidades universais de recuperação e raciocínio para LLMs. Posteriormente, o modelo é treinado em tarefas subsequentes para obter melhorias adicionais. Aplicamos tanto o Ajuste Fino Supervisionado (SFT, na sigla em inglês) quanto o Aprendizado por Reforço (RL, na sigla em inglês) para o treinamento. Para o SFT, combinamos métodos baseados em agentes e em destilação para gerar dados de treinamento, começando com um sistema multiagente composto por um planejador, reescritor, observador e seguido por um modelo professor de auto-evolução. Já para o RL, empregamos o DAPO como framework de treinamento e adotamos um sistema de recompensas híbrido, composto por recompensas de resposta e recompensas de formato. Além disso, introduzimos uma abordagem de aprendizado curricular que permite ao modelo aprender progressivamente, de instâncias mais fáceis para as mais desafiadoras, com base no número de trechos mascarados. Avaliamos a eficácia do nosso framework no cenário de questionamento de múltiplos saltos em domínio aberto. Por meio de extensos experimentos, demonstramos que o MaskSearch melhora significativamente o desempenho de agentes de busca baseados em LLMs tanto em tarefas subsequentes dentro do domínio quanto fora dele.

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.