MaskSearch: A Universal Pre-Training Framework to Enhance Agentic Search Capability
Paper: arXiv:2505.20285 | Published: 2025-05-26 | Authors: Weiqi Wu, Xin Guan, Shen Huang, Yong Jiang, Pengjun Xie, Fei Huang, et al. (9 authors total)
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 nature of their training data. MaskSearch addresses this by pre-training models for universal search capability.
Key Points
- RAMP (Retrieval Augmented Mask Prediction): During pre-training, the model must leverage search tools to fill masked spans across large amounts of pre-training data, building general retrieval and reasoning abilities rather than task-specific ones.
- Two-stage training: After pre-training, the model is fine-tuned on downstream tasks for further gains.
- SFT pipeline: Training data is generated by combining agent-based and distillation-based methods — first a multi-agent system (planner, rewriter, observer), then a self-evolving teacher model.
- RL stage: 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.
- Scenario: open-domain multi-hop question answering.
- Result: MaskSearch significantly improves LLM-based search agents on both in-domain and out-of-domain downstream tasks.
- Exact metrics are in the original PDF; this note summarizes the experimental design from the abstract and public metadata.
- 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)
- Retrieval Augmented Generation and Understanding in Vision: A Survey (arXiv:2503.18016)
Evaluation
Takeaways for Search/Rec/Personalization
1. Architecture: The agentic paradigm makes retrieval count and strategy themselves learnable, beyond the classic retrieve–rerank–generate cascade. 2. Data: Synthetic pre-training data (masked spans + search) is a scalable path to universal search skills; watch for leakage and distribution shift. 3. Transfer: Pre-training on RAMP yields out-of-domain improvements, suggesting genuine capability rather than task overfitting.
Original Abstract (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.