MaskSearch : Un Cadre Universel de Pré-Entraînement pour Améliorer la Capacité de Recherche Agentique

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

May 26, 2025

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

cs.AI

papers.abstract

Les modèles de langage à récupération augmentée (RALMs) représentent un paradigme classique où les modèles améliorent leurs capacités génératives en utilisant des connaissances externes récupérées via un module spécialisé. Les avancées récentes dans les techniques d'agents permettent aux grands modèles de langage (LLMs) d'utiliser de manière autonome des outils pour la récupération, la planification et le raisonnement. Bien que les méthodes existantes basées sur l'entraînement montrent des résultats prometteurs, leurs capacités agentiques sont limitées par les caractéristiques inhérentes aux données spécifiques à la tâche utilisées pendant l'entraînement. Pour améliorer davantage la capacité de recherche universelle des agents, nous proposons un nouveau cadre de pré-entraînement, MaskSearch. Lors de l'étape de pré-entraînement, nous introduisons la tâche de prédiction de masques à récupération augmentée (RAMP), où le modèle apprend à utiliser des outils de recherche pour combler des segments masqués sur un grand nombre de données de pré-entraînement, acquérant ainsi des capacités universelles de récupération et de raisonnement pour les LLMs. Ensuite, le modèle est entraîné sur des tâches en aval pour obtenir une amélioration supplémentaire. Nous appliquons à la fois le réglage fin supervisé (SFT) et l'apprentissage par renforcement (RL) pour l'entraînement. Pour le SFT, nous combinons des méthodes basées sur les agents et sur la distillation pour générer des données d'entraînement, en commençant par un système multi-agents composé d'un planificateur, d'un rédacteur, d'un observateur, suivi d'un modèle enseignant auto-évolutif. Pour le RL, nous utilisons DAPO comme cadre d'entraînement et adoptons un système de récompense hybride composé de récompenses de réponse et de récompenses de format. De plus, nous introduisons une approche d'apprentissage curriculaire qui permet au modèle d'apprendre progressivement, des instances plus faciles aux plus difficiles, en fonction du nombre de segments masqués. Nous évaluons l'efficacité de notre cadre dans le scénario de réponse à des questions multi-sauts en domaine ouvert. À travers des expériences approfondies, nous démontrons que MaskSearch améliore significativement les performances des agents de recherche basés sur les LLMs, aussi bien sur les tâches en aval en domaine qu'hors domaine.

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.