MaskSearch: Un Framework Universale di Pre-Addestramento per Potenziare la Capacità di Ricerca Agente

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

May 26, 2025

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

cs.AI

Abstract

I Modelli Linguistici Potenziati dal Recupero (RALMs) rappresentano un paradigma classico in cui i modelli migliorano le capacità generative utilizzando conoscenze esterne recuperate tramite un modulo specializzato. I recenti progressi nelle tecniche di Agente consentono ai Modelli Linguistici di Grande Scala (LLMs) di utilizzare autonomamente strumenti per il recupero, la pianificazione e il ragionamento. Sebbene i metodi esistenti basati sul training mostrino promettenti risultati, le loro capacità agentiche sono limitate dalle caratteristiche intrinseche dei dati specifici del compito utilizzati durante l'addestramento. Per potenziare ulteriormente la capacità di ricerca universale degli agenti, proponiamo un nuovo framework di pre-training, MaskSearch. Nella fase di pre-training, introduciamo il compito di Predizione Maschera Potenziata dal Recupero (RAMP), in cui il modello impara a sfruttare strumenti di ricerca per riempire porzioni mascherate su un ampio numero di dati di pre-training, acquisendo così capacità universali di recupero e ragionamento per gli LLMs. Successivamente, il modello viene addestrato su compiti downstream per ottenere ulteriori miglioramenti. Applichiamo sia il Fine-tuning Supervisionato (SFT) che l'Apprendimento per Rinforzo (RL) per l'addestramento. Per il SFT, combiniamo metodi basati su agente e su distillazione per generare dati di training, iniziando con un sistema multi-agente composto da un pianificatore, un riscrittore, un osservatore e seguito da un modello insegnante auto-evolutivo. Per il RL, utilizziamo DAPO come framework di addestramento e adottiamo un sistema di ricompensa ibrido composto da ricompense per le risposte e ricompense per il formato. Inoltre, introduciamo un approccio di apprendimento curriculare che consente al modello di apprendere progressivamente da istanze più semplici a più complesse in base al numero di porzioni mascherate. Valutiamo l'efficacia del nostro framework nello scenario di risposta a domande multi-hop in dominio aperto. Attraverso esperimenti estensivi, dimostriamo che MaskSearch migliora significativamente le prestazioni degli agenti di ricerca basati su LLM sia su compiti downstream in dominio che fuori dominio.

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.