MaskSearch: A Universal Pre-Training Framework to Enhance Agentic Search Capability
Paper: https://arxiv.org/abs/2505.20285 Authors: Weiqi Wu, Xin Guan, Shen Huang, Yong Jiang, Pengjun Xie, Fei Huang, et al. (9 authors) Published: 2025-05-26 Category: Agentic Search
TL;DR
MaskSearch is a novel pre-training framework that gives LLM-based agents universal search and reasoning capability. Its central idea is a self-supervised task called Retrieval Augmented Mask Prediction (RAMP): the model must use search tools to fill in masked spans across massive pre-training corpora. Because mask filling is task-agnostic, the learned retrieval and reasoning skills transfer broadly to downstream agentic search tasks.
Motivation
Retrieval-Augmented Language Models (RALMs) enhance generation with knowledge fetched by a retrieval module, and recent agent techniques let LLMs autonomously call tools for retrieval, planning, and reasoning. However, existing training-based methods are constrained by the task-specific nature of their training data, limiting the generality of the resulting agentic search abilities. MaskSearch addresses this with a scalable pre-training objective.
Method
1. Pre-training: RAMP
- Mask spans in large amounts of pre-training data.
- The model must leverage search tools (query formulation, retrieval, result reading) to reconstruct the masked content.
- This yields universal retrieval and reasoning capabilities rather than narrow, task-specific behavior.
- A multi-agent system consisting of a planner, rewriter, and observer.
- A self-evolving teacher model with distillation.
- Training framework: DAPO.
- Hybrid reward: answer rewards + format rewards.
- Curriculum learning: instances are ordered from easy to hard based on the number of masked spans, letting the model progress gradually.
- A Systematic Framework for Enterprise Knowledge Retrieval (arXiv:2512.05411)
- Retrieval Augmented Generation and Understanding in Vision: A Survey (arXiv:2503.18016)
- 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)
- MaskSearch: A Universal Pre-Training Framework to Enhance Agentic Search Capability. arXiv:2505.20285
2. Supervised Fine-Tuning (SFT)
Training data is generated by combining two approaches:
3. Reinforcement Learning (RL)
4. Downstream adaptation
After pre-training, the model is fine-tuned on downstream tasks for further gains.
Evaluation
The framework is evaluated on open-domain multi-hop question answering. Extensive experiments show that MaskSearch significantly enhances the performance of LLM-based search agents on both in-domain and out-of-domain downstream tasks compared to baselines.
Key Takeaways
1. Mask prediction as a pre-training objective for agentic search — RAMP turns web-scale corpora into training signal for tool-using search behavior without needing task-specific labels. 2. SFT + RL pipeline — multi-agent data generation and self-evolving teachers supply supervised data; DAPO with hybrid rewards refines behavior. 3. Curriculum matters — scaling difficulty by the number of masked spans enables progressive learning. 4. Generalization — improvements hold beyond the training distribution, suggesting genuine transferable search capability.
Related Entries
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.