MaskSearch: A Universal Pre-Training Framework to Enhance Agentic Search Capability
Source: arXiv:2505.20285 · Published 2025-05-26 · Authors: Weiqi Wu, Xin Guan, Shen Huang, Yong Jiang, Pengjun Xie, Fei Huang, et al. (9 authors)
Overview
Retrieval-Augmented Language Models (RALMs) enhance generation with external knowledge, and recent agent techniques let LLMs autonomously use tools for retrieval, planning, and reasoning. However, existing training-based methods are limited by the task-specific data used during training. MaskSearch addresses this with a universal pre-training framework for agentic search.
Key points
- RAMP task: Retrieval Augmented Mask Prediction — during pre-training, the model learns to leverage search tools to fill masked spans across large-scale pre-training data, acquiring universal retrieval and reasoning capabilities.
- Two-stage training: after pre-training, the model is fine-tuned on downstream tasks for further gains.
- SFT pipeline: combines agent-based and distillation-based data generation, starting with a multi-agent system (planner, rewriter, observer), followed by a self-evolving teacher model.
- RL setup: uses DAPO as the training framework with a hybrid reward (answer rewards + format rewards).
- Curriculum learning: instances progress from easy to hard based on the number of masked spans.
- Evaluation: on open-domain multi-hop question answering, MaskSearch significantly improves LLM-based search agents on both in-domain and out-of-domain downstream tasks.
- Treats retrieval frequency and strategy itself as a learnable object, moving beyond fixed cascade pipelines (retrieve → rerank → generate).
- Provides a scalable data synthesis recipe (multi-agent generation + distillation) for agentic search training.
- Combines SFT and RL with curriculum design, a practical recipe transferable to industrial search systems.
- Synergizing RAG and Reasoning: A Systematic Review (arXiv:2504.15909)
- AceSearcher: Bootstrapping Reasoning and Search for LLMs via RL (arXiv:2509.24193)
- Agentic Information Retrieval (arXiv:2410.09713)
Original abstract
> 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.
Why it matters
Limitations noted
Quantitative results should be verified against the original PDF tables. As with most agentic search work, latency, cost, evaluation reliability, and open-web safety remain open challenges for production deployment.