MaskSearch: Een Universeel Pre-Trainingsraamwerk om Agentgerichte Zoekcapaciteit te Verbeteren

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

May 26, 2025

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

cs.AI

Samenvatting

Retrieval-Augmented Language Models (RALMs) vertegenwoordigen een klassiek paradigma waarbij modellen hun generatieve capaciteiten verbeteren door gebruik te maken van externe kennis die wordt opgehaald via een gespecialiseerde module. Recente vooruitgang in Agent-technieken stelt Large Language Models (LLMs) in staat om autonoom tools te gebruiken voor het ophalen, plannen en redeneren. Hoewel bestaande op training gebaseerde methoden veelbelovend zijn, worden hun agentische vaardigheden beperkt door inherente kenmerken van de taakspecifieke gegevens die tijdens de training worden gebruikt. Om de universele zoekcapaciteit van agents verder te verbeteren, stellen we een nieuw pre-trainingsframework voor, genaamd MaskSearch. In de pre-trainingsfase introduceren we de Retrieval Augmented Mask Prediction (RAMP)-taak, waarbij het model leert om zoektools te gebruiken om gemaskeerde segmenten in een groot aantal pre-trainingsgegevens in te vullen, waardoor het universele ophaal- en redeneervaardigheden voor LLMs verwerft. Daarna wordt het model getraind op downstreamtaken om verdere verbetering te bereiken. We passen zowel Supervised Fine-tuning (SFT) als Reinforcement Learning (RL) toe voor de training. Voor SFT combineren we agent-gebaseerde en distillatie-gebaseerde methoden om trainingsgegevens te genereren, te beginnen met een multi-agent systeem bestaande uit een planner, herschrijver, observator, gevolgd door een zelf-evoluerend leraarmodel. Voor RL gebruiken we DAPO als het trainingsframework en nemen we een hybride beloningssysteem over dat bestaat uit antwoordbeloningen en formaatbeloningen. Daarnaast introduceren we een curriculumleerbenadering die het model in staat stelt om geleidelijk te leren van eenvoudigere naar meer uitdagende gevallen op basis van het aantal gemaskeerde segmenten. We evalueren de effectiviteit van ons framework in het scenario van open-domein multi-hop vraagbeantwoording. Door uitgebreide experimenten tonen we aan dat MaskSearch de prestaties van LLM-gebaseerde zoekagents aanzienlijk verbetert op zowel in-domein als out-of-domein downstreamtaken.

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.