MaskSearch: A Universal Pre-Training Framework to Enhance Agentic Search
  Capability
Authors: name: Weiqi Wu; name: Xin Guan; name: Shen Huang; name: Yong Jiang; name: Pengjun Xie; name: Fei Huang; name: Jiuxin Cao; name: Hai Zhao; name: Jingren Zhou
Published: 2025-05-26 17:58:50

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.